background-removal-arena / utils /resize_processed_images.py
tdurbor's picture
Consolidate image preparation pipeline
fe62bf5
raw
history blame
3.64 kB
from PIL import Image, ExifTags
import os
from concurrent.futures import ThreadPoolExecutor
def create_directory(path):
"""Create a directory if it doesn't exist."""
os.makedirs(path, exist_ok=True)
def correct_orientation(img):
"""Correct image orientation using EXIF data."""
try:
for orientation in ExifTags.TAGS.keys():
if ExifTags.TAGS[orientation] == 'Orientation':
break
exif = img._getexif()
if exif is not None:
orientation = exif.get(orientation, None)
if orientation == 3:
img = img.rotate(180, expand=True)
elif orientation == 6:
img = img.rotate(270, expand=True)
elif orientation == 8:
img = img.rotate(90, expand=True)
except (AttributeError, KeyError, IndexError):
# Cases: image doesn't have getexif
pass
return img
def resize_image(input_path, output_path, target_width):
"""Resize an image to the target width while maintaining aspect ratio."""
with Image.open(input_path) as img:
# Correct orientation
img = correct_orientation(img)
# Calculate the new height to maintain the aspect ratio
width_percent = target_width / img.width
target_height = int(img.height * width_percent)
# Resize the image
img = img.resize((target_width, target_height), Image.LANCZOS)
# Save the resized image in the same format as the input
img.save(output_path, format=img.format)
def process_image_file(file_path, result_path, target_width):
"""Process a single image file."""
if not os.path.exists(result_path):
print(f"Resizing {file_path} to {result_path}")
resize_image(file_path, result_path, target_width)
else:
print(f"Skipped {file_path}, already resized.")
def process_images(input_directory, output_directory, target_width):
"""Process and resize images from the input directory to the output directory."""
create_directory(output_directory)
with ThreadPoolExecutor() as executor:
for root, _, files in os.walk(input_directory):
for file in files:
if file.lower().endswith(('.png', '.jpg', '.jpeg')):
file_path = os.path.join(root, file)
result_file_name = os.path.splitext(file)[0] + os.path.splitext(file)[1]
result_path = os.path.join(output_directory, result_file_name)
executor.submit(process_image_file, file_path, result_path, target_width)
def main():
"""Main function to resize images in specified subdirectories."""
# Process images in the processed directory
base_input_directory = "../data/processed"
base_output_directory = "../data/resized"
target_width = 800 # Set the desired width for web display
# List of subdirectories to process
subdirectories = ["bria", "photoroom", "clipdrop", "removebg"]
for subdir in subdirectories:
input_directory = os.path.join(base_input_directory, subdir)
output_directory = os.path.join(base_output_directory, subdir)
process_images(input_directory, output_directory, target_width)
# Additionally, process images in the resized-original-images directory
original_input_directory = "../data/resized-original-images"
original_output_directory = "../data/web-original-images"
process_images(original_input_directory, original_output_directory, target_width)
print("Image resizing completed.")
if __name__ == "__main__":
main()