sam_vit_h_4b8939.pth

gsl_utils.py

@@ -5,17 +5,22 @@ import numpy as np
 from PIL import Image, ImageChops, ImageEnhance
 import cv2
 from simple_lama_inpainting import SimpleLama
+from segment_anything import build_sam, SamPredictor
 from transformers import pipeline
 from huggingface_hub import hf_hub_download
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def load_groundingdino_model(device='cpu'):
-    model = pipeline(model="IDEA-Research/grounding-dino-base", task="zero-shot-object-detection", device=device)
+    model = pipeline(model="ShilongLiu/GroundingDINO", 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)
+    return SamPredictor(sam_model)
+
 groundingdino_model = load_groundingdino_model(device=device)
-sam_predictor = None  # Initialize this properly using build_sam
+sam_predictor = load_sam_model(device=device)
 simple_lama = SimpleLama()
 
 def detect(image, model, text_prompt='insect . flower . cloud', box_threshold=0.15, text_threshold=0.15):
@@ -24,7 +29,6 @@ def detect(image, model, text_prompt='insect . flower . cloud', box_threshold=0.
     return results
 
 def segment(image, sam_model, boxes):
-    # sam_moded initialized with build_sam
     sam_model.set_image(image)
     H, W, _ = image.shape
     boxes_xyxy = torch.Tensor(boxes) * torch.Tensor([W, H, W, H])
@@ -59,14 +63,13 @@ def dilate_mask(mask, dilate_factor=15):
     return mask
 
 def gsl_process_image(image):
-    # image numpy array
+    # img numpy array
     if not isinstance(image, np.ndarray):
         image = np.array(image)
 
-    # load as a PIL
+    # load img as a PIL
     image_pil = Image.fromarray(image)
 
-    # detect insects
     detected_boxes = detect(image_pil, groundingdino_model)
     boxes = [[d['box']['xmin'], d['box']['ymin'], d['box']['xmax'], d['box']['ymax']] for d in detected_boxes]
     segmented_frame_masks = segment(image, sam_predictor, boxes)