Update app.py

app.py

@@ -8,43 +8,35 @@ from PIL import Image
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
+# Remove all CUDA-specific configurations
+torch.autocast(device_type="cpu", dtype=torch.float32).__enter__()
+
 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}")
-    
     transparent_background = Image.open(first_frame_path).convert('RGBA')
     w, h = transparent_background.size
-    
     fraction = 0.02
     radius = int(fraction * min(w, h))
-    
     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)
-
     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
 
-# Remove all CUDA-specific configurations
-torch.autocast(device_type="cpu", dtype=torch.float32).__enter__()
-
 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)
@@ -54,9 +46,9 @@ def show_mask(mask, ax, random_color=False, borders=True):
     mask = mask.astype(np.uint8)
     mask_image =  mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
     if borders:
-        contours, _ = cv2.findContours(mask,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 
+        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) 
+        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):
@@ -73,65 +65,82 @@ def show_box(box, ax):
 def show_masks(image, masks, scores, point_coords=None, box_coords=None, input_labels=None, borders=True):
     combined_images = []  
     mask_images = []      
-
     for i, (mask, score) in enumerate(zip(masks, scores)):
         plt.figure(figsize=(10, 10))
         plt.imshow(image)
         show_mask(mask, plt.gca(), borders=borders)
         plt.axis('off')
-
         combined_filename = f"combined_image_{i+1}.jpg"
         plt.savefig(combined_filename, format='jpg', bbox_inches='tight')
         combined_images.append(combined_filename)
         plt.close()
-
         mask_image = np.zeros_like(image, dtype=np.uint8)
         mask_layer = (mask > 0).astype(np.uint8) * 255
         for c in range(3):
             mask_image[:, :, c] = mask_layer
-
         mask_filename = f"mask_image_{i+1}.png"
         Image.fromarray(mask_image).save(mask_filename)
         mask_images.append(mask_filename)
-
     return combined_images, mask_images
 
-def sam_process(input_image, checkpoint, tracking_points, trackings_input_label):
-    image = Image.open(input_image)
-    image = np.array(image.convert("RGB"))
+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:
+        return cv2.erode(mask, kernel, iterations=1)
+
+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.array(image.convert("RGB"))
     checkpoint_map = {
         "tiny": ("./checkpoints/sam2_hiera_tiny.pt", "sam2_hiera_t.yaml"),
         "small": ("./checkpoints/sam2_hiera_small.pt", "sam2_hiera_s.yaml"),
         "base-plus": ("./checkpoints/sam2_hiera_base_plus.pt", "sam2_hiera_b+.yaml"),
         "large": ("./checkpoints/sam2_hiera_large.pt", "sam2_hiera_l.yaml")
     }
-    
     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(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]
-
-    results, mask_results = show_masks(image, masks, scores, 
-                                     point_coords=input_point, 
-                                     input_labels=input_label, 
+    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:
@@ -149,36 +158,43 @@ with gr.Blocks() as demo:
                     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, default=0, label="Expand/Contract (pixels)")
+                    expand = gr.Radio(["Expand", "Contract"], default="Expand", label="Action")
+                with gr.Row():
+                    feathering_enabled = gr.Checkbox(default=False, label="Enable Feathering")
+                    feather_size = gr.Slider(minimum=1, maximum=50, default=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],
+        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_api=False, show_error=True)