Anti-UAV比赛测试代码解读
"""
baseline for 1st Anti-UAV
https://anti-uav.github.io/
Qiang Wang
2020.02.16
"""
from __future__ import absolute_import
import os
import glob
import json
import cv2
import numpy as np
from siamfc import TrackerSiamFC
def iou(bbox1, bbox2):
"""
Calculates the intersection-over-union of two bounding boxes.
Args:
bbox1 (numpy.array, list of floats): bounding box in format x,y,w,h.
bbox2 (numpy.array, list of floats): bounding box in format x,y,w,h.
Returns:
int: intersection-over-onion of bbox1, bbox2
"""
bbox1 = [float(x) for x in bbox1]
bbox2 = [float(x) for x in bbox2]
(x0_1, y0_1, w1_1, h1_1) = bbox1
(x0_2, y0_2, w1_2, h1_2) = bbox2
x1_1 = x0_1 + w1_1
x1_2 = x0_2 + w1_2
y1_1 = y0_1 + h1_1
y1_2 = y0_2 + h1_2
# get the overlap rectangle
overlap_x0 = max(x0_1, x0_2)
overlap_y0 = max(y0_1, y0_2)
overlap_x1 = min(x1_1, x1_2)
overlap_y1 = min(y1_1, y1_2)
# check if there is an overlap
if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0:
return 0
# if yes, calculate the ratio of the overlap to each ROI size and the unified size
size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
size_union = size_1 + size_2 - size_intersection
return size_intersection / size_union
def not_exist(pred):
return (len(pred) == 1 and pred[0] == 0) or len(pred) == 0
def eval(out_res, label_res):
measure_per_frame = []
for _pred, _gt, _exist in zip(out_res, label_res['gt_rect'], label_res['exist']):
measure_per_frame.append(not_exist(_pred) if not _exist else iou(_pred, _gt) if len(_pred) > 1 else 0)
return np.mean(measure_per_frame)
def main(mode='IR', visulization=False):
assert mode in ['IR', 'RGB'], 'Only Support IR or RGB to evalute'
# setup tracker 设定跟踪器的预训练模型，并获得跟踪器
net_path = 'model.pth'
tracker = TrackerSiamFC(net_path=net_path)
# setup experiments
video_paths = glob.glob(os.path.join('dataset', 'test-challenge', '*')) # 数据集的路径
video_num = len(video_paths) # 视频序列的数量
output_dir = os.path.join('results', tracker.name) # 输出的路径
if not os.path.exists(output_dir):
os.makedirs(output_dir)
overall_performance = []
# run tracking experiments and report performance
for video_id, video_path in enumerate(video_paths, start=1):
video_name = os.path.basename(video_path)
video_file = os.path.join(video_path, '%s.mp4'%mode) # 视频序列文件
res_file = os.path.join(video_path, '%s_label.json'%mode) # ground truth文件
with open(res_file, 'r') as f:
label_res = json.load(f)
init_rect = label_res['gt_rect'][0] # 目标初始bounding box
capture = cv2.VideoCapture(video_file)
frame_id = 0
out_res = []
while True:
ret, frame = capture.read()
if not ret:
capture.release()
break
if frame_id == 0:
tracker.init(frame, init_rect) # initialization 跟踪器初始化
out = init_rect
out_res.append(init_rect)
else:
out = tracker.update(frame) # tracking
out_res.append(out.tolist())
# 可视化ground truth，但test-challenge文件夹只有首帧信息，所以这里False
if visulization:
_gt = label_res['gt_rect'][frame_id]
_exist = label_res['exist'][frame_id]
if _exist:
cv2.rectangle(frame, (int(_gt[0]), int(_gt[1])), (int(_gt[0] + _gt[2]), int(_gt[1] + _gt[3])),
(0, 255, 0))
cv2.putText(frame, 'exist' if _exist else 'not exist',
(frame.shape[1] // 2 - 20, 30), 1, 2, (0, 255, 0) if _exist else (0, 0, 255), 2)
cv2.rectangle(frame, (int(out[0]), int(out[1])), (int(out[0] + out[2]), int(out[1] + out[3])),
(0, 255, 255))
cv2.imshow(video_name, frame)
cv2.waitKey(1)
frame_id += 1
if visulization:
cv2.destroyAllWindows()
# save result
output_file = os.path.join(output_dir, '%s_%s.txt' % (video_name, mode))
# 编码成JSON数据
with open(output_file, 'w') as f:
json.dump({'res': out_res}, f)
# 结果评估，跳转到eval函数
mixed_measure = eval(out_res, label_res)
overall_performance.append(mixed_measure)
print('[%03d/%03d] %20s %5s Fixed Measure: %.03f' % (video_id, video_num, video_name, mode, mixed_measure))
print('[Overall] %5s Mixed Measure: %.03f\n' % (mode, np.mean(overall_performance)))
if __name__ == '__main__':
main(mode='IR', visulization=False)
函数作用：结果评估
def eval(out_res, label_res):
measure_per_frame = []
# 如果groundtruth中，目标不存在，执行not_exist函数；如果groundtruth中，目标存在，执行iou函数，计算预测结果与ground truth的iou。
for _pred, _gt, _exist in zip(out_res, label_res['gt_rect'], label_res['exist']):
measure_per_frame.append(not_exist(_pred) if not _exist else iou(_pred, _gt) if len(_pred) > 1 else 0)
return np.mean(measure_per_frame)
函数作用：计算预测结果与ground truth的iou
def iou(bbox1, bbox2):
"""
Calculates the intersection-over-union of two bounding boxes.
Args:
bbox1 (numpy.array, list of floats): bounding box in format x,y,w,h.
bbox2 (numpy.array, list of floats): bounding box in format x,y,w,h.
Returns:
int: intersection-over-onion of bbox1, bbox2
"""
bbox1 = [float(x) for x in bbox1]
bbox2 = [float(x) for x in bbox2]
(x0_1, y0_1, w1_1, h1_1) = bbox1
(x0_2, y0_2, w1_2, h1_2) = bbox2
# 得到右下角坐标
x1_1 = x0_1 + w1_1
x1_2 = x0_2 + w1_2
y1_1 = y0_1 + h1_1
y1_2 = y0_2 + h1_2
# get the overlap rectangle
overlap_x0 = max(x0_1, x0_2)
overlap_y0 = max(y0_1, y0_2)
overlap_x1 = min(x1_1, x1_2)
overlap_y1 = min(y1_1, y1_2)
# check if there is an overlap
if overlap_x1 - overlap_x0 <= 0 or overlap_y1 - overlap_y0 <= 0: # 如果不存在重叠区
return 0
# if yes, calculate the ratio of the overlap to each ROI size and the unified size
size_1 = (x1_1 - x0_1) * (y1_1 - y0_1)
size_2 = (x1_2 - x0_2) * (y1_2 - y0_2)
size_intersection = (overlap_x1 - overlap_x0) * (overlap_y1 - overlap_y0)
size_union = size_1 + size_2 - size_intersection
return size_intersection / size_union
def not_exist(pred):
return (len(pred) == 1 and pred[0] == 0) or len(pred) == 0