| # import gradio as gr | |
| # from gradio_image_prompter import ImagePrompter | |
| # import os | |
| # import torch | |
| # def prompter(prompts): | |
| # image = prompts["image"] # Get the image from prompts | |
| # points = prompts["points"] # Get the points from prompts | |
| # # Print the collected inputs for debugging or logging | |
| # print("Image received:", image) | |
| # print("Points received:", points) | |
| # import torch | |
| # from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| # device = torch.device("cpu") | |
| # predictor = SAM2ImagePredictor.from_pretrained( | |
| # "facebook/sam2-hiera-base-plus", device=device | |
| # ) | |
| # with torch.inference_mode(): | |
| # predictor.set_image(image) | |
| # # masks, _, _ = predictor.predict([[point[0], point[1]] for point in points]) | |
| # input_point = [[point[0], point[1]] for point in points] | |
| # input_label = [1] | |
| # masks, _, _ = predictor.predict( | |
| # point_coords=input_point, point_labels=input_label | |
| # ) | |
| # print("Predicted Mask:", masks) | |
| # return image, points | |
| # # Define the Gradio interface | |
| # demo = gr.Interface( | |
| # fn=prompter, # Use the custom prompter function | |
| # inputs=ImagePrompter( | |
| # show_label=False | |
| # ), # ImagePrompter for image input and point selection | |
| # outputs=[ | |
| # gr.Image(show_label=False), # Display the image | |
| # gr.Dataframe(label="Points"), # Display the points in a DataFrame | |
| # ], | |
| # title="Image Point Collector", | |
| # description="Upload an image, click on it, and get the coordinates of the clicked points.", | |
| # ) | |
| # # Launch the Gradio app | |
| # demo.launch() | |
| # import gradio as gr | |
| # from gradio_image_prompter import ImagePrompter | |
| # import torch | |
| # from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| # def prompter(prompts): | |
| # image = prompts["image"] # Get the image from prompts | |
| # points = prompts["points"] # Get the points from prompts | |
| # # Print the collected inputs for debugging or logging | |
| # print("Image received:", image) | |
| # print("Points received:", points) | |
| # device = torch.device("cpu") | |
| # # Load the SAM2ImagePredictor model | |
| # predictor = SAM2ImagePredictor.from_pretrained( | |
| # "facebook/sam2-hiera-base-plus", device=device | |
| # ) | |
| # # Perform inference | |
| # with torch.inference_mode(): | |
| # predictor.set_image(image) | |
| # input_point = [[point[0], point[1]] for point in points] | |
| # input_label = [1] * len(points) # Assuming all points are foreground | |
| # masks, _, _ = predictor.predict( | |
| # point_coords=input_point, point_labels=input_label | |
| # ) | |
| # # The masks are returned as a list of numpy arrays | |
| # print("Predicted Mask:", masks) | |
| # # Assuming there's only one mask returned, you can adjust if there are multiple | |
| # predicted_mask = masks[0] | |
| # print(len(image)) | |
| # print(len(predicted_mask)) | |
| # # Create annotations for AnnotatedImage | |
| # annotations = [(predicted_mask, "Predicted Mask")] | |
| # return image, annotations | |
| # # Define the Gradio interface | |
| # demo = gr.Interface( | |
| # fn=prompter, # Use the custom prompter function | |
| # inputs=ImagePrompter( | |
| # show_label=False | |
| # ), # ImagePrompter for image input and point selection | |
| # outputs=gr.AnnotatedImage(), # Display the image with the predicted mask | |
| # title="Image Point Collector with Mask Overlay", | |
| # description="Upload an image, click on it, and get the predicted mask overlayed on the image.", | |
| # ) | |
| # # Launch the Gradio app | |
| # demo.launch() | |
| import gradio as gr | |
| from gradio_image_prompter import ImagePrompter | |
| import torch | |
| import numpy as np | |
| from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| from PIL import Image | |
| def prompter(prompts): | |
| image = np.array(prompts["image"]) # Convert the image to a numpy array | |
| points = prompts["points"] # Get the points from prompts | |
| # Print the collected inputs for debugging or logging | |
| print("Image received:", image) | |
| print("Points received:", points) | |
| device = torch.device("cpu") | |
| # Load the SAM2ImagePredictor model | |
| predictor = SAM2ImagePredictor.from_pretrained( | |
| "facebook/sam2-hiera-base-plus", device=device | |
| ) | |
| # Perform inference with multimask_output=True | |
| with torch.inference_mode(): | |
| predictor.set_image(image) | |
| input_point = [[point[0], point[1]] for point in points] | |
| input_label = [1] * len(points) # Assuming all points are foreground | |
| masks, _, _ = predictor.predict( | |
| point_coords=input_point, point_labels=input_label, multimask_output=True | |
| ) | |
| # Prepare individual images with separate overlays | |
| overlay_images = [] | |
| for i, mask in enumerate(masks): | |
| print(f"Predicted Mask {i+1}:", mask) | |
| red_mask = np.zeros_like(image) | |
| red_mask[:, :, 0] = mask.astype(np.uint8) * 255 # Apply the red channel | |
| red_mask = Image.fromarray(red_mask) | |
| # Convert the original image to a PIL image | |
| original_image = Image.fromarray(image) | |
| # Blend the original image with the red mask | |
| blended_image = Image.blend(original_image, red_mask, alpha=0.5) | |
| # Add the blended image to the list | |
| overlay_images.append(blended_image) | |
| return overlay_images | |
| # Define the Gradio interface | |
| demo = gr.Interface( | |
| fn=prompter, # Use the custom prompter function | |
| inputs=ImagePrompter( | |
| show_label=False | |
| ), # ImagePrompter for image input and point selection | |
| outputs=[ | |
| gr.Image(show_label=False) for _ in range(3) | |
| ], # Display up to 3 overlay images | |
| title="Image Point Collector with Multiple Separate Mask Overlays", | |
| description="Upload an image, click on it, and get each predicted mask overlaid separately in red on individual images.", | |
| ) | |
| # Launch the Gradio app | |
| demo.launch() | |