attemp was made

sam2_mask.py

@@ -13,32 +13,20 @@ 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}")
-    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
-    # 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 get_point(point_type, tracking_points, trackings_input_label, original_image, evt):
+    x, y = evt.index
+    tracking_points.append((x, y))
+    trackings_input_label.append(1 if point_type == "include" else 0)
+    
+    # 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
 
 # use bfloat16 for the entire notebook
 torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
@@ -108,96 +96,62 @@ def show_masks(image, masks, scores, point_coords=None, box_coords=None, input_l
     return combined_images, mask_images
 
 @spaces.GPU()
-def sam_process(input_image, tracking_points, trackings_input_label):
-    image = Image.open(input_image)
-    image = np.array(image.convert("RGB"))
-    # if checkpoint == "tiny":
-    #     sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt"
-    #     model_cfg = "sam2_hiera_t.yaml"
-    # elif checkpoint == "small":
-    #     sam2_checkpoint = "./checkpoints/sam2_hiera_small.pt"
-    #     model_cfg = "sam2_hiera_s.yaml"
-    # elif checkpoint == "base-plus":
-    #     sam2_checkpoint = "./checkpoints/sam2_hiera_base_plus.pt"
-    #     model_cfg = "sam2_hiera_b+.yaml"
-    # elif checkpoint == "large":
-    #     sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt"
-    #     model_cfg = "sam2_hiera_l.yaml"
+def sam_process(original_image, points, labels):
+    # Convert image to numpy array for SAM2 processing
+    image = np.array(original_image)
     predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
-    # print(predictor)
-    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-        predictor.set_image(image)
-        input_point = np.array(tracking_points.value)
-        input_label = np.array(trackings_input_label.value)
-        print(predictor._features["image_embed"].shape, predictor._features["image_embed"][-1].shape)
-        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]
-        logits = logits[sorted_ind]
-        print(masks.shape)
-        results, mask_results = show_masks(image, masks, scores, point_coords=input_point, input_labels=input_label, borders=True)
-        print(results)
-        return results[0], mask_results[0]
+    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)
+    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_path = gr.State()
+    first_frame = gr.State()  # Tracks original image
     tracking_points = gr.State([])
     trackings_input_label = gr.State([])
+    
     with gr.Column():
         gr.Markdown("# SAM2 Image Predictor")
-        gr.Markdown("This is a simple demo for image segmentation with SAM2.")
-        gr.Markdown("""Instructions:
-        1. Upload your image
-        2. With 'include' point type selected, Click on the object to mask
-        3. Switch to 'exclude' point type if you want to specify an area to avoid
-        4. Submit !
-        """)
+        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():
-            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="tiny")
-                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
-        )
+            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(
-            fn=preprocess_image,
-            inputs=[points_map],
-            outputs=[first_frame_path, tracking_points, trackings_input_label, input_image],
-            queue=False
+            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(
-            fn=get_point,
-            inputs=[point_type, tracking_points, trackings_input_label, first_frame_path],
-            outputs=[tracking_points, trackings_input_label, points_map],
-            queue=False
+            get_point,
+            inputs=[point_type, tracking_points, trackings_input_label, first_frame],
+            outputs=[tracking_points, trackings_input_label, points_map]
         )
-        submit_btn.click(
-            fn=sam_process,
+        submit_button.click(
+            sam_process,
             inputs=[input_image, tracking_points, trackings_input_label],
-            outputs=[output_result, output_result_mask]
+            outputs=output_image
         )
-    return input_image, points_map, output_result_mask
+    
+    return input_image, points_map, output_image