App - V3 - Fully Complete

app.py

@@ -5,11 +5,12 @@ import numpy as np
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 from uuid import uuid4
 import os
-from huggingface_hub import upload_folder, login
+from huggingface_hub import upload_folder
 from PIL import Image as PILImage
 from datasets import Dataset, Features, Array2D, Image
 import shutil
-import time
+import random
+from datasets import load_dataset
 
 MODEL = "facebook/sam2-hiera-large"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -17,7 +18,7 @@ PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
 
 DESTINATION_DS = "amaye15/object-segmentation"
 
-login(os.getenv("TOKEN"))
+# login(os.getenv("TOKEN"))
 
 IMAGE = None
 MASKS = None
@@ -25,6 +26,21 @@ MASKED_IMAGES = None
 INDEX = None
 
 
+ds_name = ["amaye15/product_labels"]  #  "amaye15/Products-10k", "amaye15/receipts"
+choices = ["test", "train"]
+max_len = None
+
+ds_stream = load_dataset(random.choice(ds_name), streaming=True)
+
+
+ds_split = ds_stream[random.choice(choices)]
+
+ds_iter = ds_split.iter(batch_size=1)
+
+for idx, val in enumerate(ds_iter):
+    max_len = idx
+
+
 def prompter(prompts):
 
     image = np.array(prompts["image"])  # Convert the image to a numpy array
@@ -116,114 +132,139 @@ def save_selected_mask(image, mask, output_dir="output"):
 
     shutil.rmtree(folder_path)
 
-    iframe_code = "Success - Check out the 'Results' tab."
+    iframe_code = """## Success! 🎉🤖✅
+
+You've successfully contributed to the dataset. 
+
+Please note that because new data has been added to the dataset, it may take a couple of minutes to render. 
+
+Check it out here:
+
+[Object Segmentation Dataset](https://huggingface.co/datasets/amaye15/object-segmentation)
+"""
 
     return iframe_code
 
-    # time.sleep(5)
 
-    # # Add a random query parameter to force reload
-    # random_param = uuid4()
-    # iframe_code = f"""
-    # <iframe
-    #   src="https://huggingface.co/datasets/{DESTINATION_DS}/embed/viewer/default/train"
-    #   frameborder="0"
-    #   width="100%"
-    #   height="560px"
-    # ></iframe>
-    # """
+def get_random_image():
+    """Get a random image from the dataset."""
+    global max_len
+    random_idx = random.choice(range(max_len))
+    image_data = list(ds_split.skip(random_idx).take(1))[0]["pixel_values"]
+    formatted_image = {
+        "image": np.array(image_data),
+        "points": [],
+    }  # Create the correct format
+    return formatted_image
 
 
 # Define the Gradio Blocks app
 with gr.Blocks() as demo:
+    gr.Markdown("# Object Segmentation- Image Point Collector and Mask Overlay Tool")
+    gr.Markdown(
+        """
+        This application utilizes **Segment Anything V2 (SAM2)** to allow you to upload an image or select a random image from a dataset and interactively generate segmentation masks based on multiple points you select on the image. 
+
+        ### How It Works:
+        1. **Upload or Select an Image**: You can either upload your own image or use a random image from the dataset.
+        2. **Point Selection**: Click on the image to indicate points of interest. You can add multiple points, and these will be used collectively to generate segmentation masks using SAM2.
+        3. **Mask Generation**: The app will generate up to three different segmentation masks for the selected points, each displayed separately with a red overlay.
+        4. **Mask Selection**: Carefully review the generated masks and select the one that best fits your needs. **It's important to choose the correct mask, as your selection will be saved and used for further processing.**
+        5. **Save and Contribute**: Save the selected mask along with the image to a dataset, contributing to a shared dataset on Hugging Face.
+
+        **Disclaimer**: All images and masks you work with will be collected and stored in a public dataset. Please ensure that you are comfortable with your selections and the data you provide before saving.
+        
+        This tool is particularly useful for creating precise object segmentation masks for computer vision tasks, such as training models or generating labeled datasets.
+        """
+    )
 
