Update sam2_mask.py

sam2_mask.py

@@ -12,16 +12,16 @@ 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 POINTS: {tracking_points.value}")
+    tracking_points.append(evt.index)
+    print(f"TRACKING POINTS: {tracking_points}")
     if point_type == "include":
-        trackings_input_label.value.append(1)
+        trackings_input_label.append(1)
     elif point_type == "exclude":
-        trackings_input_label.value.append(0)
-    print(f"TRACKING INPUT LABELS: {trackings_input_label.value}")
+        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
@@ -30,16 +30,16 @@ def get_point(point_type, tracking_points, trackings_input_label, first_frame_pa
     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)
+    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, track, radius, (255, 0, 0, 255), -1)
+            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
-    
+
 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)
@@ -99,21 +99,21 @@ def show_masks(image, masks, scores, point_coords=None, box_coords=None, input_l
         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)
-    if not input_point.size or not input_label.size:
-        print("No points or labels provided, returning None")
-        return None
-    masks, scores, _= predictor.predict(input_point, input_label, multimask_output=False)
+    masks, scores, _ = predictor.predict(input_point, input_label, multimask_output=False)
     sorted_indices = np.argsort(scores)[::-1]
     masks = masks[sorted_indices]
     # Generate mask image
@@ -129,12 +129,14 @@ def create_sam2_tab():
     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
         
         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")
         
@@ -144,16 +146,19 @@ def create_sam2_tab():
             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],