kiss: simplifying everything

app.py

@@ -9,10 +9,11 @@ from controlnet_union import ControlNetModel_Union
 from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
-from sam2_mask import create_sam2_tab, sam_process
+from sam2_mask import create_sam2_mask_interface
 
 #from sam2.sam2_image_predictor import SAM2ImagePredictor
 
+sam2_mask_tab = create_sam2_mask_interface()
 
 MODELS = {
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
@@ -486,9 +487,8 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         use_as_input_button_outpaint = gr.Button("Use as Input Image", visible=False)
                         history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
                         preview_image = gr.Image(label="Preview")
-        with gr.TabItem("SAM2 Masking"):
-            input_image, points_map, output_result_mask = create_sam2_tab()
-
+        with gr.TabItem("SAM2 Mask"):
+            sam2_mask_tab
         with gr.TabItem("Misc"):
             with gr.Column():
                 clear_cache_button = gr.Button("Clear CUDA Cache")

sam2_mask.py

@@ -1,204 +1,62 @@
-import spaces
+# K-I-S-S
+
 import gradio as gr
-import os
-os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"
+from gradio_image_prompter import ImagePrompter
+from sam2.sam2_image_predictor import SAM2ImagePredictor
 import torch
 import numpy as np
-import cv2
-import matplotlib.pyplot as plt
-from PIL import Image, ImageFilter
-from sam2.build_sam import build_sam2
-from sam2.sam2_image_predictor import SAM2ImagePredictor
-from gradio_image_prompter import ImagePrompter
-
-def preprocess_image(image):
-    return image, gr.State([]), gr.State([]), image
-
-def get_point(point_type, tracking_points, trackings_input_label, first_frame_path, evt: gr.SelectData):
-    print(f"You selected {evt.value} at {evt.index} from {evt.target}")
-    tracking_points.append(evt.index)
-    print(f"TRACKING POINTS: {tracking_points}")
-    if point_type == "include":
-        trackings_input_label.append(1)
-    elif point_type == "exclude":
-        trackings_input_label.append(0)
-    print(f"TRACKING INPUT LABELS: {trackings_input_label}")
-    # Open the image and get its dimensions
-    transparent_background = Image.open(first_frame_path).convert('RGBA')
-    w, h = transparent_background.size
-    # Define the circle radius as a fraction of the smaller dimension
-    fraction = 0.02  # You can adjust this value as needed
-    radius = int(fraction * min(w, h))
-    # Create a transparent layer to draw on
-    transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
-    for index, track in enumerate(tracking_points):
-        if trackings_input_label[index] == 1:
-            cv2.circle(transparent_layer, tuple(track), radius, (0, 255, 0, 255), -1)
-        else:
-            cv2.circle(transparent_layer, tuple(track), radius, (255, 0, 0, 255), -1)
-    # Convert the transparent layer back to an image
-    transparent_layer = Image.fromarray(transparent_layer, 'RGBA')
-    selected_point_map = Image.alpha_composite(transparent_background, transparent_layer)
-    return tracking_points, trackings_input_label, selected_point_map
-    
-# use bfloat16 for the entire notebook
-torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
-
-if torch.cuda.get_device_properties(0).major >= 8:
-    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-    
-def show_mask(mask, ax, random_color=False, borders=True):
-    if random_color:
-        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+from PIL import Image as PILImage
+
+# Initialize SAM2 predictor
+MODEL = "facebook/sam2-hiera-large"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
+
+def predict_masks(image, points):
+    """Predict a single mask from the image based on selected points."""
+    image_np = np.array(image)
+    points_list = [[point["x"], point["y"]] for point in points]
+    input_labels = [1] * len(points_list)
+
+    with torch.inference_mode():
+        PREDICTOR.set_image(image_np)
+        masks, _, _ = PREDICTOR.predict(
+            point_coords=points_list, point_labels=input_labels, multimask_output=False
+        )
+
+    # Prepare the overlay image
+    red_mask = np.zeros_like(image_np)
+    if masks and len(masks) > 0:
+        red_mask[:, :, 0] = masks[0].astype(np.uint8) * 255  # Apply the red channel
+        red_mask = PILImage.fromarray(red_mask)
+        original_image = PILImage.fromarray(image_np)
+        blended_image = PILImage.blend(original_image, red_mask, alpha=0.5)
+        return np.array(blended_image)
     else:
-        color = np.array([30/255, 144/255, 255/255, 0.6])
-    h, w = mask.shape[-2:]
-    mask = mask.astype(np.uint8)
-    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
-    if borders:
-        import cv2
-        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
-        # Try to smooth contours
-        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
-        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
-    ax.imshow(mask_image)
-
-def show_points(coords, labels, ax, marker_size=375):
-    pos_points = coords[labels == 1]
-    neg_points = coords[labels == 0]
-    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
-    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
-
-def show_box(box, ax):
-    x0, y0 = box[0], box[1]
-    w, h = box[2] - box[0], box[3] - box[1]
-    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))
-
-def show_masks(image, masks, scores, point_coords=None, box_coords=None, input_labels=None, borders=True):
-    combined_images = []  # List to store filenames of images with masks overlaid
-    mask_images = []      # List to store filenames of separate mask images
-    for i, (mask, score) in enumerate(zip(masks, scores)):
-        # ---- Original Image with Mask Overlaid ----
-        plt.figure(figsize=(10, 10))
-        plt.imshow(image)
-        show_mask(mask, plt.gca(), borders=borders)  # Draw the mask with borders
-        if box_coords is not None:
-            show_box(box_coords, plt.gca())
-        if len(scores) > 1:
-            plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
-        plt.axis('off')
-        # Save the figure as a JPG file
-        combined_filename = f"combined_image_{i+1}.jpg"
-        plt.savefig(combined_filename, format='jpg', bbox_inches='tight')
-        combined_images.append(combined_filename)
-        plt.close()  # Close the figure to free up memory
-        # ---- Separate Mask Image (White Mask on Black Background) ----
-        # Create a black image
-        mask_image = np.zeros_like(image, dtype=np.uint8)
-        # The mask is a binary array where the masked area is 1, else 0.
-        # Convert the mask to a white color in the mask_image
-        mask_layer = (mask > 0).astype(np.uint8) * 255
-        for c in range(3):  # Assuming RGB, repeat mask for all channels
-            mask_image[:, :, c] = mask_layer
-        # Save the mask image
-        mask_filename = f"mask_image_{i+1}.png"
-        Image.fromarray(mask_image).save(mask_filename)
-        mask_images.append(mask_filename)
-        plt.close()  # Close the figure to free up memory
-    return combined_images, mask_images
-
-@spaces.GPU()
-def sam_process(original_image, points, labels):
-    
-    print(f"Points: {points}")
-    print(f"Labels: {labels}")
-    image = Image.open(original_image)
-    image = np.array(image.convert("RGB"))
-    
-    if not points or not labels:
-        print("No points or labels provided, returning None")
-        return None
-    # Convert image to numpy array for SAM2 processing
-    # image = np.array(original_image)
-    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
-    predictor.set_image(image)
-    input_point = np.array(points.value)
-    input_label = np.array(labels.value)
-    
-    print(predictor._features["image_embed"].shape, predictor._features["image_embed"][-1].shape)
-
-    masks, scores, logits = predictor.predict(
-        point_coords=input_point,
-        point_labels=input_label,
-        multimask_output=False,
-    )
-    sorted_indices = np.argsort(scores)[::-1]
-    masks = masks[sorted_indices]
-    scores = scores[sorted_indices]
-    logits = logits[sorted_indices]
-    print(masks.shape)
-
-    results, mask_results = show_masks(image, masks, scores, point_coords=input_point, input_labels=input_label, borders=True)
-    print(results)
-    
-    return results[0], mask_results[0]
-
-def create_sam2_tab():
-    first_frame = gr.State()  # Tracks original image
-    tracking_points = gr.State([])
-    trackings_input_label = gr.State([])
-    
-    with gr.Column():
+        return image_np
+
+def create_sam2_mask_interface():
+    """Create the Gradio interface for SAM2 mask generation."""
+    with gr.Blocks() as sam2_mask_tab:
+        gr.Markdown("# Object Segmentation with SAM2")
+        gr.Markdown(
+            """
+            This application utilizes **Segment Anything V2 (SAM2)** to allow you to upload an image and interactively generate a segmentation mask based on multiple points you select on the image.
+            """
+        )
         with gr.Row():
             with gr.Column():
