Delete tool

tool/detector.py

@@ -1,93 +0,0 @@
-import torch
-import numpy as np
-import cv2
-import PIL
-
-from groundingdino.models import build_model as build_grounding_dino
-from groundingdino.util.slconfig import SLConfig
-from groundingdino.util.utils import clean_state_dict
-from groundingdino.util.inference import annotate, load_image, predict
-import groundingdino.datasets.transforms as T
-
-from torchvision.ops import box_convert
-
-class Detector:
-    def __init__(self, device):
-        config_file = "src/groundingdino/groundingdino/config/GroundingDINO_SwinT_OGC.py"
-        grounding_dino_ckpt = './ckpt/groundingdino_swint_ogc.pth'
-        args = SLConfig.fromfile(config_file) 
-        args.device = device
-        self.deivce = device
-        self.gd = build_grounding_dino(args)
-
-        checkpoint = torch.load(grounding_dino_ckpt, map_location='cpu')
-        log = self.gd.load_state_dict(clean_state_dict(checkpoint['model']), strict=False)
-        print("Model loaded from {} \n => {}".format(grounding_dino_ckpt, log))
-        self.gd.eval()
-    
-    def image_transform_grounding(self, init_image):
-        transform = T.Compose([
-            T.RandomResize([800], max_size=1333),
-            T.ToTensor(),
-            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-        ])
-        image, _ = transform(init_image, None) # 3, h, w
-        return init_image, image
-
-    def image_transform_grounding_for_vis(self, init_image):
-        transform = T.Compose([
-            T.RandomResize([800], max_size=1333),
-        ])
-        image, _ = transform(init_image, None) # 3, h, w
-        return image
-
-    def transfer_boxes_format(self, boxes, height, width):
-        boxes = boxes * torch.Tensor([width, height, width, height])
-        boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy")
-
-        transfered_boxes = []
-        for i in range(len(boxes)):
-            box = boxes[i]
-            transfered_box = [[int(box[0]), int(box[1])], [int(box[2]), int(box[3])]]
-            transfered_boxes.append(transfered_box)
-        
-        transfered_boxes = np.array(transfered_boxes)
-        return transfered_boxes
-        
-    @torch.no_grad()
-    def run_grounding(self, origin_frame, grounding_caption, box_threshold, text_threshold):
-        '''
-            return:
-                annotated_frame:nd.array
-                transfered_boxes: nd.array [N, 4]: [[x0, y0], [x1, y1]]
-        '''
-        height, width, _ = origin_frame.shape
-        img_pil = PIL.Image.fromarray(origin_frame)
-        re_width, re_height = img_pil.size
-        _, image_tensor = self.image_transform_grounding(img_pil)
-        # img_pil = self.image_transform_grounding_for_vis(img_pil)
-
-        # run grounidng
-        boxes, logits, phrases = predict(self.gd, image_tensor, grounding_caption, box_threshold, text_threshold, device=self.deivce)
-        annotated_frame = annotate(image_source=np.asarray(img_pil), boxes=boxes, logits=logits, phrases=phrases)[:, :, ::-1]
-        annotated_frame = cv2.resize(annotated_frame, (width, height), interpolation=cv2.INTER_LINEAR)
-        
-        # transfer boxes to sam-format 
-        transfered_boxes = self.transfer_boxes_format(boxes, re_height, re_width)
-        return annotated_frame, transfered_boxes
-
-if __name__ == "__main__":
-    detector = Detector("cuda")
-    origin_frame = cv2.imread('./debug/point.png')
-    origin_frame = cv2.cvtColor(origin_frame, cv2.COLOR_BGR2RGB)
-    grounding_caption = "swan.water"
-    box_threshold = 0.25
-    text_threshold = 0.25
-
-    annotated_frame, boxes = detector.run_grounding(origin_frame, grounding_caption, box_threshold, text_threshold)
-    cv2.imwrite('./debug/x.png', annotated_frame)
-
-    for i in range(len(boxes)):
-        bbox = boxes[i]
-        origin_frame = cv2.rectangle(origin_frame, bbox[0], bbox[1], (0, 0, 255))
-    cv2.imwrite('./debug/bbox_frame.png', origin_frame)

tool/segmentor.py

