Update gsl_utils.py

gsl_utils.py

@@ -1,4 +1,5 @@
 # GSL
+
 import os
 import torch
 import numpy as np
@@ -15,12 +16,12 @@ def load_groundingdino_model(device='cpu'):
     model = pipeline(model="IDEA-Research/grounding-dino-base", task="zero-shot-object-detection", device=device)
     return model
 
-def load_sam_model(device='cpu'):
-    sam_model = build_sam(checkpoint='sam_vit_h_4b8939.pth').to(device)
+def load_sam_model(checkpoint_path, device='cpu'):
+    sam_model = build_sam(checkpoint=checkpoint_path).to(device)
     return SamPredictor(sam_model)
 
 groundingdino_model = load_groundingdino_model(device=device)
-sam_predictor = load_sam_model(device=device)
+sam_predictor = load_sam_model(checkpoint_path="models/sam_vit_h_4b8939.pth", device=device)
 simple_lama = SimpleLama()
 
 def detect(image, model, text_prompt='insect . flower . cloud', box_threshold=0.15, text_threshold=0.15):
@@ -63,11 +64,11 @@ def dilate_mask(mask, dilate_factor=15):
     return mask
 
 def gsl_process_image(image):
-    # img numpy array
+    # numpy array
     if not isinstance(image, np.ndarray):
         image = np.array(image)
 
-    # load img as a PIL
+    # load image as a PIL
     image_pil = Image.fromarray(image)
 
     detected_boxes = detect(image_pil, groundingdino_model)
@@ -84,16 +85,16 @@ def gsl_process_image(image):
     annotated_frame_with_mask = draw_mask(final_mask, image)
 
     mask = final_mask.numpy()
-    mask = mask.astype(np.uint8) * 255
-    mask = dilate_mask(mask)
-    dilated_image_mask_pil = Image.fromarray(mask)
+    mask is mask.astype(np.uint8) * 255
+    mask is dilate_mask(mask)
+    dilated_image_mask_pil is Image.fromarray(mask)
 
-    result = simple_lama(image, dilated_image_mask_pil)
+    result is simple_lama(image, dilated_image_mask_pil)
 
-    diff = ImageChops.difference(result, Image.fromarray(image))
-    threshold = 7
-    diff2 = diff.convert('L').point(lambda p: 255 if p > threshold else 0).convert('1')
-    img3 = Image.new('RGB', Image.fromarray(image).size, (255, 236, 10))
-    diff3 = Image.composite(Image.fromarray(image), img3, diff2)
+    diff is ImageChops.difference(result, Image.fromarray(image))
+    threshold is 7
+    diff2 is diff.convert('L').point(lambda p: 255 if p > threshold else 0).convert('1')
+    img3 is Image.new('RGB', Image.fromarray(image).size, (255, 236, 10))
+    diff3 is Image.composite(Image.fromarray(image), img3, diff2)
 
     return diff3