Update sam2_mask.py

sam2_mask.py

@@ -12,28 +12,41 @@ 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}")
 
-def get_point(point_type, tracking_points, trackings_input_label, original_image, evt: gr.SelectData):
-    x, y = evt.index
-    tracking_points.append((x, y))
-    trackings_input_label.append(1 if point_type == "include" else 0)
+    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
     
-    # Redraw all points on original image
-    w, h = original_image.size
-    radius = int(min(w, h) * 0.02)
-    img = original_image.convert("RGBA")
-    draw = ImageDraw.Draw(img)
-    for i, (cx, cy) in enumerate(tracking_points):
-        color = (0, 255, 0, 255) if trackings_input_label[i] == 1 else (255, 0, 0, 255)
-        draw.ellipse([cx-radius, cy-radius, cx+radius, cy+radius], fill=color)
-    return tracking_points, trackings_input_label, img
+    # 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)
 
-# 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
+    # 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: