Change dataset to checkered + processing pipeline

app.py

@@ -31,7 +31,7 @@ google_analytics_tracking_id = os.getenv("GOOGLE_ANALYTICS_TRACKING_ID")
 logging.basicConfig(level=logging.INFO)
 
 # Load datasets and initialize database
-dataset = load_dataset("bgsys/background-removal-arena-green-v0.2", split='train')
+dataset = load_dataset("bgsys/background-removal-arena-green_v0_clothing_checkered", split='train')
 fill_database_once()
 
 # Directory setup for JSON dataset
@@ -136,23 +136,20 @@ def get_notice_markdown():
     """
 
 def compute_mask_difference(segmented_a, segmented_b):
-    """Compute the absolute difference between two image masks, ignoring green background."""
-    mask_a = np.asarray(segmented_a)
-    mask_b = np.asarray(segmented_b)
-
-    # Define the green background color
-    green_background = (0, 255, 0, 255)
-
-    # Create a binary mask where non-green and non-transparent pixels are marked as 1
-    mask_a_1d = np.where(
-        (mask_a[..., :3] != green_background[:3]).any(axis=-1) & (mask_a[..., 3] != 0), 1, 0
-    )
-    mask_b_1d = np.where(
-        (mask_b[..., :3] != green_background[:3]).any(axis=-1) & (mask_b[..., 3] != 0), 1, 0
-    )
-
-    # Compute the absolute difference between the masks
-    return np.abs(mask_a_1d - mask_b_1d)
+    """Compute the difference between two images across all channels."""
+    mask_a = np.asarray(segmented_a, dtype=np.float32) / 255.0
+    mask_b = np.asarray(segmented_b, dtype=np.float32) / 255.0
+    
+    # Compute absolute difference across all channels (RGBA)
+    difference = np.abs(mask_a - mask_b)
+    
+    # Take the maximum difference across channels at each pixel
+    max_diff = np.max(difference, axis=-1)
+    
+    # Apply non-linear transformation to enhance visibility of smaller differences
+    # Using square root to compress the range while maintaining continuity
+    # Multiply by 3 to make it more visible
+    return np.sqrt(max_diff) * 3
 
 js = r"""
 function load_zoom() {
@@ -169,7 +166,7 @@ function load_zoom() {
     });
 
     // Choose a scale factor
-    const scale = 2;
+    const scale = 5;
 
     function handleMouseMove(e) {
         const rect = e.currentTarget.getBoundingClientRect();
@@ -297,7 +294,7 @@ def gradio_interface():
                         image_a = gr.Image(value=segmented_a, label="Image A", width=image_width, height=image_height)
                         image_b = gr.Image(value=segmented_b, label="Image B", width=image_width, height=image_height)
                         input_image_display = gr.AnnotatedImage(
-                            value=(input_image, [(mask_difference > 0, button_name)]), 
+                            value=(input_image, [(mask_difference, button_name)]), 
                             width=image_width,
                             height=image_height
                         )

image_processing_pipeline.py

@@ -11,9 +11,9 @@ from utils.removebg import iterate_over_directory as removebg_iterate
 from utils.photoroom import iterate_over_directory as photoroom_iterate
 from utils.bria_rmbg20 import iterate_over_directory as bria_iterate
 from utils.clipdrop import iterate_over_directory as clipdrop_iterate
-from utils.add_green_background import process_directory as add_green_background_process
 from utils.upload_to_dataset import upload_to_dataset
 from utils.resize_processed_images import process_images as downsize_processed_images
+from utils.add_checkered_background import process_directory as add_checkered_background_process
 
 def check_env_variables():
     """Check if the necessary environment variables are loaded."""
@@ -30,7 +30,7 @@ def check_env_variables():
 
 def copy_images(source_dir, dest_dir):
     os.makedirs(dest_dir, exist_ok=True)
-    valid_extensions = ('.png', '.jpg', '.jpeg')
+    valid_extensions = ('.png', '.jpg', '.jpeg', '.webp')
 
     # Walk through the source directory
     for root, _, files in os.walk(source_dir):
@@ -57,7 +57,6 @@ def main():
     parser = argparse.ArgumentParser(description="Image Processing Pipeline")
     parser.add_argument("--input-dir", type=str, default="original-images", help="Input directory for images")
     parser.add_argument("--work-dir", type=str, default="workdir", help="Working directory for intermediate images")
-    parser.add_argument("--output-dir", type=str, default="final-images", help="Output directory for final images")
     parser.add_argument("--dataset-name", type=str, help="Name of the dataset to upload to Hugging Face Hub")
     parser.add_argument("--push-dataset", action="store_true", help="Push the dataset to the Hugging Face Hub")
 
@@ -66,10 +65,10 @@ def main():
     # Define intermediate directories within the work directory
     input_resized_dir = os.path.join(args.work_dir, "resized")
     bg_removed_dir = os.path.join(args.work_dir, "background-removed")
-    green_bg_dir = os.path.join(args.work_dir, "green-background")
+    checkered_bg_dir = os.path.join(args.work_dir, "checkered-background")
 
     # Ensure all directories exist
-    for directory in [input_resized_dir, bg_removed_dir, green_bg_dir]:
+    for directory in [input_resized_dir, bg_removed_dir, checkered_bg_dir]:
         os.makedirs(directory, exist_ok=True)
 
     # Step 4: Move images to final output directory
@@ -100,23 +99,13 @@ def main():
         executor.submit(bria_iterate, input_resized_dir, bg_removal_dirs["bria"])
         executor.submit(clipdrop_iterate, input_resized_dir, bg_removal_dirs["clipdrop"])
 
-    print("Adding green background...")
-    add_green_background_process(bg_removed_dir, green_bg_dir)
-
-    print("Resizing processed images...")
-    target_width = 800 
-    subdirectories = ["bria", "photoroom", "clipdrop", "removebg"]
-    os.makedirs(args.output_dir, exist_ok=True)
-    for subdir in subdirectories:
-        input_directory = os.path.join(green_bg_dir, subdir)
-        output_directory = os.path.join(args.output_dir, subdir)
-        downsize_processed_images(input_directory, output_directory, target_width)
-    
-    original_output_directory = os.path.join(args.output_dir, "web-original-images")
-    downsize_processed_images(input_resized_dir, original_output_directory, target_width)
+
+    print("Adding checkered background...")
+    add_checkered_background_process(bg_removed_dir, checkered_bg_dir)
+
 
     if args.dataset_name:
-        upload_to_dataset(original_output_directory, args.output_dir, args.dataset_name, dry_run=not args.push_dataset)
+        upload_to_dataset(input_resized_dir, checkered_bg_dir, args.dataset_name, dry_run=not args.push_dataset)
     else:
         print("Please provide a dataset name using --dataset-name")