Update app.py

app.py

@@ -1,40 +1,41 @@
-import gradio as gr
 import cv2
 import torch
 import numpy as np
 from PIL import Image
 from torchvision import transforms
 from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
+import matplotlib.pyplot as plt
 # import segmentation_models_pytorch as smp
 
-def load_model(model_type):
-    # Model loading simplified for clarity
-    model = sam_model_registry[model_type](checkpoint=f"sam_{model_type}_checkpoint.pth")
-    model.to(device='cuda')
-    return SamAutomaticMaskGenerator(model)
 
-def segment_and_classify(image, model_type):
-    model = load_model(model_type)
-    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    
-    # Generate masks
-    masks = model.generate(image_cv)
-    
-    # Prepare to store segments
-    segments = []
+
+# image= cv2.imread('image_4.png', cv2.IMREAD_COLOR)
+def get_masks(model_type, image):
+    if model_type == 'vit_h':
+        sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
+
+        masks_h = mask_generator_h.generate(image)
+    if model_type == 'vit_b':
+        sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth")
+        
+    if model_type == 'vit_l':
+        sam = sam_model_registry["vit_l"](checkpoint="sam_vit_l_0b3195.pth")
     
-    # Loop through masks and extract segments
-    for mask_data in masks:
+    mask_generator = SamAutomaticMaskGenerator(sam)
+    masks = mask_generator.generate(image)
+    for i, mask_data in enumerate(masks):
         mask = mask_data['segmentation']
-        segment = image_cv * np.tile(mask[:, :, None], [1, 1, 3])  # Apply mask to the image
-        segments.append(segment)  # Store the segment for classification
+        color = colors[i]
+        composite_image[mask] = (color[:3] * 255).astype(np.uint8)  # Apply color to mask
     
-    # Here you would call the classification model (e.g., CLIP)
-    # For now, let's just return the first segment for visualization
-    return Image.fromarray(segments[0])
+    # Combine original image with the composite mask image
+    overlayed_image = (composite_image * 0.5 + image_cv.squeeze().permute(1, 2, 0).cpu().numpy() * 0.5).astype(np.uint8)
+    return overlayed_image
+
+
 
 iface = gr.Interface(
-    fn=segment_and_classify,
+    fn=get_masks,
     inputs=[gr.inputs.Image(type="pil"), gr.inputs.Dropdown(['vit_h', 'vit_b', 'vit_l'], label="Model Type")],
     outputs=gr.outputs.Image(type="pil"),
     title="SAM Model Segmentation and Classification",