def _mask_post_processing(self, mask): target_mask = (mask > cfg.TRACK.MASK_THERSHOLD) target_mask = target_mask.astype(np.uint8) if cv2.__version__[-5] == '4': contours, _ = cv2.findContours(target_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) else: _, contours, _ = cv2.findContours(target_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) cnt_area = [cv2.contourArea(cnt) for cnt in contours] if len(contours) != 0 and np.max(cnt_area) > 100: contour = contours[np.argmax(cnt_area)] polygon = contour.reshape(-1, 2) prbox = cv2.boxPoints(cv2.minAreaRect(polygon)) rbox_in_img = prbox else: # empty mask location = cxy_wh_2_rect(self.center_pos, self.size) rbox_in_img = np.array([[location[0], location[1]], [location[0] + location[2], location[1]], [location[0] + location[2], location[1] + location[3]], [location[0], location[1] + location[3]]]) return rbox_in_img
def _mask_post_processing(self, mask): target_mask = (mask > cfg.TRACK.MASK_THERSHOLD) target_mask = target_mask.astype(np.uint8) if cv2.__version__[-5] == '4': contours, _ = cv2.findContours(target_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) else: _, contours, _ = cv2.findContours(target_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) cnt_area = [cv2.contourArea(cnt) for cnt in contours] if len(contours) != 0 and np.max(cnt_area) > 100: contour = contours[np.argmax(cnt_area)] polygon = contour.reshape(-1, 2) ## the following code estimate the shape angle with ellipse ## then fit a axis-aligned bounding box on the rotated image ellipseBox = cv2.fitEllipse(polygon) # get the center of the ellipse and the angle angle = ellipseBox[-1] #print(angle) center = np.array(ellipseBox[0]) axes = np.array(ellipseBox[1]) # get the ellipse box ellipseBox = cv2.boxPoints(ellipseBox) #compute the rotation matrix rot_mat = cv2.getRotationMatrix2D((center[0], center[1]), angle, 1.0) # rotate the ellipse box one = np.ones([ellipseBox.shape[0], 3, 1]) one[:, :2, :] = ellipseBox.reshape(-1, 2, 1) ellipseBox = np.matmul(rot_mat, one).reshape(-1, 2) # to xmin ymin xmax ymax xs = ellipseBox[:, 0] xmin, xmax = np.min(xs), np.max(xs) ys = ellipseBox[:, 1] ymin, ymax = np.min(ys), np.max(ys) ellipseBox = [xmin, ymin, xmax, ymax] # rotate the contour one = np.ones([polygon.shape[0], 3, 1]) one[:, :2, :] = polygon.reshape(-1, 2, 1) polygon = np.matmul(rot_mat, one).astype(int).reshape(-1, 2) # remove points outside of the ellipseBox logi = polygon[:, 0] <= xmax logi = np.logical_and(polygon[:, 0] >= xmin, logi) logi = np.logical_and(polygon[:, 1] >= ymin, logi) logi = np.logical_and(polygon[:, 1] <= ymax, logi) polygon = polygon[logi, :] x, y, w, h = cv2.boundingRect(polygon) bRect = [x, y, x + w, y + h] # get the intersection of ellipse box and the rotated box x1, y1, x2, y2 = ellipseBox[0], ellipseBox[1], ellipseBox[ 2], ellipseBox[3] tx1, ty1, tx2, ty2 = bRect[0], bRect[1], bRect[2], bRect[3] xx1 = min(max(tx1, x1, 0), target_mask.shape[1] - 1) yy1 = min(max(ty1, y1, 0), target_mask.shape[0] - 1) xx2 = max(min(tx2, x2, target_mask.shape[1] - 1), 0) yy2 = max(min(ty2, y2, target_mask.shape[0] - 1), 0) rotated_mask = cv2.warpAffine( target_mask, rot_mat, (target_mask.shape[1], target_mask.shape[0])) #refinement alpha_factor = cfg.TRACK.FACTOR while True: if np.sum(rotated_mask[int(yy1):int(yy2), int(xx1)]) < (yy2 - yy1) * alpha_factor: temp = xx1 + (xx2 - xx1) * 0.02 if not (temp >= target_mask.shape[1] - 1 or xx2 - xx1 < 1): xx1 = temp else: break else: break while True: if np.sum(rotated_mask[int(yy1):int(yy2), int(xx2)]) < (yy2 - yy1) * alpha_factor: temp = xx2 - (xx2 - xx1) * 0.02 if not (temp <= 0 or xx2 - xx1 < 1): xx2 = temp else: break else: break while True: if np.sum(rotated_mask[int(yy1), int(xx1):int(xx2)]) < ( xx2 - xx1) * alpha_factor: temp = yy1 + (yy2 - yy1) * 0.02 if not (temp >= target_mask.shape[0] - 1 or yy2 - yy1 < 1): yy1 = temp else: break else: break while True: if np.sum(rotated_mask[int(yy2), int(xx1):int(xx2)]) < ( xx2 - xx1) * alpha_factor: temp = yy2 - (yy2 - yy1) * 0.02 if not (temp <= 0 or yy2 - yy1 < 1): yy2 = temp else: break else: break prbox = np.array([[xx1, yy1], [xx2, yy1], [xx2, yy2], [xx1, yy2]]) # inverse of the rotation matrix M_inv = cv2.invertAffineTransform(rot_mat) # project the points back to image coordinate one = np.ones([prbox.shape[0], 3, 1]) one[:, :2, :] = prbox.reshape(-1, 2, 1) prbox = np.matmul(M_inv, one).reshape(-1, 2) rbox_in_img = prbox else: # empty mask location = cxy_wh_2_rect(self.center_pos, self.size) rbox_in_img = np.array( [[location[0], location[1]], [location[0] + location[2], location[1]], [location[0] + location[2], location[1] + location[3]], [location[0], location[1] + location[3]]]) return rbox_in_img
def main(): # load config cfg.merge_from_file(args.config) # create model model = ModelBuilder() # load model model = load_pretrain(model, args.snapshot).cuda().eval() # build tracker tracker = build_tracker(model) #model_name = args.snapshot.split('/')[-1].split('.')[0] #total_lost = 0 #cur_dir = os.path.dirname(os.path.realpath(__file__)) #dataset_root = os.path.join(cur_dir, '../testing_dataset', args.dataset) video_path = '/home/yuuzhao/Documents/project/pysot/testing_dataset/VOT2016' #lists = open('/home/lichao/tracking/LaSOT_Evaluation_Toolkit/sequence_evaluation_config/' + setfile + '.txt', 'r') #list_file = [line.strip() for line in lists] category = os.listdir(video_path) category.sort() # create dataset #dataset = DatasetFactory.create_dataset(name=args.dataset,dataset_root=dataset_root,load_img=False) template_acc = [] template_cur = [] init0 = [] init = [] pre = [] gt = [] # init0 is reset init print("Category & Video:") for tmp_cat in category: tmp_cat_path = temp_path + '/' + tmp_cat if not os.path.isdir(tmp_cat_path): os.makedirs(tmp_cat_path) print("Category:", tmp_cat) video = os.listdir(join(video_path, tmp_cat)) video.sort() #video_cut = video[0:frames_of_each_video] frame = 0 #for picture in video_cut: # 这个循环或许该去掉 # print("Frame:", picture) gt_path = join(video_path, tmp_cat, 'groundtruth.txt') ground_truth = np.loadtxt(gt_path, delimiter=',') # num_frames = len(ground_truth); # num_frames = min(num_frames, frame_max) num_frames = frames_of_each_video # print("num_frames: ",num_frames) img_path = join(video_path, tmp_cat) # print("imgpath",img_path) imgFiles = [ join(img_path, '%08d.jpg') % i for i in range(1, num_frames + 1) ] while frame < num_frames: print("frame:", frame) Polygon = ground_truth[frame] cx, cy, w, h = get_axis_aligned_bbox(Polygon) gt_rect = [cx, cy, w, h] image_file = imgFiles[frame] # target_pos, target_sz = np.array([cx, cy]), np.array([w, h]) img = cv2.imread(image_file) # HxWxC if frame == 0: tracker.init(img, gt_rect) if w * h != 0: # image_file = imgFiles[frame] # target_pos, target_sz = np.array([cx, cy]), np.array([w, h]) # img = cv2.imread(image_file) # HxWxC zf_acc = tracker.get_zf(img, gt_rect) output = tracker.track(img) pre_rect = output['bbox'] zf_pre = tracker.get_zf(img, pre_rect) template_acc.append(zf_acc) template_cur.append((zf_pre)) print("ACC&PRE") init0.append(0) init.append(frame) frame_reset = 0 pre.append(0) gt.append(1) while frame < (num_frames - 1): print("while ", frame, "<", num_frames) frame = frame + 1 frame_reset = frame_reset + 1 image_file = imgFiles[frame] if not image_file: break Polygon = ground_truth[frame] cx, cy, w, h = get_axis_aligned_bbox(Polygon) gt_rect = [cx, cy, w, h] img = cv2.imread(image_file) # HxWxC zf_acc = tracker.get_zf(img, gt_rect) output = tracker.track(img) pre_rect = output['bbox'] zf_pre = tracker.get_zf(img, pre_rect) # print("zf_pre:",zf_pre.shape) # print("zf_acc:",zf_acc.shape) # pdb.set_trace() template_acc.append(zf_acc) template_cur.append(zf_pre) init0.append(frame_reset) init.append(frame) pre.append(1) if frame == (num_frames - 1): # last frame print("if frame == num_frames-1") gt.append(0) else: gt.append(1) pre_rect_arr = np.array(pre_rect) cx, cy, w, h = get_axis_aligned_bbox(pre_rect_arr) target_pos, target_siz = np.array([cx, cy]), np.array([w, h]) res = cxy_wh_2_rect(target_pos, target_siz) if reset: cx, cy, w, h = get_axis_aligned_bbox( ground_truth[frame]) gt_rect = [cx, cy, w, h] gt_rect = np.array(gt_rect) iou = overlap_ratio(gt_rect, res) if iou <= 0: break else: print("else") template_acc.append( torch.zeros([1, 3, 127, 127], dtype=torch.float32)) template_cur.append( torch.zeros([1, 3, 127, 127], dtype=torch.float32)) init0.append(0) init.append(frame) pre.append(1) if frame == (num_frames - 1): # last frame gt.append(0) else: gt.append(1) frame = frame + 1 # skip #写出一次 #print("template_acc:",template_acc) #print("template_cur:",template_cur) #print("init:", init) #print("init0:",init0) #print("pre:",pre) #template_acc_con = np.concatenate(template_acc); #template_cur_con = np.concatenate(template_cur) print("write for each video") np.save(tmp_cat_path + '/template', template_acc) np.save(tmp_cat_path + '/templatei', template_cur) np.save(tmp_cat_path + '/init0', init0) np.save(tmp_cat_path + '/init', init) np.save(tmp_cat_path + '/pre', pre) np.save(tmp_cat_path + '/gt', gt) print("template")