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")
