Update app.py

app.py

@@ -1,12 +1,10 @@
 import gradio as gr
-
 import os
-os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "1"
 import torch
 import numpy as np
 import cv2
 import matplotlib.pyplot as plt
-from PIL import Image, ImageFilter
+from PIL import Image
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
@@ -25,15 +23,12 @@ def get_point(point_type, tracking_points, trackings_input_label, first_frame_pa
         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
+    fraction = 0.02
     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):
@@ -42,21 +37,15 @@ def get_point(point_type, tracking_points, trackings_input_label, first_frame_pa
         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
-    
-# 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
-    
-def show_mask(mask, ax, random_color=False, borders = True):
+# 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)
     else:
@@ -65,9 +54,7 @@ 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:
-        import cv2
         contours, _ = cv2.findContours(mask,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) 
-        # Try to smooth contours
         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) 
     ax.imshow(mask_image)
@@ -84,94 +71,66 @@ def show_box(box, ax):
     ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))    
 
 def show_masks(image, masks, scores, point_coords=None, box_coords=None, input_labels=None, borders=True):
-    combined_images = []  # List to store filenames of images with masks overlaid
-    mask_images = []      # List to store filenames of separate mask images
+    combined_images = []  
+    mask_images = []      
 
     for i, (mask, score) in enumerate(zip(masks, scores)):
-        # ---- Original Image with Mask Overlaid ----
         plt.figure(figsize=(10, 10))
         plt.imshow(image)
-        show_mask(mask, plt.gca(), borders=borders)  # Draw the mask with borders
-        """
-        if point_coords is not None:
-            assert input_labels is not None
-            show_points(point_coords, input_labels, plt.gca())
-        """
-        if box_coords is not None:
-            show_box(box_coords, plt.gca())
-        if len(scores) > 1:
-            plt.title(f"Mask {i+1}, Score: {score:.3f}", fontsize=18)
+        show_mask(mask, plt.gca(), borders=borders)
         plt.axis('off')
 
-        # Save the figure as a JPG file
         combined_filename = f"combined_image_{i+1}.jpg"
         plt.savefig(combined_filename, format='jpg', bbox_inches='tight')
         combined_images.append(combined_filename)
+        plt.close()
 
-        plt.close()  # Close the figure to free up memory
-
-        # ---- Separate Mask Image (White Mask on Black Background) ----
-        # Create a black image
         mask_image = np.zeros_like(image, dtype=np.uint8)
-        
-        # The mask is a binary array where the masked area is 1, else 0.
-        # Convert the mask to a white color in the mask_image
         mask_layer = (mask > 0).astype(np.uint8) * 255
-        for c in range(3):  # Assuming RGB, repeat mask for all channels
+        for c in range(3):
             mask_image[:, :, c] = mask_layer
 
-        # Save the mask image
         mask_filename = f"mask_image_{i+1}.png"
         Image.fromarray(mask_image).save(mask_filename)
         mask_images.append(mask_filename)
 
-        plt.close()  # Close the figure to free up memory
-
     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"))
 
-    if checkpoint == "tiny":
-        sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt"
-        model_cfg = "sam2_hiera_t.yaml"
-    elif checkpoint == "samll":
-        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"
+    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_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
+    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)
 
-    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)
+    results, mask_results = show_masks(image, masks, scores, 
+                                     point_coords=input_point, 
+                                     input_labels=input_label, 
+                                     borders=True)
 
     return results[0], mask_results[0]
 
@@ -180,23 +139,12 @@ with gr.Blocks() as demo:
     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("# SAM2 Image Predictor (CPU Version)")
+        gr.Markdown("This version runs entirely on CPU")
         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
-                )
+                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")
@@ -207,30 +155,30 @@ with gr.Blocks() as demo:
                 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],
+        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
+        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
+        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],
-        outputs = [output_result, output_result_mask]
+        fn=sam_process,
+        inputs=[input_image, checkpoint, tracking_points, trackings_input_label],
+        outputs=[output_result, output_result_mask]
     )
 
 demo.launch(show_api=False, show_error=True)