Update sam2_mask.py

sam2_mask.py

@@ -1,59 +1,42 @@
-import spaces
 import gradio as gr
 import os
-os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"
 import torch
 import numpy as np
 import cv2
+import huggingface_hub
 import matplotlib.pyplot as plt
-from PIL import Image, ImageFilter
+from PIL import Image
 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
+
+# Remove all CUDA-specific configurations
+torch.autocast(device_type="cpu", dtype=torch.float32).__enter__()
 
 def preprocess_image(image):
-    return 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}")
-    
-    # Extract x, y coordinates from evt.index
-    x, y = evt.index
-    
-    # Add the point as [x, y]
-    tracking_points.append([x, y])
-    print(f"TRACKING POINTS: {tracking_points}")
-    
+    tracking_points.value.append(evt.index)
+    print(f"TRACKING POINT: {tracking_points.value}")
     if point_type == "include":
-        trackings_input_label.append(1)
+        trackings_input_label.value.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
+        trackings_input_label.value.append(0)
+    print(f"TRACKING INPUT LABEL: {trackings_input_label.value}")
     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
+    fraction = 0.02
     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)
+    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, tuple(track), radius, (255, 0, 0, 255), -1)
-    
-    # Convert the transparent layer back to an image
+            cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
     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
 
 def show_mask(mask, ax, random_color=False, borders=True):
@@ -63,156 +46,159 @@ def show_mask(mask, ax, random_color=False, borders=True):
         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)
+    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.findContours(mask,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
         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]
+def show_points(coords, labels, ax, marker_size=200):
+    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_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)
+    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))
+    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
+    combined_images = []  
+    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)
+        show_mask(mask, plt.gca(), borders=borders)
         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
+        plt.close()
         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
+        for c in range(3):
             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}")
-    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)
-    input_label = np.array(labels)
-    
-    try:
-        masks, scores, _ = predictor.predict(input_point, input_label, multimask_output=False)
-    except Exception as e:
-        print(f"Error during prediction: {e}")
-        return None
-    
-    sorted_indices = np.argsort(scores)[::-1]
-    masks = masks[sorted_indices]
-    
-    if masks and len(masks) > 0:
-        mask = masks[0] * 255
-        mask_image = Image.fromarray(mask.astype(np.uint8))
-        return mask_image
+def expand_contract_mask(mask, px, expand=True):
+    kernel = np.ones((px, px), np.uint8)
+    if expand:
+        return cv2.dilate(mask, kernel, iterations=1)
     else:
-        print("No masks generated, returning None")
-        return None
+        return cv2.erode(mask, kernel, iterations=1)
 
-def create_sam2_tab():
-    first_frame = gr.State()  # Tracks original image
+def feather_mask(mask, feather_size=10):
+    feathered_mask = mask.copy()
+    Feathered_region = mask > 0
+    Feathered_region = cv2.dilate(Feathered_region.astype(np.uint8), np.ones((feather_size, feather_size), np.uint8), iterations=1)
+    Feathered_region = Feathered_region & (~mask.astype(bool))
+    
+    for i in range(1, feather_size + 1):
+        weight = i / (feather_size + 1)
+        feathered_mask[Feathered_region] = feathered_mask[Feathered_region] * (1 - weight) + weight
+
+    return feathered_mask
+
+def process_mask(mask, expand_contract_px, expand, feathering_enabled, feather_size):
+    if expand_contract_px > 0:
+        mask = expand_contract_mask(mask, expand_contract_px, expand)
+    if feathering_enabled:
+        mask = feather_mask(mask, feather_size)
+    return mask
+
+def sam_process(input_image, checkpoint, tracking_points, trackings_input_label, expand_contract_px, expand, feathering_enabled, feather_size):
+    image = Image.open(input_image)
+    image = np.array(image.convert("RGB"))
+    sam21_hfmap = {
+        "tiny": "facebook/sam2.1-hiera-tiny",
+        "small": "facebook/sam2.1-hiera-small",
+        "base-plus": "facebook/sam2.1-hiera-base-plus",
+        "large": "facebook/sam2.1-hiera-large",
+    }
+    # sam2_checkpoint, model_cfg = checkpoint_map[checkpoint]
+    # Use CPU for both model and computations
+    # sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cpu")
+    predictor = SAM2ImagePredictor.from_pretrained(sam21_hfmap[checkpoint], device="cpu")
+
+    # predictor = SAM2ImagePredictor(sam2_model)
+    predictor.set_image(image)
+    input_point = np.array(tracking_points.value)
+    input_label = np.array(trackings_input_label.value)
+    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]
+    processed_masks = []
+    for mask in masks:
+        processed_mask = process_mask(mask, expand_contract_px, expand, feathering_enabled, feather_size)
+        processed_masks.append(processed_mask)
+    results, mask_results = show_masks(image, processed_masks, scores,
+                                     point_coords=input_point,
+                                     input_labels=input_label,
+                                     borders=True)
+    return results[0], mask_results[0]
+
+with gr.Blocks() as demo:
+    first_frame_path = gr.State()
     tracking_points = gr.State([])
     trackings_input_label = gr.State([])
