major: add sam2 masking tab

.gitignore

@@ -1 +1,2 @@
 .env
+__pycache__/*

app.py

@@ -9,13 +9,10 @@ from controlnet_union import ControlNetModel_Union
 from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
-from sam2.sam2_image_predictor import SAM2ImagePredictor
+from sam2_mask import create_sam2_tab, sam_process
+
+#from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-# predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large", device="cuda")
-# print(predictor)
-# with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-#     predictor.set_image(<your_image>)
-#     masks, _, _ = predictor.predict(<input_prompts>)
 
 MODELS = {
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
@@ -479,15 +476,6 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                                     )
                                 with gr.Column():
                                     preview_button = gr.Button("Preview alignment and mask")
-                        gr.Examples(
-                            examples=[
-                                ["./examples/example_1.webp", 1280, 720, "Middle"],
-                                ["./examples/example_2.jpg", 1440, 810, "Left"],
-                                ["./examples/example_3.jpg", 1024, 1024, "Top"],
-                                ["./examples/example_3.jpg", 1024, 1024, "Bottom"],
-                            ],
-                            inputs=[input_image_outpaint, width_slider, height_slider, alignment_dropdown],
-                        )
                     with gr.Column():
                         result_outpaint = ImageSlider(
                             interactive=False,
@@ -496,6 +484,9 @@ 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("Misc"):
             with gr.Column():
                 clear_cache_button = gr.Button("Clear CUDA Cache")

sam2_mask.py

@@ -0,0 +1,201 @@
+import gradio as gr
+import os
+os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"
+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
+
+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.value.append(evt.index)
+    print(f"TRACKING POINT: {tracking_points.value}")
+    if point_type == "include":
+        trackings_input_label.value.append(1)
+    elif point_type == "exclude":
+        trackings_input_label.value.append(0)
+    print(f"TRACKING INPUT LABEL: {trackings_input_label.value}")
+    # 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.value):
+        if trackings_input_label.value[index] == 1:
+            cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
+        else:
+            cv2.circle(transparent_layer, 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)
+    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
+
+def sam_process(input_image, tracking_points, trackings_input_label):
+    image = Image.open(input_image)
+    image = np.array(image.convert("RGB"))
+    # if checkpoint == "tiny":
+    #     sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt"
+    #     model_cfg = "sam2_hiera_t.yaml"
+    # elif checkpoint == "small":
+    #     sam2_checkpoint = "./checkpoints/sam2_hiera_small.pt"
+    #     model_cfg = "sam2_hiera_s.yaml"
+    # elif checkpoint == "base-plus":
+    #     sam2_checkpoint = "./checkpoints/sam2_hiera_base_plus.pt"
+    #     model_cfg = "sam2_hiera_b+.yaml"
+    # elif checkpoint == "large":
+    #     sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt"
+    #     model_cfg = "sam2_hiera_l.yaml"
+    predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
+    # print(predictor)
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        predictor.set_image(image)
+        input_point = np.array(tracking_points.value)
+        input_label = np.array(trackings_input_label.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_ind = np.argsort(scores)[::-1]
+        masks = masks[sorted_ind]
+        scores = scores[sorted_ind]
+        logits = logits[sorted_ind]
+        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]
+    # sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
+    # predictor = SAM2ImagePredictor(sam2_model)
+
+
+
+def create_sam2_tab():
+    first_frame_path = gr.State()
+    tracking_points = gr.State([])
+    trackings_input_label = gr.State([])
+    with gr.Column():
+        gr.Markdown("# SAM2 Image Predictor")
+        gr.Markdown("This is a simple demo for image segmentation with SAM2.")
+        gr.Markdown("""Instructions:
+        1. Upload your image
+        2. With 'include' point type selected, Click on the object to mask
+        3. Switch to 'exclude' point type if you want to specify an area to avoid
+        4. Submit !
+        """)
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="input image", interactive=False, type="filepath", visible=False)                 
+                points_map = gr.Image(
+                    label="points map",
+                    type="filepath",
+                    interactive=True
+                )
+                with gr.Row():
+                    point_type = gr.Radio(label="point type", choices=["include", "exclude"], value="include")
+                    clear_points_btn = gr.Button("Clear Points")
+                checkpoint = gr.Dropdown(label="Checkpoint", choices=["tiny", "small", "base-plus", "large"], value="tiny")
+                submit_btn = gr.Button("Submit")
+            with gr.Column():
+                output_result = gr.Image()
+                output_result_mask = gr.Image()
+        clear_points_btn.click(
+            fn=admin_preprocess_image,
+            inputs=input_image,
+            outputs=[first_frame_path, tracking_points, trackings_input_label, points_map],
+            queue=False
+        )
+        points_map.upload(
+            fn=admin_preprocess_image,
+            inputs=[points_map],
+            outputs=[first_frame_path, tracking_points, trackings_input_label, input_image],
+            queue=False
+        )
+        points_map.select(
+            fn=admin_get_point,
+            inputs=[point_type, tracking_points, trackings_input_label, first_frame_path],
+            outputs=[tracking_points, trackings_input_label, points_map],
+            queue=False
+        )
+        submit_btn.click(
+            fn=sam_process,
+            inputs=[input_image, checkpoint, tracking_points, trackings_input_label],
+            outputs=[output_result, output_result_mask]
+        )
+    return input_image, points_map, output_result_mask