app.py

import gradio as gr
from gradio_image_prompter import ImagePrompter
import torch
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 PIL import Image as PILImage
from datasets import Dataset, Features, Array2D, Image
import shutil
import random
from datasets import load_dataset

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"


token = os.getenv("TOKEN")
if token:
    login(token)

IMAGE = None
MASKS = None
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
    points = prompts["points"]  # Get the points from prompts

    # 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 = PILImage.fromarray(red_mask)

        # Convert the original image to a PIL image
        original_image = PILImage.fromarray(image)

        # Blend the original image with the red mask
        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, 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


def select_mask(
    selected_mask_index,
    mask1,
    mask2,
    mask3,
):
    masks = [mask1, mask2, mask3]
    global INDEX
    INDEX = selected_mask_index
    return masks[selected_mask_index]


def save_selected_mask(image, mask, output_dir="output"):

    output_dir = os.path.join(os.getcwd(), output_dir)

    os.makedirs(output_dir, exist_ok=True)

    folder_id = str(uuid4())

    folder_path = os.path.join(output_dir, folder_id)

    os.makedirs(folder_path, exist_ok=True)

    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])
            ),
        }
    )

    ds = Dataset.from_list([data], features=features)
    ds.to_parquet(data_path)

    upload_folder(
        folder_path=output_dir,
        repo_id=DESTINATION_DS,
        repo_type="dataset",
    )

    shutil.rmtree(folder_path)

    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


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.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)

    save_button = gr.Button("Save Selected Mask and Image")
    iframe_display = gr.Markdown()

    # 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
    )

    # 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,
#     )