App - V2 - Improved File Formats & UI

app.py

@@ -3,24 +3,25 @@ from gradio_image_prompter import ImagePrompter
 import torch
 import numpy as np
 from sam2.sam2_image_predictor import SAM2ImagePredictor
-from PIL import Image
 from uuid import uuid4
 import os
 from huggingface_hub import upload_folder, login
+from PIL import Image as PILImage
+from datasets import Dataset, Features, Array2D, Image
 import shutil
+import time
 
 MODEL = "facebook/sam2-hiera-large"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
 
+DESTINATION_DS = "amaye15/object-segmentation"
 
-login(os.getenv("TOKEN"))
-
-GLOBALS = {}
-
+# login(os.getenv("TOKEN"))
 
 IMAGE = None
 MASKS = None
+MASKED_IMAGES = None
 INDEX = None
 
 
@@ -44,20 +45,20 @@ def prompter(prompts):
         print(f"Predicted Mask {i+1}:", mask.shape)
         red_mask = np.zeros_like(image)
         red_mask[:, :, 0] = mask.astype(np.uint8) * 255  # Apply the red channel
-        red_mask = Image.fromarray(red_mask)
+        red_mask = PILImage.fromarray(red_mask)
 
         # Convert the original image to a PIL image
-        original_image = Image.fromarray(image)
+        original_image = PILImage.fromarray(image)
 
         # Blend the original image with the red mask
-        blended_image = Image.blend(original_image, red_mask, alpha=0.5)
+        blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
 
         # Add the blended image to the list
         overlay_images.append(blended_image)
 
-    global IMAGE, MASKS
-
+    global IMAGE, MASKS, MASKED_IMAGES
     IMAGE, MASKS = image, masks
+    MASKED_IMAGES = [np.array(img) for img in overlay_images]
 
     return overlay_images[0], overlay_images[1], overlay_images[2], masks
 
@@ -80,82 +81,63 @@ def save_selected_mask(image, mask, output_dir="output"):
 
     os.makedirs(output_dir, exist_ok=True)
 
-    # Generate a unique UUID for the folder name
     folder_id = str(uuid4())
 
-    # Create a path for the new folder
     folder_path = os.path.join(output_dir, folder_id)
 
-    # Ensure the folder is created
     os.makedirs(folder_path, exist_ok=True)
 
-    # Define the paths for saving the image and mask
-    image_path = os.path.join(folder_path, "image.npy")
-    mask_path = os.path.join(folder_path, "mask.npy")
-
-    # Save the image and mask to the respective paths
-    with open(image_path, "wb") as f:
-        np.save(f, IMAGE)
+    data_path = os.path.join(folder_path, "data.parquet")
+
+    data = {
+        "image": IMAGE,
+        "masked_image": MASKED_IMAGES[INDEX],
+        "mask": MASKS[INDEX],
+    }
+
+    features = Features(
+        {
+            "image": Image(),
+            "masked_image": Image(),
+            "mask": Array2D(
+                dtype="int64", shape=(MASKS[INDEX].shape[0], MASKS[INDEX].shape[1])
+            ),
+        }
+    )
 
-    with open(mask_path, "wb") as f:
-        np.save(f, MASKS[INDEX])
+    ds = Dataset.from_list([data], features=features)
+    ds.to_parquet(data_path)
 
-        # Upload the folder to the Hugging Face Hub
     upload_folder(
         folder_path=output_dir,
-        # path_in_repo=path_in_repo,
-        repo_id="amaye15/object-segmentation",
+        repo_id=DESTINATION_DS,
         repo_type="dataset",
-        # ignore_patterns="**/logs/*.txt",  # Adjust this if needed
     )
 
     shutil.rmtree(folder_path)
 
-    return f"Image and mask saved to {folder_path}."
+    iframe_code = "Success - Check out the 'Results' tab."
 
+    return iframe_code
 
-def save_dataset_name(key, dataset_name):
-    global GLOBALS
-    GLOBALS[key] = dataset_name
+    # time.sleep(5)
 
-    iframe_code = f"""
-    <iframe
-      src="https://huggingface.co/datasets/{dataset_name}/embed/viewer/default/train"
-      frameborder="0"
-      width="100%"
-      height="560px"
-    ></iframe>
-    """
-    return f"Huggingface Dataset: {dataset_name}", iframe_code
+    # # 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>
+    # """
 
 
 # Define the Gradio Blocks app
 with gr.Blocks() as demo:
-    with gr.Tab("Setup"):
-        with gr.Row():
-            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,
-                )
-
-    with gr.Tab("Object Mask - Point Prompt"):
+    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."
@@ -185,13 +167,14 @@ with gr.Blocks() as demo:
             # selected_mask_output = gr.Image(show_label=False)
 
         save_button = gr.Button("Save Selected Mask and Image")
-        save_message = gr.Textbox(visible=False)
+        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,
         )
 
         # Define the action triggered by mask selection
@@ -205,9 +188,42 @@ with gr.Blocks() as demo:
         save_button.click(
             fn=save_selected_mask,
             inputs=[gr.State(), gr.State()],
-            outputs=save_message,
+            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,
+            #     )
 
 # Launch the Gradio app
 demo.launch()

check.py

@@ -1,10 +0,0 @@
-import numpy as np
-import matplotlib.pyplot as plt
-
-# Load the image data from the .npy file
-image = np.load("/Users/andrewmayes/Dev/image/image.npy")
-
-# Display the image using matplotlib
-plt.imshow(image)
-plt.axis("off")  # Turn off the axis labels
-plt.show()  # Show the image

requirements.txt

@@ -2,5 +2,7 @@ gradio
 gradio-image-prompter
 huggingface-hub
 Pillow
-opencv-python
 git+https://github.com/facebookresearch/segment-anything-2.git
+pyarrow
+fastparquet
+datasets