-    
     with gr.Column():
-        gr.Markdown("# SAM2 Image Predictor")
-        gr.Markdown("1. Upload your image\n2. Click points to mask\n3. Submit")
-        
-        points_map = gr.Image(label="Points Map", type="pil", interactive=True)
-        input_image = gr.Image(type="pil", visible=False)  # Original image
-        
+        gr.Markdown("# SAM2 Image Predictor (CPU Version)")
+        gr.Markdown("This version runs entirely on CPU")
         with gr.Row():
-            point_type = gr.Radio(["include", "exclude"], value="include", label="Point Type")
-            clear_button = gr.Button("Clear Points")
-        
-        submit_button = gr.Button("Submit")
-        output_image = gr.Image("Segmented Output")
-        
-        # Event handlers
-        points_map.upload(
-            lambda img: (img, img, [], []),
-            inputs=points_map,
-            outputs=[input_image, first_frame, tracking_points, trackings_input_label]
-        )
-        
-        clear_button.click(
-            lambda img: ([], [], img),
-            inputs=first_frame,
-            outputs=[tracking_points, trackings_input_label, points_map]
-        )
-        
-        points_map.select(
-            get_point,
-            inputs=[point_type, tracking_points, trackings_input_label, first_frame],
-            outputs=[tracking_points, trackings_input_label, points_map]
-        )
-        
-        submit_button.click(
-            sam_process,
-            inputs=[input_image, tracking_points, trackings_input_label],
-            outputs=output_image
-        )
-    
-    return input_image, points_map, output_image
+            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="base-plus")
+                with gr.Row():
+                    expand_contract_px = gr.Slider(minimum=0, maximum=50, value=0, label="Expand/Contract (pixels)")
+                    expand = gr.Radio(["Expand", "Contract"], value="Expand", label="Action")
+                with gr.Row():
+                    feathering_enabled = gr.Checkbox(value=False, label="Enable Feathering")
+                    feather_size = gr.Slider(minimum=1, maximum=50, value=10, label="Feathering Size", visible=False)
+                submit_btn = gr.Button("Submit")
+            with gr.Column():
+                output_result = gr.Image()
+                output_result_mask = gr.Image()
+    clear_points_btn.click(
+        fn=preprocess_image,
+        inputs=input_image,
+        outputs=[first_frame_path, tracking_points, trackings_input_label, points_map],
+        queue=False
+    )
+    points_map.upload(
+        fn=preprocess_image,
+        inputs=[points_map],
+        outputs=[first_frame_path, tracking_points, trackings_input_label, input_image],
+        queue=False
+    )
+    points_map.select(
+        fn=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, expand_contract_px, expand, feathering_enabled, feather_size],
+        outputs=[output_result, output_result_mask]
+    )
+    feathering_enabled.change(
+        fn=lambda enabled: gr.update(visible=enabled),
+        inputs=[feathering_enabled],
+        outputs=[feather_size]
+    )
+
+demo.launch(show_error=True)