EoL_Understanding / utils.py
valegro's picture
Update utils.py
010b878 verified
raw
history blame contribute delete
774 Bytes
import importlib, torch, cv2, numpy as np
from huggingface_hub import hf_hub_download
from segment_anything import sam_model_registry, SamPredictor
from groundingdino.util.inference import Model as GDModel
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@torch.inference_mode()
def load_sam():
ckpt = hf_hub_download("ybelkada/segment-anything", "sam_vit_b.pth")
sam = sam_model_registry["vit_b"](checkpoint=ckpt)
return SamPredictor(sam.to(DEVICE))
@torch.inference_mode()
def load_groundingdino():
ckpt = hf_hub_download(
"GroundingDINO/groundingdino-swint-ogc",
"groundingdino_swint_ogc.pth"
)
return GDModel(model_config_path=None, model_checkpoint_path=ckpt, device=DEVICE)
SAM = load_sam()
GD = load_groundingdino()