Update app.py

app.py

@@ -114,10 +114,11 @@ def build_groundingdino():
     return groundingdino
 
 
-
+model_sam = build_sam()
+model_groundingdino = build_groundingdino()
 
 '''Predictions'''
-def predict_dino(image_pil, text_prompt, box_threshold, text_threshold, model_groundingdino):
+def predict_dino(image_pil, text_prompt, box_threshold, text_threshold):
     image_trans = transform_image(image_pil)
     boxes, logits, phrases = predict(model = model_groundingdino,
                                      image = image_trans,
@@ -130,7 +131,7 @@ def predict_dino(image_pil, text_prompt, box_threshold, text_threshold, model_gr
     return boxes, logits, phrases
 
 
-def predict_sam(image_pil, boxes, model_sam):
+def predict_sam(image_pil, boxes):
     image_array = np.asarray(image_pil)
     model_sam.set_image(image_array)
     transformed_boxes = model_sam.transform.apply_boxes_torch(boxes, image_array.shape[:2])
@@ -143,11 +144,11 @@ def predict_sam(image_pil, boxes, model_sam):
     return masks.cpu()
 
 
-def mask_predict(image_pil, text_prompt, box_threshold=0.3, text_threshold=0.25, models = [model_groundingdino , model_sam]):
-    boxes, logits, phrases = predict_dino(image_pil, text_prompt, box_threshold, text_threshold, models[0])
+def mask_predict(image_pil, text_prompt, box_threshold=0.3, text_threshold=0.25):
+    boxes, logits, phrases = predict_dino(image_pil, text_prompt, box_threshold, text_threshold)
     masks = torch.tensor([])
     if len(boxes) > 0:
-        masks = predict_sam(image_pil, boxes, models[1])
+        masks = predict_sam(image_pil, boxes)
         masks = masks.squeeze(1)
     return masks, boxes, phrases, logits
 
@@ -189,9 +190,7 @@ def visualize_results(img1, img2, task):
 # text_prompt = 'wooden stool'
 # image_path = '/kaggle/input/avataar/stool.jpeg'
 # output_image_path = '/kaggle/working'
-        
-model_sam = build_sam()
-model_groundingdino = build_groundingdino()
+    
 
 def main_fun(image_pil, x_units, y_units, text_prompt):
 #     x_units = 200
@@ -199,7 +198,7 @@ def main_fun(image_pil, x_units, y_units, text_prompt):
 #     text_prompt = 'wooden stool'
     
 #     image_pil = load_image(image_path)
-    masks, boxes, phrases, logits = mask_predict(image_pil, text_prompt=text_prompt, box_threshold=0.23, text_threshold=0.25,models = [model_groundingdino , model_sam])
+    masks, boxes, phrases, logits = mask_predict(image_pil, text_prompt=text_prompt, box_threshold=0.23, text_threshold=0.25)
     output = draw_image(image_pil, masks, boxes, alpha=0.4)