def eval_seq(dataloader,
data_type,
result_filename,
save_dir=None,
show_image=True,
args=None):
if save_dir is not None:
mkdirs(save_dir)
tracker = OnlineTracker(metric_net=args.metric, ide=args.ide)
timer = Timer()
results = []
wait_time = 1
for frame_id, batch in enumerate(dataloader):
if frame_id % 20 == 0:
logger.info('Processing frame {} ({:.2f} fps)'.format(
frame_id, 1. / max(1e-5, timer.average_time)))
frame, det_tlwhs, det_scores, _, _ = batch
# run tracking
timer.tic()
online_targets = tracker.update(frame, det_tlwhs, None)
online_tlwhs = []
online_ids = []
for t in online_targets:
online_tlwhs.append(t.tlwh)
online_ids.append(t.track_id)
timer.toc()
# save results
results.append((frame_id + 1, online_tlwhs, online_ids))
online_im = vis.plot_tracking(frame,
online_tlwhs,
online_ids,
frame_id=frame_id,
fps=1. / timer.average_time)
if show_image:
cv2.imshow('online_im', online_im)
if save_dir is not None:
cv2.imwrite(os.path.join(save_dir, '{:05d}.jpg'.format(frame_id)),
online_im)
key = cv2.waitKey(wait_time)
key = chr(key % 128).lower()
if key == 'q':
exit(0)
elif key == 'p':
cv2.waitKey(0)
elif key == 'a':
wait_time = int(not wait_time)
# save results
write_results(result_filename, results, data_type)
class MOTDT(Expert):
def __init__(self, min_height, min_det_score):
super(MOTDT, self).__init__("MOTDT")
self.min_height = min_height
if min_det_score is None:
min_det_score = -np.inf
self.min_det_score = min_det_score
def initialize(self, seq_info):
super(MOTDT, self).initialize(seq_info)
self.tracker = OnlineTracker()
def track(self, img_path, dets):
super(MOTDT, self).track(img_path, dets)
im, tlwhs, scores = self.preprocess(img_path, dets)
online_targets = self.tracker.update(im, tlwhs, None)
results = []
for t in online_targets:
tracking_id = t.track_id
bbox = t.tlwh
results.append([tracking_id, bbox[0], bbox[1], bbox[2], bbox[3]])
return results
def preprocess(self, img_path, dets):
im = imread(img_path) # rgb
im = im[:, :, ::-1] # bgr
if dets is None:
return im, [], []
tlwhs = dets[:, 2:6]
scores = dets[:, 6]
keep = (tlwhs[:, 3] >= self.min_height) & (scores > self.min_det_score)
tlwhs = tlwhs[keep]
scores = scores[keep]
return im, tlwhs, scores
def evaluation(challenges=[],
resultDir="./results",
configs={},
visualize=False):
if not os.path.exists(resultDir):
os.makedirs(resultDir)
with open(os.path.join(resultDir, "info.txt"), 'w') as info:
for key in sorted(configs.keys()):
info.write(key + " : " + str(configs[key]) + "\r\n")
motdt = OnlineTracker(models={
"personDetect": yolov2,
"faceEmbed": facenet,
"superResolution": srgan,
"faceDetect": ssd
},
configs=configs)
for challenge in challenges:
motdt.clear()
evalName = challenge.split('/')[-2]
with open(os.path.join(resultDir, evalName + ".txt"), 'w') as result:
imageNames = os.listdir(challenge)
imageNames.sort()
for imageName in imageNames:
frame = cv2.imread(os.path.join(challenge, imageName),
cv2.IMREAD_COLOR)
# Tracking
online_targets = motdt.update(frame)
regions = []
ids = []
detects_face = []
detects_person = []
for t in online_targets:
regions.append(t.tlbr)
ids.append(t.track_id)
# write evaluation result
for region, id in zip(regions, ids):
print('%d,%d,%f,%f,%f,%f,1,-1,-1,-1' %
(int(imageName.split(".")[0]), id, region[0],
region[1], region[2] - region[0],
region[3] - region[1]),
file=result)
# Visualize
if visualize:
for region, id in zip(regions, ids):
cv2.rectangle(frame, (int(region[0]), int(region[1])),
(int(region[2]), int(region[3])),
(255, 255, 255), 2)
cv2.putText(frame, str(id),
(int(region[0]), int(region[1])), 0,
5e-3 * 200, (0, 255, 0), 2)
for detect in detects_person:
cv2.rectangle(frame, (int(detect[0]), int(detect[1])),
(int(detect[2]), int(detect[3])),
(255, 0, 0), 2)
for detect in detects_face:
cv2.rectangle(frame, (int(detect[0]), int(detect[1])),
(int(detect[2]), int(detect[3])),
(0, 0, 255), 2)
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
return
def eval_video(**kwargs):
logger.setLevel(logging.INFO)
cap = cv2.VideoCapture(kwargs['video_source'])
fps = cap.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(
*'MP4V') # int(cap.get(cv2.CAP_PROP_FOURCC))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_count = -1
iter_count = 0
each_frame = kwargs['each_frame']
save_dir = kwargs['save_dir']
frames_limit = kwargs['frames_limit']
video_writer = None
video_output = kwargs['video_output']
if video_output is not None:
logger.info(
f'Write video to {video_output} ({width}x{height}, {fps/each_frame} fps) ...'
)
video_writer = cv2.VideoWriter(video_output,
fourcc,
fps / each_frame,
frameSize=(width, height))
write_report_to = None
data = {}
if kwargs['report_output']:
write_report_to = kwargs['report_output']
tracker = OnlineTracker(**kwargs)
timer = Timer()
results = []
wait_time = 1
drv = driver.load_driver('tensorflow')
logger.info(f'init person detection driver...')
person_detect_driver = drv()
person_detect_model = kwargs['person_detect_model']
logger.info(f'loading person detection model {person_detect_model}...')
person_detect_driver.load_model(person_detect_model)
logger.info(f'person detection model {person_detect_model} loaded')
try:
while True:
frame_count += 1
if frames_limit is not None and frame_count > frames_limit:
logger.warn('frames limit {} reached'.format(frames_limit))
break
# read each X bgr frame
frame = cap.read() # bgr
if frame_count % each_frame > 0:
continue
if isinstance(frame, tuple):
frame = frame[1]
if frame is None:
logger.warn('video capturing finished')
break
if iter_count % 20 == 0:
logger.info(
'Processing frame {} (iteration {}) ({:.2f} fps)'.format(
frame_count, iter_count,
1. / max(1e-5, timer.average_time)))
det_tlwhs, det_scores = detect_persons_tf(person_detect_driver,
frame,
threshold=.5)
# run tracking
timer.tic()
online_targets = tracker.update(frame, det_tlwhs, None)
online_tlwhs = []
online_ids = []
for t in online_targets:
online_tlwhs.append(t.tlwh)
online_ids.append(t.track_id)
timer.toc()
if write_report_to:
for i, id in enumerate(online_ids):
if id not in data:
data[id] = {
'intervals': [],
'images': [],
'last_image': None,
}
di = data[id]['intervals']
if len(di) == 0 or di[-1][1] < frame_count - each_frame:
if len(di) > 0 and di[-1][0] == di[-1][1]:
di = di[:-1]
di.append([frame_count, frame_count])
else:
di[-1][1] = frame_count
if not data[id]['last_image'] or data[id][
'last_image'] < frame_count - fps * 10:
data[id]['last_image'] = frame_count
tlwh = [max(0, int(o)) for o in online_tlwhs[i]]
pers_img = frame[tlwh[1]:tlwh[1] + tlwh[3],
tlwh[0]:tlwh[0] + tlwh[2]].copy()
if max(pers_img.shape[0], pers_img.shape[1]) > 100:
coef = max(pers_img.shape[0],
pers_img.shape[1]) / 100
pers_img = cv2.resize(
pers_img, (int(pers_img.shape[1] / coef),
int(pers_img.shape[0] / coef)))
_, pers_img = cv2.imencode('.jpeg', pers_img)
data[id]['images'].append(
base64.b64encode(pers_img).decode())
# save results
frame_id = frame_count # or make it incremental?
results.append((frame_id + 1, online_tlwhs, online_ids))
online_im = vis.plot_tracking(frame,
online_tlwhs,
online_ids,
frame_id=frame_id,
fps=1. / timer.average_time)
for tlwh in det_tlwhs:
cv2.rectangle(
online_im,
(tlwh[0], tlwh[1]), # (left, top)
(tlwh[0] + tlwh[2], tlwh[1] + tlwh[3]), # (right, bottom)
(0, 255, 0),
1,
)
if kwargs['show_image']:
cv2.imshow('online_im', online_im)
if save_dir is not None:
save_to = os.path.join(save_dir, '{:05d}.jpg'.format(frame_id))
cv2.imwrite(save_to, online_im)
if video_writer is not None:
video_writer.write(cv2.resize(online_im, (width, height)))
key = cv2.waitKey(wait_time)
key = chr(key % 128).lower()
if key in [ord('q'), 202,
27]: # 'q' or Esc or 'q' in russian layout
exit(0)
elif key == 'p':
cv2.waitKey(0)
elif key == 'a':
wait_time = int(not wait_time)
iter_count += 1
except (KeyboardInterrupt, SystemExit) as e:
logger.info('Caught %s: %s' % (e.__class__.__name__, e))
finally:
cv2.destroyAllWindows()
if video_writer is not None:
logger.info('Written video to %s.' % video_output)
video_writer.release()
if write_report_to:
for i in data:
di = data[i]
di['index'] = i
di['duration'] = sum([i[1] - i[0] for i in di['intervals']])
di['duration_sec'] = '{:.2f}'.format(di['duration'] / fps)
di['intervals_str'] = ', '.join([
'{:.2f}-{:.2f}'.format(i[0] / fps, i[1] / fps)
for i in di['intervals']
])
data = data.values()
data = sorted(data, key=lambda x: x['duration'], reverse=True)
# prepare html
tpl = jinja2.Template(template)
html = tpl.render(data=data)
with open(write_report_to, 'w') as f:
f.write(html)
update_data({'#documents.persons.html': html}, use_mlboard,
mlboard)
def initialize(self, seq_info):
super(MOTDT, self).initialize(seq_info)
self.tracker = OnlineTracker()