Update app.py

app.py

@@ -12,14 +12,13 @@ import gradio as gr
 
 # image= cv2.imread('image_4.png', cv2.IMREAD_COLOR)
 def get_masks(model_type, image):
-    if model_type == 'vit_h':
+    if model_type.all() == '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':
+        
+    if model_type,all() == 'vit_b':
         sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth")
         
-    if model_type == 'vit_l':
+    if model_type.all() == 'vit_l':
         sam = sam_model_registry["vit_l"](checkpoint="sam_vit_l_0b3195.pth")
     
     mask_generator = SamAutomaticMaskGenerator(sam)