| 1024 | return updated_image_path |
| 1025 | |
| 1026 | class Text2Box: |
| 1027 | def __init__(self, device): |
| 1028 | print(f"Initializing ObjectDetection to {device}") |
| 1029 | self.device = device |
| 1030 | self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 |
| 1031 | self.model_checkpoint_path = os.path.join("checkpoints","groundingdino") |
| 1032 | self.model_config_path = os.path.join("checkpoints","grounding_config.py") |
| 1033 | self.download_parameters() |
| 1034 | self.box_threshold = 0.3 |
| 1035 | self.text_threshold = 0.25 |
| 1036 | self.grounding = (self.load_model()).to(self.device) |
| 1037 | |
| 1038 | def download_parameters(self): |
| 1039 | url = "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth" |
| 1040 | if not os.path.exists(self.model_checkpoint_path): |
| 1041 | wget.download(url,out=self.model_checkpoint_path) |
| 1042 | config_url = "https://raw.githubusercontent.com/IDEA-Research/GroundingDINO/main/groundingdino/config/GroundingDINO_SwinT_OGC.py" |
| 1043 | if not os.path.exists(self.model_config_path): |
| 1044 | wget.download(config_url,out=self.model_config_path) |
| 1045 | def load_image(self,image_path): |
| 1046 | # load image |
| 1047 | image_pil = Image.open(image_path).convert("RGB") # load image |
| 1048 | |
| 1049 | transform = T.Compose( |
| 1050 | [ |
| 1051 | T.RandomResize([512], max_size=1333), |
| 1052 | T.ToTensor(), |
| 1053 | T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| 1054 | ] |
| 1055 | ) |
| 1056 | image, _ = transform(image_pil, None) # 3, h, w |
| 1057 | return image_pil, image |
| 1058 | |
| 1059 | def load_model(self): |
| 1060 | args = SLConfig.fromfile(self.model_config_path) |
| 1061 | args.device = self.device |
| 1062 | model = build_model(args) |
| 1063 | checkpoint = torch.load(self.model_checkpoint_path, map_location="cpu") |
| 1064 | load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) |
| 1065 | print(load_res) |
| 1066 | _ = model.eval() |
| 1067 | return model |
| 1068 | |
| 1069 | def get_grounding_boxes(self, image, caption, with_logits=True): |
| 1070 | caption = caption.lower() |
| 1071 | caption = caption.strip() |
| 1072 | if not caption.endswith("."): |
| 1073 | caption = caption + "." |
| 1074 | image = image.to(self.device) |
| 1075 | with torch.no_grad(): |
| 1076 | outputs = self.grounding(image[None], captions=[caption]) |
| 1077 | logits = outputs["pred_logits"].cpu().sigmoid()[0] # (nq, 256) |
| 1078 | boxes = outputs["pred_boxes"].cpu()[0] # (nq, 4) |
| 1079 | logits.shape[0] |
| 1080 | |
| 1081 | # filter output |
| 1082 | logits_filt = logits.clone() |
| 1083 | boxes_filt = boxes.clone() |
nothing calls this directly
no outgoing calls
no test coverage detected