Update app.py

app.py

@@ -6,34 +6,51 @@ from torchvision import transforms
 from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 import matplotlib.pyplot as plt
 import gradio as gr
+
 # import segmentation_models_pytorch as smp
 
+##set the device to cuda for sam model
+# device = torch.device('cuda')
 
 
 # image= cv2.imread('image_4.png', cv2.IMREAD_COLOR)
-def get_masks(model_type, image):
-    if model_type.all() == 'vit_h':
+def get_masks( image, model_type):
+    print(image)
+    # image_pil = Image.fromarray(image.astype('uint8'), 'RGB')
+    # print(image_pil)
+    if model_type == 'vit_h':
         sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
-        
-    if model_type,all() == 'vit_b':
+    if model_type == 'vit_b':
         sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth")
         
-    if model_type.all() == 'vit_l':
+    if model_type == 'vit_l':
         sam = sam_model_registry["vit_l"](checkpoint="sam_vit_l_0b3195.pth")
+    else:
+        sam=  sam_model_registry["vit_l"](checkpoint="sam_vit_l_0b3195.pth")
     
+    # print(image.shape)
+    #set the device to cuda for sam model
+    # sam.to(device= device)
+
     mask_generator = SamAutomaticMaskGenerator(sam)
     masks = mask_generator.generate(image)
+    composite_image = np.zeros_like(image)
+    colors = plt.cm.jet(np.linspace(0, 1, len(masks)))  # Generate distinct colors
+
     for i, mask_data in enumerate(masks):
         mask = mask_data['segmentation']
         color = colors[i]
         composite_image[mask] = (color[:3] * 255).astype(np.uint8)  # Apply color to mask
+    print(composite_image.shape, image.shape)
     
     # 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)
+    overlayed_image = (composite_image * 0.5 + torch.from_numpy(image).resize(738, 1200, 3).cpu().numpy() * 0.5).astype(np.uint8)
+
     return overlayed_image
 
 
 
+
 iface = gr.Interface(
     fn=get_masks,
     inputs=["image", gr.components.Dropdown(choices=['vit_h', 'vit_b', 'vit_l'], label="Model Type")],
@@ -42,5 +59,4 @@ iface = gr.Interface(
     description="Upload an image, select a model type, and receive the segmented and classified parts."
 )
 
-
 iface.launch()