Fix original image downsizing / Add image size check before upload

image_processing_pipeline.py

@@ -113,7 +113,7 @@ def main():
         downsize_processed_images(input_directory, output_directory, target_width)
     
     original_output_directory = os.path.join(args.output_dir, "web-original-images")
-    downsize_processed_images(original_images_dir, original_output_directory, target_width)
+    downsize_processed_images(input_resized_dir, original_output_directory, target_width)
 
     if args.dataset_name:
         upload_to_dataset(original_output_directory, args.output_dir, args.dataset_name, dry_run=not args.push_dataset)

utils/resize_processed_images.py

@@ -31,14 +31,13 @@ def resize_image(input_path, output_path, target_width):
     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)
 

utils/upload_to_dataset.py

@@ -4,6 +4,8 @@ import os
 from collections import defaultdict
 import pandas as pd
 import argparse
+from PIL import Image as PILImage
+import sys
 
 def upload_to_dataset(original_images_dir, processed_images_dir, dataset_name, dry_run=False):
     # Define the dataset features with dedicated columns for each model
@@ -53,8 +55,21 @@ def upload_to_dataset(original_images_dir, processed_images_dir, dataset_name, d
         "original_filename": []
     }
 
+    errors = []
+
     for filename, entry in data.items():
         if "original_image" in entry:
+            # Check if all images have the same size
+            try:
+                original_size = PILImage.open(entry["original_image"]).size
+                for source in ["clipdrop_image", "bria_image", "photoroom_image", "removebg_image"]:
+                    if entry[source] is not None:
+                        processed_size = PILImage.open(entry[source]).size
+                        if processed_size != original_size:
+                            errors.append(f"Size mismatch for {filename}: {source} image size {processed_size} does not match original size {original_size}.")
+            except Exception as e:
+                errors.append(f"Error processing {filename}: {e}")
+
             dataset_dict["original_image"].append(entry["original_image"])
             dataset_dict["clipdrop_image"].append(entry["clipdrop_image"])
             dataset_dict["bria_image"].append(entry["bria_image"])
@@ -62,6 +77,11 @@ def upload_to_dataset(original_images_dir, processed_images_dir, dataset_name, d
             dataset_dict["removebg_image"].append(entry["removebg_image"])
             dataset_dict["original_filename"].append(filename)
 
+    if errors:
+        for error in errors:
+            print(error)
+        sys.exit(1)
+
     # Save the data dictionary to a CSV file for inspection
     df = pd.DataFrame.from_dict(dataset_dict)
     df.to_csv("image_data.csv", index=False)