aikenml/data_mining

1import torch2import numpy as np3import cv24import PIL5 6from groundingdino.models import build_model as build_grounding_dino7from groundingdino.util.slconfig import SLConfig8from groundingdino.util.utils import clean_state_dict9from groundingdino.util.inference import annotate, load_image, predict10import groundingdino.datasets.transforms as T11 12from torchvision.ops import box_convert13 14class Detector:15    def __init__(self, device):16        config_file = "src/groundingdino/groundingdino/config/GroundingDINO_SwinT_OGC.py"17        grounding_dino_ckpt = './ckpt/groundingdino_swint_ogc.pth'18        args = SLConfig.fromfile(config_file) 19        args.device = device20        self.deivce = device21        self.gd = build_grounding_dino(args)22 23        checkpoint = torch.load(grounding_dino_ckpt, map_location='cpu')24        log = self.gd.load_state_dict(clean_state_dict(checkpoint['model']), strict=False)25        print("Model loaded from {} \n => {}".format(grounding_dino_ckpt, log))26        self.gd.eval()27    28    def image_transform_grounding(self, init_image):29        transform = T.Compose([30            T.RandomResize([800], max_size=1333),31            T.ToTensor(),32            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])33        ])34        image, _ = transform(init_image, None) # 3, h, w35        return init_image, image36 37    def image_transform_grounding_for_vis(self, init_image):38        transform = T.Compose([39            T.RandomResize([800], max_size=1333),40        ])41        image, _ = transform(init_image, None) # 3, h, w42        return image43 44    def transfer_boxes_format(self, boxes, height, width):45        boxes = boxes * torch.Tensor([width, height, width, height])46        boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy")47 48        transfered_boxes = []49        for i in range(len(boxes)):50            box = boxes[i]51            transfered_box = [[int(box[0]), int(box[1])], [int(box[2]), int(box[3])]]52            transfered_boxes.append(transfered_box)53        54        transfered_boxes = np.array(transfered_boxes)55        return transfered_boxes56        57    @torch.no_grad()58    def run_grounding(self, origin_frame, grounding_caption, box_threshold, text_threshold):59        '''60            return:61                annotated_frame:nd.array62                transfered_boxes: nd.array [N, 4]: [[x0, y0], [x1, y1]]63        '''64        height, width, _ = origin_frame.shape65        img_pil = PIL.Image.fromarray(origin_frame)66        re_width, re_height = img_pil.size67        _, image_tensor = self.image_transform_grounding(img_pil)68        # img_pil = self.image_transform_grounding_for_vis(img_pil)69 70        # run grounidng71        boxes, logits, phrases = predict(self.gd, image_tensor, grounding_caption, box_threshold, text_threshold, device=self.deivce)72        annotated_frame = annotate(image_source=np.asarray(img_pil), boxes=boxes, logits=logits, phrases=phrases)[:, :, ::-1]73        annotated_frame = cv2.resize(annotated_frame, (width, height), interpolation=cv2.INTER_LINEAR)74        75        # transfer boxes to sam-format 76        transfered_boxes = self.transfer_boxes_format(boxes, re_height, re_width)77        return annotated_frame, transfered_boxes78 79if __name__ == "__main__":80    detector = Detector("cuda")81    origin_frame = cv2.imread('./debug/point.png')82    origin_frame = cv2.cvtColor(origin_frame, cv2.COLOR_BGR2RGB)83    grounding_caption = "swan.water"84    box_threshold = 0.2585    text_threshold = 0.2586 87    annotated_frame, boxes = detector.run_grounding(origin_frame, grounding_caption, box_threshold, text_threshold)88    cv2.imwrite('./debug/x.png', annotated_frame)89 90    for i in range(len(boxes)):91        bbox = boxes[i]92        origin_frame = cv2.rectangle(origin_frame, bbox[0], bbox[1], (0, 0, 255))93    cv2.imwrite('./debug/bbox_frame.png', origin_frame)94