use green backgrounds + fix button naming

app.py

@@ -21,7 +21,7 @@ from huggingface_hub import CommitScheduler
 token = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
 # Load datasets
-dataset = load_dataset("bgsys/background-removal-arena-test", split='train')
+dataset = load_dataset("bgsys/background-removal-arena-green", split='train')
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -131,13 +131,20 @@ def get_notice_markdown():
     """
 
 def compute_mask_difference(segmented_a, segmented_b):
-    """Compute the absolute difference between two image masks."""
+    """Compute the absolute difference between two image masks, ignoring green background."""
     mask_a = np.asarray(segmented_a)
     mask_b = np.asarray(segmented_b)
 
-    # Create a binary mask where non-transparent pixels are marked as 1
-    mask_a_1d = np.where(mask_a[..., 3] != 0, 1, 0)
-    mask_b_1d = np.where(mask_b[..., 3] != 0, 1, 0)
+    # 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)
@@ -146,6 +153,7 @@ def gradio_interface():
     """Create and return the Gradio interface."""
     with gr.Blocks() as demo:
         gr.Markdown("# Background Removal Arena")
+        button_name = "Difference between masks"
 
         with gr.Tabs() as tabs:
             with gr.Tab("⚔️ Arena (battle)", id=0):
@@ -168,7 +176,7 @@ def gradio_interface():
                         height=500
                     )
                     input_image_display = gr.AnnotatedImage(
-                        value=(input_image, [(mask_difference > 0, "Difference between masks")]),
+                        value=(input_image, [(mask_difference > 0, button_name)]),
                         label="Input Image",
                         width=500,
                         height=500
@@ -218,11 +226,8 @@ def gradio_interface():
 
                     try:
                         logging.debug("Adding vote data to the database: %s", vote_data)
-
-                        # Only add vote if running in space
-                        if is_running_in_space():
-                            result = add_vote(vote_data)
-                            logging.info("Vote successfully recorded in the database with ID: %s", result["id"])
+                        result = add_vote(vote_data)
+                        logging.info("Vote successfully recorded in the database with ID: %s", result["id"])
                     except Exception as e:
                         logging.error("Error recording vote: %s", str(e))
 
@@ -236,7 +241,7 @@ def gradio_interface():
                     # Update the notice markdown with the new vote count
                     new_notice_markdown = get_notice_markdown()
 
-                    return (fpath_input.value, (new_input_image, [(mask_difference, "Mask")]), new_segmented_a,
+                    return (fpath_input.value, (new_input_image, [(mask_difference, button_name)]), new_segmented_a,
                              new_segmented_b, model_a_name.value, model_b_name.value, new_notice_markdown)
            
             with gr.Tab("🏆 Leaderboard", id=1) as leaderboard_tab:

db.py

@@ -1,4 +1,5 @@
 import os
+import logging
 from sqlalchemy import create_engine, Column, Integer, String, DateTime
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, Session
@@ -48,6 +49,7 @@ def add_vote(vote_data):
         db.add(db_vote)
         db.commit()
         db.refresh(db_vote)
+        logging.info("Vote registered with ID: %s, using database: %s", db_vote.id, DATABASE_URL)
         return {"id": db_vote.id, "user_id": db_vote.user_id, "timestamp": db_vote.timestamp}
 
 # Function to get all votes

utils/add_green_background.py

@@ -0,0 +1,38 @@
+import os
+from PIL import Image
+
+def add_green_background_to_image(image_path, output_path, background_color=(0, 255, 0)):
+    """Add a green background to an image and save it as PNG."""
+    with Image.open(image_path) as img:
+        img = img.convert("RGBA")
+        background = Image.new("RGBA", img.size, background_color + (255,))
+        combined = Image.alpha_composite(background, img)
+        combined.save(output_path, "PNG")
+
+def process_directory(input_dir, output_dir, background_color=(0, 255, 0)):
+    """Recursively process a directory to add a green background to all images and convert them to PNG."""
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    for root, _, files in os.walk(input_dir):
+        for file in files:
+            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
+                input_path = os.path.join(root, file)
+                relative_path = os.path.relpath(input_path, input_dir)
+                output_path = os.path.join(output_dir, os.path.splitext(relative_path)[0] + '.png')
+
+                # Ensure the output directory exists
+                os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+                # Check if the output file already exists
+                if not os.path.exists(output_path):
+                    # Add green background to the image and convert to PNG
+                    add_green_background_to_image(input_path, output_path, background_color)
+                    print(f"Processed: {input_path} -> {output_path}")
+                else:
+                    print(f"Skipped: {output_path} already exists")
+
+# Example usage
+input_directory = "../../background-removal-arena-v0/train/data/resized"
+output_directory = "../data/resized-green/"
+process_directory(input_directory, output_directory)

utils/upload-to-dataset.py

@@ -3,74 +3,82 @@ from huggingface_hub import HfApi
 import os
 from collections import defaultdict
 import pandas as pd
+import argparse
 
-# Define the path to your images
-IMAGE_DIR = "../../background-removal-arena-v0/train/data/resized"
+def upload_to_dataset(image_dir, dataset_name):
+    # Define the dataset features with dedicated columns for each model
+    features = Features({
+        "original_image": Image(),  # Original image feature
+        "clipdrop_image": Image(),  # Clipdrop segmented image
+        "bria_image": Image(),      # Bria segmented image
+        "photoroom_image": Image(), # Photoroom segmented image
+        "removebg_image": Image(),  # RemoveBG segmented image
+        "original_filename": Value("string")  # Original filename
+    })
 
-# Define the dataset features with dedicated columns for each model
-features = Features({
-    "original_image": Image(),  # Original image feature
-    "clipdrop_image": Image(),  # Clipdrop segmented image
-    "bria_image": Image(),      # Bria segmented image
-    "photoroom_image": Image(), # Photoroom segmented image
-    "removebg_image": Image(),  # RemoveBG segmented image
-    "original_filename": Value("string")  # Original filename
-})
+    # Load image paths and metadata
+    data = defaultdict(lambda: {
+        "clipdrop_image": None,
+        "bria_image": None,
+        "photoroom_image": None,
+        "removebg_image": None
+    })
 
-# Load image paths and metadata
-data = defaultdict(lambda: {
-    "clipdrop_image": None,
-    "bria_image": None,
-    "photoroom_image": None,
-    "removebg_image": None
-})
+    # Walk into the web-original-images folder
+    web_original_images_dir = os.path.join(image_dir, "web-original-images")
+    for root, _, files in os.walk(web_original_images_dir):
+        for f in files:
+            if f.endswith(('.png', '.jpg', '.jpeg')):
+                original_image_path = os.path.join(root, f)
+                data[f]["original_image"] = original_image_path
+                data[f]["original_filename"] = f
 
-# Walk into the web-original-images folder
-web_original_images_dir = os.path.join(IMAGE_DIR, "web-original-images")
-for root, _, files in os.walk(web_original_images_dir):
-    for f in files:
-        if f.endswith(('.png', '.jpg', '.jpeg')):
-            original_image_path = os.path.join(root, f)
-            data[f]["original_image"] = original_image_path
-            data[f]["original_filename"] = f
+                # Check for corresponding images in other directories
+                for source in ["clipdrop", "bria", "photoroom", "removebg"]:
+                    # Check for processed images ending in .png or .jpg
+                    for ext in ['.png', '.jpg']:
+                        processed_image_filename = os.path.splitext(f)[0] + ext
+                        source_image_path = os.path.join(image_dir, source, processed_image_filename)
+        
+                        if os.path.exists(source_image_path):
+                            data[f][f"{source}_image"] = source_image_path
+                            break  # Stop checking other extensions if a file is found
 
-            # Check for corresponding images in other directories
-            for source in ["clipdrop", "bria", "photoroom", "removebg"]:
-                # Check for processed images ending in .png or .jpg
-                for ext in ['.png', '.jpg']:
-                    processed_image_filename = os.path.splitext(f)[0] + ext
-                    source_image_path = os.path.join(IMAGE_DIR, source, processed_image_filename)
-    
-                    if os.path.exists(source_image_path):
-                        data[f][f"{source}_image"] = source_image_path
-                        break  # Stop checking other extensions if a file is found
+    # Convert the data to a dictionary of lists
+    dataset_dict = {
+        "original_image": [],
+        "clipdrop_image": [],
+        "bria_image": [],
+        "photoroom_image": [],
+        "removebg_image": [],
+        "original_filename": []
+    }
 
-# Convert the data to a dictionary of lists
-dataset_dict = {
-    "original_image": [],
-    "clipdrop_image": [],
-    "bria_image": [],
-    "photoroom_image": [],
-    "removebg_image": [],
-    "original_filename": []
-}
+    for filename, entry in data.items():
+        if "original_image" in entry:
+            dataset_dict["original_image"].append(entry["original_image"])
+            dataset_dict["clipdrop_image"].append(entry["clipdrop_image"])
+            dataset_dict["bria_image"].append(entry["bria_image"])
+            dataset_dict["photoroom_image"].append(entry["photoroom_image"])
+            dataset_dict["removebg_image"].append(entry["removebg_image"])
+            dataset_dict["original_filename"].append(filename)
 
-for filename, entry in data.items():
-    if "original_image" in entry:
-        dataset_dict["original_image"].append(entry["original_image"])
-        dataset_dict["clipdrop_image"].append(entry["clipdrop_image"])
-        dataset_dict["bria_image"].append(entry["bria_image"])
-        dataset_dict["photoroom_image"].append(entry["photoroom_image"])
-        dataset_dict["removebg_image"].append(entry["removebg_image"])
-        dataset_dict["original_filename"].append(filename)
+    # 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)
 
-# 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)
+    # Create a Dataset
+    dataset = Dataset.from_dict(dataset_dict, features=features)
 
-# Create a Dataset
-dataset = Dataset.from_dict(dataset_dict, features=features)
+    # Push the dataset to Hugging Face Hub
+    api = HfApi()
+    dataset.push_to_hub(dataset_name, token=api.token)
 
-# Push the dataset to Hugging Face Hub
-api = HfApi()
-dataset.push_to_hub("bgsys/background-removal-arena-test", token=api.token)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Upload images to a Hugging Face dataset.")
+    parser.add_argument("image_dir", type=str, help="Directory containing the images.")
+    parser.add_argument("dataset_name", type=str, help="Name of the dataset to upload to Hugging Face Hub.")
+    
+    args = parser.parse_args()
+    
+    upload_to_dataset(args.image_dir, args.dataset_name)