@@ -1,96 +0,0 @@
-import torch
-import cv2
-import numpy as np
-from sam.segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
-
-class Segmentor:
-    def __init__(self, sam_args):
-        """
-        sam_args:
-            sam_checkpoint: path of SAM checkpoint
-            generator_args: args for everything_generator
-            gpu_id: device
-        """
-        self.device = sam_args["gpu_id"]
-        self.sam = sam_model_registry[sam_args["model_type"]](checkpoint=sam_args["sam_checkpoint"])
-        self.sam.to(device=self.device)
-        self.everything_generator = SamAutomaticMaskGenerator(model=self.sam, **sam_args['generator_args'])
-        self.interactive_predictor = self.everything_generator.predictor
-        self.have_embedded = False
-        
-    @torch.no_grad()
-    def set_image(self, image):
-        # calculate the embedding only once per frame.
-        if not self.have_embedded:
-            self.interactive_predictor.set_image(image)
-            self.have_embedded = True
-    @torch.no_grad()
-    def interactive_predict(self, prompts, mode, multimask=True):
-        assert self.have_embedded, 'image embedding for sam need be set before predict.'        
-        
-        if mode == 'point':
-            masks, scores, logits = self.interactive_predictor.predict(point_coords=prompts['point_coords'], 
-                                point_labels=prompts['point_modes'], 
-                                multimask_output=multimask)
-        elif mode == 'mask':
-            masks, scores, logits = self.interactive_predictor.predict(mask_input=prompts['mask_prompt'], 
-                                multimask_output=multimask)
-        elif mode == 'point_mask':
-            masks, scores, logits = self.interactive_predictor.predict(point_coords=prompts['point_coords'], 
-                                point_labels=prompts['point_modes'], 
-                                mask_input=prompts['mask_prompt'], 
-                                multimask_output=multimask)
-                                
-        return masks, scores, logits
-        
-    @torch.no_grad()
-    def segment_with_click(self, origin_frame, coords, modes, multimask=True):
-        '''
-            
-            return: 
-                mask: one-hot 
-        '''
-        self.set_image(origin_frame)
-
-        prompts = {
-            'point_coords': coords,
-            'point_modes': modes,
-        }
-        masks, scores, logits = self.interactive_predict(prompts, 'point', multimask)
-        mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
-        prompts = {
-            'point_coords': coords,
-            'point_modes': modes,
-            'mask_prompt': logit[None, :, :]
-        }
-        masks, scores, logits = self.interactive_predict(prompts, 'point_mask', multimask)
-        mask = masks[np.argmax(scores)]
-
-        return mask.astype(np.uint8)
-
-    def segment_with_box(self, origin_frame, bbox, reset_image=False):
-        if reset_image:
-            self.interactive_predictor.set_image(origin_frame)
-        else:
-            self.set_image(origin_frame)
-        # coord = np.array([[int((bbox[1][0] - bbox[0][0]) / 2.),  int((bbox[1][1] - bbox[0][1]) / 2)]])
-        # point_label = np.array([1])
-
-        masks, scores, logits = self.interactive_predictor.predict(
-            point_coords=None,
-            point_labels=None,
-            box=np.array([bbox[0][0], bbox[0][1], bbox[1][0], bbox[1][1]]),
-            multimask_output=True
-        )
-        mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
-
-        masks, scores, logits = self.interactive_predictor.predict(
-            point_coords=None,
-            point_labels=None,
-            box=np.array([[bbox[0][0], bbox[0][1], bbox[1][0], bbox[1][1]]]),
-            mask_input=logit[None, :, :],
-            multimask_output=True
-        )
-        mask = masks[np.argmax(scores)]
-        
-        return [mask]

tool/transfer_tools.py

@@ -1,51 +0,0 @@
-import cv2
-import numpy as np 
-
-def mask2bbox(mask):
-    if len(np.where(mask > 0)[0]) == 0:
-        print(f'not mask')
-        return np.array([[0, 0], [0, 0]]).astype(np.int64)
-
-    x_ = np.sum(mask, axis=0)
-    y_ = np.sum(mask, axis=1)
-
-    x0 = np.min(np.nonzero(x_)[0])
-    x1 = np.max(np.nonzero(x_)[0])
-    y0 = np.min(np.nonzero(y_)[0])
-    y1 = np.max(np.nonzero(y_)[0])
-
-    return np.array([[x0, y0], [x1, y1]]).astype(np.int64)
-
-def draw_outline(mask, frame):
-    _, binary_mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)
-
-    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-    cv2.drawContours(frame, contours, -1, (0, 0, 255), 2)
-
-    return frame
-
-def draw_points(points, modes, frame):
-    neg_points = points[np.argwhere(modes==0)[:, 0]]
-    pos_points = points[np.argwhere(modes==1)[:, 0]]
-
-    for i in range(len(neg_points)):
-        point = neg_points[i]
-        cv2.circle(frame, (point[0], point[1]), 8, (255, 80, 80), -1)
-    
-    for i in range(len(pos_points)):
-        point = pos_points[i]
-        cv2.circle(frame, (point[0], point[1]), 8, (0, 153, 255), -1)
-
-    return frame
-
-if __name__ == '__main__':
-    mask = cv2.imread('./debug/mask.jpg', cv2.IMREAD_GRAYSCALE)
-    frame = cv2.imread('./debug/frame.jpg')
-    draw_frame = draw_outline(mask, frame)
-    
-    cv2.imwrite('./debug/outline.jpg', draw_frame)
-
-    # bbox = mask2bbox(mask)
-    # draw_0 = cv2.rectangle(mask, bbox[0], bbox[1], (0, 0, 255))
-    # cv2.imwrite('./debug/rect.png', draw_0)