-    with gr.Tab("Object Segmentation - Point Prompt"):
-        gr.Markdown("# Image Point Collector with Multiple Separate Mask Overlays")
-        gr.Markdown(
-            "Upload an image, click on it, and get each predicted mask overlaid separately in red on individual images."
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.State()
+            # Input: ImagePrompter for uploaded image
+            upload_image_input = ImagePrompter(show_label=False)
+
+            random_image_button = gr.Button("Use Random Image")
+
+            submit_button = gr.Button("Submit")
+
+    with gr.Row():
+        with gr.Column():
+            # Outputs: Up to 3 overlay images
+            image_output_1 = gr.Image(show_label=False)
+        with gr.Column():
+            image_output_2 = gr.Image(show_label=False)
+        with gr.Column():
+            image_output_3 = gr.Image(show_label=False)
+
+    # Dropdown for selecting the correct mask
+    with gr.Row():
+        mask_selector = gr.Radio(
+            label="Select the correct mask",
+            choices=["Mask 1", "Mask 2", "Mask 3"],
+            type="index",
         )
+        # selected_mask_output = gr.Image(show_label=False)
 
-        with gr.Row():
-            with gr.Column():
-                # Input: ImagePrompter
-                image_input = ImagePrompter(show_label=False)
-                submit_button = gr.Button("Submit")
-        with gr.Row():
-            with gr.Column():
-                # Outputs: Up to 3 overlay images
-                image_output_1 = gr.Image(show_label=False)
-            with gr.Column():
-                image_output_2 = gr.Image(show_label=False)
-            with gr.Column():
-                image_output_3 = gr.Image(show_label=False)
-
-        # Dropdown for selecting the correct mask
-        with gr.Row():
-            mask_selector = gr.Radio(
-                label="Select the correct mask",
-                choices=["Mask 1", "Mask 2", "Mask 3"],
-                type="index",
-            )
-            # selected_mask_output = gr.Image(show_label=False)
-
-        save_button = gr.Button("Save Selected Mask and Image")
-        iframe_display = gr.Markdown()
-
-        # Define the action triggered by the submit button
-        submit_button.click(
-            fn=prompter,
-            inputs=image_input,
-            outputs=[image_output_1, image_output_2, image_output_3, gr.State()],
-            show_progress=True,
-        )
+    save_button = gr.Button("Save Selected Mask and Image")
+    iframe_display = gr.Markdown()
 
-        # Define the action triggered by mask selection
-        mask_selector.change(
-            fn=select_mask,
-            inputs=[mask_selector, image_output_1, image_output_2, image_output_3],
-            outputs=gr.State(),
-        )
+    # Logic for the random image button
+    random_image_button.click(
+        fn=get_random_image,
+        inputs=None,
+        outputs=upload_image_input,  # Pass the formatted random image to ImagePrompter
+    )
 
-        # Define the action triggered by the save button
-        save_button.click(
-            fn=save_selected_mask,
-            inputs=[gr.State(), gr.State()],
-            outputs=iframe_display,
-            show_progress=True,
-        )
-    with gr.Tab("Results"):
-        with gr.Row():
-            gr.HTML(
-                f"""
-                <iframe
-                  src="https://huggingface.co/datasets/{DESTINATION_DS}/embed/viewer/default/train"
-                  frameborder="0"
-                  width="100%"
-                  height="560px"
-                ></iframe>
-                """
-            )
-            # with gr.Column():
-            #     source = gr.Textbox(label="Source Dataset")
-            #     source_display = gr.Markdown()
-            #     iframe_display = gr.HTML()
-
-            #     source.change(
-            #         save_dataset_name,
-            #         inputs=(gr.State("source_dataset"), source),
-            #         outputs=(source_display, iframe_display),
-            #     )
-
-            # with gr.Column():
-
-            #     destination = gr.Textbox(label="Destination Dataset")
-            #     destination_display = gr.Markdown()
-
-            #     destination.change(
-            #         save_dataset_name,
-            #         inputs=(gr.State("destination_dataset"), destination),
-            #         outputs=destination_display,
-            #     )
+    # Logic to use uploaded image
+    upload_image_input.change(
+        fn=lambda img: img, inputs=upload_image_input, outputs=image_input
+    )
+    # Define the action triggered by the submit button
+    submit_button.click(
+        fn=prompter,
+        inputs=upload_image_input,  # The final image input (whether uploaded or random)
+        outputs=[image_output_1, image_output_2, image_output_3, gr.State()],
+        show_progress=True,
+    )
+
+    # Define the action triggered by mask selection
+    mask_selector.change(
+        fn=select_mask,
+        inputs=[mask_selector, image_output_1, image_output_2, image_output_3],
+        outputs=gr.State(),
+    )
+
+    # Define the action triggered by the save button
+    save_button.click(
+        fn=save_selected_mask,
+        inputs=[gr.State(), gr.State()],
+        outputs=iframe_display,
+        show_progress=True,
+    )
 
 # Launch the Gradio app
 demo.launch()
+
+
+# with gr.Column():
+#     source = gr.Textbox(label="Source Dataset")
+#     source_display = gr.Markdown()
+#     iframe_display = gr.HTML()
+
+#     source.change(
+#         save_dataset_name,
+#         inputs=(gr.State("source_dataset"), source),
+#         outputs=(source_display, iframe_display),
+#     )
+
+# with gr.Column():
+
+#     destination = gr.Textbox(label="Destination Dataset")
+#     destination_display = gr.Markdown()
+
+#     destination.change(
+#         save_dataset_name,
+#         inputs=(gr.State("destination_dataset"), destination),
+#         outputs=destination_display,
+#     )