Version 1 - Working

app.py

@@ -1,183 +1,211 @@
-# import gradio as gr
-# from gradio_image_prompter import ImagePrompter
+import gradio as gr
+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
+import shutil
 
-# import os
-# import torch
+MODEL = "facebook/sam2-hiera-large"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
 
 
-# def prompter(prompts):
-#     image = prompts["image"]  # Get the image from prompts
-#     points = prompts["points"]  # Get the points from prompts
+GLOBALS = {}
 
-#     # Print the collected inputs for debugging or logging
-#     print("Image received:", image)
-#     print("Points received:", points)
 
-#     import torch
-#     from sam2.sam2_image_predictor import SAM2ImagePredictor
+IMAGE = None
+MASKS = None
+INDEX = None
 
-#     device = torch.device("cpu")
 
-#     predictor = SAM2ImagePredictor.from_pretrained(
-#         "facebook/sam2-hiera-base-plus", device=device
-#     )
+def prompter(prompts):
 
-#     with torch.inference_mode():
-#         predictor.set_image(image)
-#         # masks, _, _ = predictor.predict([[point[0], point[1]] for point in points])
-#         input_point = [[point[0], point[1]] for point in points]
-#         input_label = [1]
-#         masks, _, _ = predictor.predict(
-#             point_coords=input_point, point_labels=input_label
-#         )
-#     print("Predicted Mask:", masks)
+    image = np.array(prompts["image"])  # Convert the image to a numpy array
+    points = prompts["points"]  # Get the points from prompts
 
-#     return image, points
+    # Perform inference with multimask_output=True
+    with torch.inference_mode():
+        PREDICTOR.set_image(image)
+        input_point = [[point[0], point[1]] for point in points]
+        input_label = [1] * len(points)  # Assuming all points are foreground
+        masks, _, _ = PREDICTOR.predict(
+            point_coords=input_point, point_labels=input_label, multimask_output=True
+        )
 
+    # Prepare individual images with separate overlays
+    overlay_images = []
+    for i, mask in enumerate(masks):
+        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)
 
-# # Define the Gradio interface
-# demo = gr.Interface(
-#     fn=prompter,  # Use the custom prompter function
-#     inputs=ImagePrompter(
-#         show_label=False
-#     ),  # ImagePrompter for image input and point selection
-#     outputs=[
-#         gr.Image(show_label=False),  # Display the image
-#         gr.Dataframe(label="Points"),  # Display the points in a DataFrame
-#     ],
-#     title="Image Point Collector",
-#     description="Upload an image, click on it, and get the coordinates of the clicked points.",
-# )
+        # Convert the original image to a PIL image
+        original_image = Image.fromarray(image)
 
-# # Launch the Gradio app
-# demo.launch()
+        # Blend the original image with the red mask
+        blended_image = Image.blend(original_image, red_mask, alpha=0.5)
 
+        # Add the blended image to the list
+        overlay_images.append(blended_image)
 
-# import gradio as gr
-# from gradio_image_prompter import ImagePrompter
-# import torch
-# from sam2.sam2_image_predictor import SAM2ImagePredictor
+    global IMAGE, MASKS
 
+    IMAGE, MASKS = image, masks
 
-# def prompter(prompts):
-#     image = prompts["image"]  # Get the image from prompts
-#     points = prompts["points"]  # Get the points from prompts
+    return overlay_images[0], overlay_images[1], overlay_images[2], masks
 
-#     # Print the collected inputs for debugging or logging
-#     print("Image received:", image)
-#     print("Points received:", points)
 
-#     device = torch.device("cpu")
+def select_mask(
+    selected_mask_index,
+    mask1,
+    mask2,
+    mask3,
+):
+    masks = [mask1, mask2, mask3]
+    global INDEX
+    INDEX = selected_mask_index
+    return masks[selected_mask_index]
 
-#     # Load the SAM2ImagePredictor model
-#     predictor = SAM2ImagePredictor.from_pretrained(
-#         "facebook/sam2-hiera-base-plus", device=device
-#     )
 
-#     # Perform inference
-#     with torch.inference_mode():
-#         predictor.set_image(image)
-#         input_point = [[point[0], point[1]] for point in points]
-#         input_label = [1] * len(points)  # Assuming all points are foreground
-#         masks, _, _ = predictor.predict(
-#             point_coords=input_point, point_labels=input_label
-#         )
+def save_selected_mask(image, mask, output_dir="output"):
 
-#     # The masks are returned as a list of numpy arrays
-#     print("Predicted Mask:", masks)
+    output_dir = os.path.join(os.getcwd(), output_dir)
 
-#     # Assuming there's only one mask returned, you can adjust if there are multiple
-#     predicted_mask = masks[0]
+    os.makedirs(output_dir, exist_ok=True)
 
-#     print(len(image))
+    # Generate a unique UUID for the folder name
+    folder_id = str(uuid4())
 
