valegro commited on
Commit
cf5fa6f
·
verified ·
1 Parent(s): f6c3c05

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +23 -0
utils.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch, cv2
2
+ from huggingface_hub import hf_hub_download
3
+ from segment_anything import sam_model_registry, SamPredictor
4
+ from groundingdino.util.inference import Model as GDModel
5
+
6
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
7
+
8
+ @torch.inference_mode()
9
+ def load_sam():
10
+ ckpt = hf_hub_download("ybelkada/segment-anything", "sam_vit_b.pth")
11
+ sam = sam_model_registry["vit_b"](checkpoint=ckpt)
12
+ return SamPredictor(sam.to(DEVICE))
13
+
14
+ @torch.inference_mode()
15
+ def load_groundingdino():
16
+ ckpt = hf_hub_download(
17
+ "GroundingDINO/groundingdino-swint-ogc",
18
+ "groundingdino_swint_ogc.pth"
19
+ )
20
+ return GDModel(model_config_path=None, model_checkpoint_path=ckpt, device=DEVICE)
21
+
22
+ SAM = load_sam()
23
+ GD = load_groundingdino()