Created
January 9, 2024 21:30
-
-
Save shawngraham/96cd57fde9ed6b548f4a6595f41844c2 to your computer and use it in GitHub Desktop.
I've got a bunch of old screenshots of social media posts. I wanted to crop out the text & comments and end up with a folder of just the original images in the post. This was an attempt at that. It kinda works.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import numpy as np | |
from PIL import Image | |
from skimage.color import rgb2gray | |
from skimage import filters, morphology | |
import argparse | |
def find_nonzero_extents(mask): | |
rows = np.any(mask, axis=1) | |
first_nonzero_row = next((i for i, row in enumerate(rows) if row), None) | |
last_nonzero_row = next((i for i, row in enumerate(rows[::-1]) if row), None) | |
if last_nonzero_row is not None: | |
last_nonzero_row = mask.shape[0] - last_nonzero_row | |
return first_nonzero_row, last_nonzero_row | |
def crop_image(image_np, text_bands_mask): | |
mask_rows = text_bands_mask.shape[0] | |
top_first_nonzero, top_last_nonzero = find_nonzero_extents(text_bands_mask[:mask_rows//2]) | |
bottom_first_nonzero, bottom_last_nonzero = find_nonzero_extents(text_bands_mask[mask_rows//2:]) | |
if bottom_first_nonzero is not None: | |
bottom_first_nonzero += mask_rows // 2 | |
if bottom_last_nonzero is not None: | |
bottom_last_nonzero += mask_rows // 2 | |
top_region_end = top_last_nonzero if top_last_nonzero is not None else 0 | |
bottom_region_start = bottom_first_nonzero if bottom_first_nonzero is not None else mask_rows | |
return image_np[top_region_end:bottom_region_start] | |
def detect_horizontal_text_bands(image_np, band_height_ratio=0.1): | |
# Convert image to grayscale | |
grayscale = rgb2gray(image_np[..., :3]) # Assuming RGB | |
# Define the height of bands (top and bottom) based on the image height and a ratio | |
band_height = int(grayscale.shape[0] * band_height_ratio) | |
# Extract the top and bottom bands of the image | |
top_band = grayscale[:band_height] | |
bottom_band = grayscale[-band_height:] | |
# Use edge detection to find horizontal lines/edges | |
edges_top = filters.sobel_h(top_band) | |
edges_bottom = filters.sobel_h(bottom_band) | |
# Threshold edges to binary | |
threshold_top = filters.threshold_otsu(edges_top) | |
threshold_bottom = filters.threshold_otsu(edges_bottom) | |
edges_top = edges_top > threshold_top | |
edges_bottom = edges_bottom > threshold_bottom | |
# Use horizontal structure for morphology operations to highlight horizontal lines | |
kernel_width = 30 # This can be adjusted as needed | |
kernel = morphology.rectangle(1, kernel_width) | |
morph_top = morphology.closing(edges_top, kernel) | |
morph_bottom = morphology.closing(edges_bottom, kernel) | |
# Combine the top and bottom masks to create a full mask | |
text_bands_mask = np.zeros_like(grayscale, dtype=bool) | |
text_bands_mask[:band_height] = morph_top | |
text_bands_mask[-band_height:] = morph_bottom | |
return text_bands_mask | |
# Define a new function to process each image file | |
def process_image(input_image_path, output_dir): | |
original_image = Image.open(input_image_path) | |
original_image_np = np.array(original_image) | |
# Detect text bands and crop the image | |
text_bands_mask = detect_horizontal_text_bands(original_image_np) | |
cropped_image_np = crop_image(original_image_np, text_bands_mask) | |
# Save the cropped image | |
base_name = os.path.basename(input_image_path) | |
file_name, _ = os.path.splitext(base_name) | |
cropped_image_path = os.path.join(output_dir, f"{file_name}_cropped.png") | |
Image.fromarray(cropped_image_np).save(cropped_image_path) | |
print(f"Processed and saved: {cropped_image_path}") | |
# Use argparse to handle command line arguments | |
def main(): | |
parser = argparse.ArgumentParser(description="Process a folder of images, detect text bands, and crop the images.") | |
parser.add_argument('input_folder', help='The folder containing the images to process') | |
parser.add_argument('output_folder', help='The folder where the cropped images will be saved') | |
args = parser.parse_args() | |
# Create output folder if it doesn't exist | |
if not os.path.exists(args.output_folder): | |
os.makedirs(args.output_folder) | |
# Process each file in the input folder | |
for file_name in os.listdir(args.input_folder): | |
file_path = os.path.join(args.input_folder, file_name) | |
if os.path.isfile(file_path): | |
try: | |
process_image(file_path, args.output_folder) | |
except Exception as e: | |
print(f"Error processing {file_name}: {e}") | |
if __name__ == "__main__": | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment