initial segmentation app

app.py

@@ -0,0 +1,64 @@
+from io import BytesIO
+import numpy as np
+import gradio as gr
+from PIL import Image
+import requests
+import torch
+from transformers import AutoModelForImageSegmentation, AutoFeatureExtractor
+
+from utils import annotate_masks
+from utils.sam import predict
+
+# Load the model and feature extractor
+model_name = "facebook/detr-resnet-50"
+model = AutoModelForImageSegmentation.from_pretrained(model_name)
+extractor = AutoFeatureExtractor.from_pretrained(model_name)
+
+# Function to handle segmentation
+def segment_image(image):
+    method = "sam"
+    if method == "sam":
+        point=[300,300]        
+        image_rgb = np.array(image)  # Converts PIL image directly to RGB NumPy array
+        if image_rgb.size == 0:
+            raise ValueError("The image is empty!")
+        if len(image_rgb.shape) == 2:  # Grayscale image fix
+            image_rgb = np.stack([image_rgb]*3, axis=-1)
+        elif len(image_rgb.shape) == 3 and image_rgb.shape[2] == 4:  # RGBA to RGB
+            image_rgb = image_rgb[:, :, :3]
+
+        print(f"========================Image type: {type(image_rgb)}, Shape: {image_rgb.shape}")
+
+        # Ensure correct format for SAM (RGB and np.uint8)
+        if image_rgb.dtype != np.uint8:
+            image_rgb = (image_rgb * 255).astype(np.uint8)        
+
+        masks, scores, logits = predict(image_rgb, [point])
+        return annotate_masks(image_rgb, masks)
+    else:        
+        # Prepare the image and perform segmentation
+        inputs = extractor(images=image, return_tensors="pt")
+        with torch.no_grad():
+            outputs = model(**inputs)
+        segmentation_mask = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
+    
+        # Convert the segmentation mask to an image
+        mask_image = Image.fromarray(segmentation_mask.astype('uint8'))
+        return mask_image
+
+# Create Gradio interface
+demo = gr.Interface(
+    fn=segment_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    live=True,
+    title="Image Segmentation App",
+    description="Upload an image and get the segmented output using a pre-trained model."
+)
+
+# Launch the Gradio app
+
+if __name__ == "__main__":
+    demo.launch()
+
+

requirements.txt

@@ -0,0 +1,7 @@
+git+https://github.com/facebookresearch/segment-anything.git
+gradio
+supervision
+torch --index-url https://download.pytorch.org/whl/cu124
+torchvision --index-url https://download.pytorch.org/whl/cu124
+torchaudio --index-url https://download.pytorch.org/whl/cu124
+transformers

utils/__init__.py

@@ -0,0 +1,13 @@
+import cv2
+import numpy as np
+
+def annotate_masks(image, masks):
+    canvas = np.zeros_like(image)
+    for i, mask in enumerate(masks):
+        # Generate a unique color for each mask (you can also choose specific colors)
+        color = np.random.randint(0, 256, size=3)  # Random color (R, G, B)
+
+        # Apply the color to the masked regions
+        canvas[mask == 1] = color
+    overlay_image = cv2.addWeighted(image, 0.7, canvas, 0.3, 0)
+    return overlay_image

utils/sam.py

@@ -0,0 +1,19 @@
+import numpy as np
+import torch
+from segment_anything import sam_model_registry, SamPredictor
+
+# Load the SAM model (adjust the path to the model checkpoint)
+sam_checkpoint = r"H:\dev\pantareh\data\models\sam_vit_h_4b8939.pth"
+sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
+
+
+def predict(image, points):
+    # Initialize the predictor
+    predictor = SamPredictor(sam)
+
+    predictor.set_image(image)
+
+    # Generate masks (you can adjust parameters based on your requirements)
+    masks, scores, logits = predictor.predict(point_coords=np.array(points), point_labels=np.ones(len(points)), box=None)
+    return masks, scores, logits 
+