Update sam2_mask.py

sam2_mask.py

@@ -12,42 +12,33 @@ 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):
+    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}")
-
+    print(f"TRACKING POINTS: {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}")
-    
+    print(f"TRACKING INPUT LABELS: {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
-
+    
 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)
@@ -107,26 +98,27 @@ 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}")
     # 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)
-    masks, scores, _ = predictor.predict(input_point, input_label, multimask_output=False)
+    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)
     sorted_indices = np.argsort(scores)[::-1]
     masks = masks[sorted_indices]
-    
     # Generate mask image
     mask = masks[0] * 255
     mask_image = Image.fromarray(mask.astype(np.uint8))
     return mask_image
-    # sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
-    # predictor = SAM2ImagePredictor(sam2_model)
-
 
 def create_sam2_tab():
     first_frame = gr.State()  # Tracks original image