import spaces 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 sam2.build_sam import build_sam2 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, 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__() 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): if random_color: color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0) else: color = np.array([30/255, 144/255, 255/255, 0.6]) h, w = mask.shape[-2:] 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) def show_points(coords, labels, ax, marker_size=375): pos_points = coords[labels == 1] neg_points = coords[labels == 0] ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25) def show_box(box, ax): x0, y0 = box[0], box[1] w, h = box[2] - box[0], box[3] - box[1] 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 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 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) 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() # 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 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 @spaces.GPU() 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") 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 = gr.State() # Tracks original image tracking_points = gr.State([]) trackings_input_label = gr.State([]) 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") # Event handlers points_map.upload( 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( 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], outputs=output_image ) return input_image, points_map, output_image