import sys sys.path.append("..") sys.path.append("./sam") from sam.segment_anything import sam_model_registry, SamAutomaticMaskGenerator from aot_tracker import get_aot import numpy as np from tool.segmentor import Segmentor from tool.detector import Detector from tool.transfer_tools import draw_outline, draw_points import cv2 from seg_track_anything import draw_mask class SegTracker(): def __init__(self,segtracker_args, sam_args, aot_args) -> None: """ Initialize SAM and AOT. """ self.sam = Segmentor(sam_args) self.tracker = get_aot(aot_args) self.detector = Detector(self.sam.device) self.sam_gap = segtracker_args['sam_gap'] self.min_area = segtracker_args['min_area'] self.max_obj_num = segtracker_args['max_obj_num'] self.min_new_obj_iou = segtracker_args['min_new_obj_iou'] self.reference_objs_list = [] self.object_idx = 1 self.curr_idx = 1 self.origin_merged_mask = None # init by segment-everything or update self.first_frame_mask = None # debug self.everything_points = [] self.everything_labels = [] print("SegTracker has been initialized") def seg(self,frame): ''' Arguments: frame: numpy array (h,w,3) Return: origin_merged_mask: numpy array (h,w) ''' frame = frame[:, :, ::-1] anns = self.sam.everything_generator.generate(frame) # anns is a list recording all predictions in an image if len(anns) == 0: return # merge all predictions into one mask (h,w) # note that the merged mask may lost some objects due to the overlapping self.origin_merged_mask = np.zeros(anns[0]['segmentation'].shape,dtype=np.uint8) idx = 1 for ann in anns: if ann['area'] > self.min_area: m = ann['segmentation'] self.origin_merged_mask[m==1] = idx idx += 1 self.everything_points.append(ann["point_coords"][0]) self.everything_labels.append(1) obj_ids = np.unique(self.origin_merged_mask) obj_ids = obj_ids[obj_ids!=0] self.object_idx = 1 for id in obj_ids: if np.sum(self.origin_merged_mask==id) < self.min_area or self.object_idx > self.max_obj_num: self.origin_merged_mask[self.origin_merged_mask==id] = 0 else: self.origin_merged_mask[self.origin_merged_mask==id] = self.object_idx self.object_idx += 1 self.first_frame_mask = self.origin_merged_mask return self.origin_merged_mask def update_origin_merged_mask(self, updated_merged_mask): self.origin_merged_mask = updated_merged_mask # obj_ids = np.unique(updated_merged_mask) # obj_ids = obj_ids[obj_ids!=0] # self.object_idx = int(max(obj_ids)) + 1 def reset_origin_merged_mask(self, mask, id): self.origin_merged_mask = mask self.curr_idx = id def add_reference(self,frame,mask,frame_step=0): ''' Add objects in a mask for tracking. Arguments: frame: numpy array (h,w,3) mask: numpy array (h,w) ''' self.reference_objs_list.append(np.unique(mask)) self.curr_idx = self.get_obj_num() + 1 self.tracker.add_reference_frame(frame,mask, self.curr_idx - 1, frame_step) def track(self,frame,update_memory=False): ''' Track all known objects. Arguments: frame: numpy array (h,w,3) Return: origin_merged_mask: numpy array (h,w) ''' pred_mask = self.tracker.track(frame) if update_memory: self.tracker.update_memory(pred_mask) return pred_mask.squeeze(0).squeeze(0).detach().cpu().numpy().astype(np.uint8) def get_tracking_objs(self): objs = set() for ref in self.reference_objs_list: objs.update(set(ref)) objs = list(sorted(list(objs))) objs = [i for i in objs if i!=0] return objs def get_obj_num(self): objs = self.get_tracking_objs() if len(objs) == 0: return 0 return int(max(objs)) def find_new_objs(self, track_mask, seg_mask): ''' Compare tracked results from AOT with segmented results from SAM. Select objects from background if they are not tracked. Arguments: track_mask: numpy array (h,w) seg_mask: numpy array (h,w) Return: new_obj_mask: numpy array (h,w) ''' new_obj_mask = (track_mask==0) * seg_mask new_obj_ids = np.unique(new_obj_mask) new_obj_ids = new_obj_ids[new_obj_ids!=0] # obj_num = self.get_obj_num() + 1 obj_num = self.curr_idx for idx in new_obj_ids: new_obj_area = np.sum(new_obj_mask==idx) obj_area = np.sum(seg_mask==idx) if new_obj_area/obj_area < self.min_new_obj_iou or new_obj_area < self.min_area\ or obj_num > self.max_obj_num: new_obj_mask[new_obj_mask==idx] = 0 else: new_obj_mask[new_obj_mask==idx] = obj_num obj_num += 1 return new_obj_mask def restart_tracker(self): self.tracker.restart() def seg_acc_bbox(self, origin_frame: np.ndarray, bbox: np.ndarray,): '''' Use bbox-prompt to get mask Parameters: origin_frame: H, W, C bbox: [[x0, y0], [x1, y1]] Return: refined_merged_mask: numpy array (h, w) masked_frame: numpy array (h, w, c) ''' # get interactive_mask interactive_mask = self.sam.segment_with_box(origin_frame, bbox)[0] refined_merged_mask = self.add_mask(interactive_mask) # draw mask masked_frame = draw_mask(origin_frame.copy(), refined_merged_mask) # draw bbox masked_frame = cv2.rectangle(masked_frame, bbox[0], bbox[1], (0, 0, 255)) return refined_merged_mask, masked_frame def seg_acc_click(self, origin_frame: np.ndarray, coords: np.ndarray, modes: np.ndarray, multimask=True): ''' Use point-prompt to get mask Parameters: origin_frame: H, W, C coords: nd.array [[x, y]] modes: nd.array [[1]] Return: refined_merged_mask: numpy array (h, w) masked_frame: numpy array (h, w, c) ''' # get interactive_mask interactive_mask = self.sam.segment_with_click(origin_frame, coords, modes, multimask) refined_merged_mask = self.add_mask(interactive_mask) # draw mask masked_frame = draw_mask(origin_frame.copy(), refined_merged_mask) # draw points # self.everything_labels = np.array(self.everything_labels).astype(np.int64) # self.everything_points = np.array(self.everything_points).astype(np.int64) masked_frame = draw_points(coords, modes, masked_frame) # draw outline masked_frame = draw_outline(interactive_mask, masked_frame) return refined_merged_mask, masked_frame def add_mask(self, interactive_mask: np.ndarray): ''' Merge interactive mask with self.origin_merged_mask Parameters: interactive_mask: numpy array (h, w) Return: refined_merged_mask: numpy array (h, w) ''' if self.origin_merged_mask is None: self.origin_merged_mask = np.zeros(interactive_mask.shape,dtype=np.uint8) refined_merged_mask = self.origin_merged_mask.copy() refined_merged_mask[interactive_mask > 0] = self.curr_idx return refined_merged_mask def detect_and_seg(self, origin_frame: np.ndarray, grounding_caption, box_threshold, text_threshold, box_size_threshold=1, reset_image=False): ''' Using Grounding-DINO to detect object acc Text-prompts Retrun: refined_merged_mask: numpy array (h, w) annotated_frame: numpy array (h, w, 3) ''' # backup id and origin-merged-mask bc_id = self.curr_idx bc_mask = self.origin_merged_mask # get annotated_frame and boxes annotated_frame, boxes = self.detector.run_grounding(origin_frame, grounding_caption, box_threshold, text_threshold) for i in range(len(boxes)): bbox = boxes[i] if (bbox[1][0] - bbox[0][0]) * (bbox[1][1] - bbox[0][1]) > annotated_frame.shape[0] * annotated_frame.shape[1] * box_size_threshold: continue interactive_mask = self.sam.segment_with_box(origin_frame, bbox, reset_image)[0] refined_merged_mask = self.add_mask(interactive_mask) self.update_origin_merged_mask(refined_merged_mask) self.curr_idx += 1 # reset origin_mask self.reset_origin_merged_mask(bc_mask, bc_id) return refined_merged_mask, annotated_frame if __name__ == '__main__': from model_args import segtracker_args,sam_args,aot_args Seg_Tracker = SegTracker(segtracker_args, sam_args, aot_args) # ------------------ detect test ---------------------- origin_frame = cv2.imread('/data2/cym/Seg_Tra_any/Segment-and-Track-Anything/debug/point.png') origin_frame = cv2.cvtColor(origin_frame, cv2.COLOR_BGR2RGB) grounding_caption = "swan.water" box_threshold = 0.25 text_threshold = 0.25 predicted_mask, annotated_frame = Seg_Tracker.detect_and_seg(origin_frame, grounding_caption, box_threshold, text_threshold) masked_frame = draw_mask(annotated_frame, predicted_mask) origin_frame = cv2.cvtColor(origin_frame, cv2.COLOR_RGB2BGR) cv2.imwrite('./debug/masked_frame.png', masked_frame) cv2.imwrite('./debug/x.png', annotated_frame)