-#     print(len(predicted_mask))
+    # Create a path for the new folder
+    folder_path = os.path.join(output_dir, folder_id)
 
-#     # Create annotations for AnnotatedImage
-#     annotations = [(predicted_mask, "Predicted Mask")]
+    # Ensure the folder is created
+    os.makedirs(folder_path, exist_ok=True)
 
-#     return image, annotations
+    # 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)
 
-# # Define the Gradio interface
-# demo = gr.Interface(
-#     fn=prompter,  # Use the custom prompter function
-#     inputs=ImagePrompter(
-#         show_label=False
-#     ),  # ImagePrompter for image input and point selection
-#     outputs=gr.AnnotatedImage(),  # Display the image with the predicted mask
-#     title="Image Point Collector with Mask Overlay",
-#     description="Upload an image, click on it, and get the predicted mask overlayed on the image.",
-# )
+    with open(mask_path, "wb") as f:
+        np.save(f, MASKS[INDEX])
 
-# # Launch the Gradio app
-# demo.launch()
+        # 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_type="dataset",
+        # ignore_patterns="**/logs/*.txt",  # Adjust this if needed
+    )
 
+    shutil.rmtree(folder_path)
 
-import gradio as gr
-from gradio_image_prompter import ImagePrompter
-import torch
-import numpy as np
-from sam2.sam2_image_predictor import SAM2ImagePredictor
-from PIL import Image
+    return f"Image and mask saved to {folder_path}."
 
 
-def prompter(prompts):
-    image = np.array(prompts["image"])  # Convert the image to a numpy array
-    points = prompts["points"]  # Get the points from prompts
+def save_dataset_name(key, dataset_name):
+    global GLOBALS
+    GLOBALS[key] = dataset_name
 
-    # Print the collected inputs for debugging or logging
-    print("Image received:", image)
-    print("Points received:", points)
+    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
 
-    device = torch.device("cpu")
 
-    # Load the SAM2ImagePredictor model
-    predictor = SAM2ImagePredictor.from_pretrained(
-        "facebook/sam2-hiera-base-plus", device=device
-    )
+# 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()
 
-    # Perform inference with multimask_output=True
-    with torch.inference_mode():
-        predictor.set_image(image)
-        input_point = [[point[0], point[1]] for point in points]
-        input_label = [1] * len(points)  # Assuming all points are foreground
-        masks, _, _ = predictor.predict(
-            point_coords=input_point, point_labels=input_label, multimask_output=True
-        )
+                source.change(
+                    save_dataset_name,
+                    inputs=(gr.State("source_dataset"), source),
+                    outputs=(source_display, iframe_display),
+                )
 
-    # Prepare individual images with separate overlays
-    overlay_images = []
-    for i, mask in enumerate(masks):
-        print(f"Predicted Mask {i+1}:", mask)
-        red_mask = np.zeros_like(image)
-        red_mask[:, :, 0] = mask.astype(np.uint8) * 255  # Apply the red channel
-        red_mask = Image.fromarray(red_mask)
+            with gr.Column():
 
-        # Convert the original image to a PIL image
-        original_image = Image.fromarray(image)
+                destination = gr.Textbox(label="Destination Dataset")
+                destination_display = gr.Markdown()
 
-        # Blend the original image with the red mask
-        blended_image = Image.blend(original_image, red_mask, alpha=0.5)
+                destination.change(
+                    save_dataset_name,
+                    inputs=(gr.State("destination_dataset"), destination),
+                    outputs=destination_display,
+                )
 
-        # Add the blended image to the list
-        overlay_images.append(blended_image)
+    with gr.Tab("Object Mask - 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():
+                # 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")
+        save_message = gr.Textbox(visible=False)
+
+        # 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()],
+        )
 
-    return overlay_images
-
-
-# Define the Gradio interface
-demo = gr.Interface(
-    fn=prompter,  # Use the custom prompter function
-    inputs=ImagePrompter(
-        show_label=False
-    ),  # ImagePrompter for image input and point selection
-    outputs=[
-        gr.Image(show_label=False) for _ in range(3)
-    ],  # Display up to 3 overlay images
-    title="Image Point Collector with Multiple Separate Mask Overlays",
-    description="Upload an image, click on it, and get each predicted mask overlaid separately in red on individual images.",
-)
+        # 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=save_message,
+            show_progress=True,
+        )
 
 # Launch the Gradio app
 demo.launch()

check.py

@@ -0,0 +1,10 @@
+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

create_repo.py

@@ -0,0 +1,9 @@
+from huggingface_hub import HfApi
+
+# Initialize the API
+api = HfApi()
+
+# Create a new dataset repository
+repo_url = api.create_repo(repo_id="amaye15/object-segmentation", repo_type="dataset")
+
+print(f"Dataset repository created: {repo_url}")

requirements.txt

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