Spaces:
Sleeping
Sleeping
| from transformers import pipeline, SamModel, SamProcessor | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| import requests | |
| # Image Segmentation Model | |
| sam_model = SamModel.from_pretrained("Zigeng/SlimSAM-uniform-77") | |
| sam_processor = SamProcessor.from_pretrained("Zigeng/SlimSAM-uniform-77") | |
| def show_colored_mask(mask, combined_mask, color): | |
| """ | |
| Add a single-colored mask to the combined mask. | |
| Args: | |
| mask (numpy.ndarray): Binary mask to overlay. | |
| combined_mask (numpy.ndarray): Combined RGBA mask. | |
| color (tuple): RGBA color for the mask. | |
| """ | |
| if mask.ndim == 3: # If mask has channels then take the first one | |
| mask = mask[0] | |
| mask = mask.squeeze() # Remove extra dimension | |
| mask_binary = (mask > 0).astype(np.uint8) # Ensure the mask is binary | |
| # Apply the color to the mask | |
| for c in range(3): # RGB channels | |
| combined_mask[:, :, c] = np.where(mask_binary > 0, color[c], combined_mask[:, :, c]) | |
| combined_mask[:, :, 3] = np.where(mask_binary > 0, color[3], combined_mask[:, :, 3]) # Alpha channel (transperency) | |
| def segment_image(input_image, input_points): | |
| """ | |
| Perform image segmentation and overlay masks with a single solid color. | |
| Args: | |
| input_image (PIL.Image): The input image. | |
| input_points (list): List of points [[x, y], ...]. | |
| Returns: | |
| PIL.Image: Image with masks applied in one solid red color. | |
| """ | |
| # Convert input points to a 4D tensor | |
| input_points_tensor = torch.tensor(input_points, dtype=torch.float32).unsqueeze(0).unsqueeze(1) | |
| # Process input and run the SAM model | |
| inputs = sam_processor(input_image, input_points=input_points_tensor, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = sam_model(**inputs) | |
| # Post-process masks | |
| predicted_masks = sam_processor.image_processor.post_process_masks( | |
| outputs.pred_masks, inputs["original_sizes"], inputs["reshaped_input_sizes"] | |
| ) | |
| # Define a solid red color with full opacity | |
| single_color = (255, 0, 0, 100) | |
| # Prepare a combined RGBA mask | |
| image_size = input_image.size | |
| combined_mask = np.zeros((image_size[1], image_size[0], 4), dtype=np.uint8) | |
| # Apply all masks using the single color | |
| for mask in predicted_masks[0]: | |
| mask = mask.numpy() | |
| show_colored_mask(mask, combined_mask, single_color) | |
| # Combine the mask with the original image | |
| input_image_rgba = input_image.convert("RGBA") # Red Green Blue Alpha | |
| combined_image = Image.alpha_composite(input_image_rgba, Image.fromarray(combined_mask, "RGBA")) | |
| return combined_image | |