-                sam_input_image = gr.Image(label="input image", interactive=False, type="filepath", visible=False)           
-                img_prompter = ImagePrompter(show_label=False)
-                points_map = gr.Image(
-                    label="points map", 
-                    type="filepath",
-                    interactive=True
-                )
-                with gr.Row():
-                    point_type = gr.Radio(["include", "exclude"], value="include", label="Point Type")
-                    clear_button = gr.Button("Clear Points")
+                upload_image_input = ImagePrompter(show_label=False)
                 submit_button = gr.Button("Submit")
-
             with gr.Column():
-                output_image = gr.Image("Segmented Output")
-                output_result_mask = gr.Image()
-                prompted_output = gr.Image(show_label=False)
-                prompted_data = gr.Dataframe(label="Points")
-
-    # lambda prompts: (prompts["image"], prompts["points"]),
-    # Event handlers
-    points_map.upload(
-        fn = preprocess_image,
-        inputs = [points_map], 
-        # outputs=[sam_input_image, first_frame, tracking_points, trackings_input_label],
-        outputs = [first_frame, tracking_points, trackings_input_label, sam_input_image],
-        queue=False
-    )
-    
-    clear_button.click(
-        lambda img: ([], [], img),
-        inputs=first_frame,
-        outputs=[tracking_points, trackings_input_label, points_map],
-        queue=False
-    )
-    
-    points_map.select(
-        get_point,
-        inputs=[point_type, tracking_points, trackings_input_label, first_frame],
-        outputs=[tracking_points, trackings_input_label, points_map],
-        queue = False
-    )
-    
-    submit_button.click(
-        sam_process,
-        inputs=[sam_input_image, tracking_points, trackings_input_label],
-        outputs = [output_image, output_result_mask]
-    )
-    
-    return sam_input_image, points_map, output_image
+                image_output = gr.Image(label="Segmented Image", type="pil").style(height=400)
+        
+        # Define the action triggered by the submit button
+        submit_button.click(
+            fn=predict_masks,
+            inputs=[upload_image_input.image, upload_image_input.points],
+            outputs=image_output,
+            show_progress=True,
+        )
+    
+    return sam2_mask_tab