diff --git a/deploy/pipeline/pphuman/mtmct.py b/deploy/pipeline/pphuman/mtmct.py index 50c789c75d9d4a641d0bf1c1dd45b039126bf26e..8ab72f4a4351d3872f1ae36881fc10f07653eae1 100644 --- a/deploy/pipeline/pphuman/mtmct.py +++ b/deploy/pipeline/pphuman/mtmct.py @@ -13,6 +13,7 @@ # limitations under the License. from pptracking.python.mot.visualize import plot_tracking +from python.visualize import visualize_attr import os import re import cv2 @@ -103,7 +104,8 @@ def get_mtmct_matching_results(pred_mtmct_file, secs_interval=0.5, return camera_results, cid_tid_fid_results -def save_mtmct_vis_results(camera_results, captures, output_dir): +def save_mtmct_vis_results(camera_results, captures, output_dir, + multi_res=None): # camera_results: 'cid, tid, fid, x1, y1, w, h' camera_ids = list(camera_results.keys()) @@ -126,7 +128,7 @@ def save_mtmct_vis_results(camera_results, captures, output_dir): height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = int(capture.get(cv2.CAP_PROP_FPS)) frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT)) - fourcc = cv2.VideoWriter_fourcc(* 'mp4v') + fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) frame_id = 0 while (1): @@ -143,6 +145,28 @@ def save_mtmct_vis_results(camera_results, captures, output_dir): boxes = frame_results[:, -4:] ids = frame_results[:, 1] image = plot_tracking(frame, boxes, ids, frame_id=frame_id, fps=fps) + + # add attr vis + if multi_res: + tid_list = [ + 'c' + str(idx) + '_' + 't' + str(int(j)) + for j in range(1, len(ids) + 1) + ] # c0_t1, c0_t2... + all_attr_result = [multi_res[i]["attrs"] + for i in tid_list] # all cid_tid result + if any( + all_attr_result + ): # at least one cid_tid[attrs] is not None will goes to attrs_vis + attr_res = [] + for k in tid_list: + if (frame_id - 1) >= len(multi_res[k]['attrs']): + t_attr = None + else: + t_attr = multi_res[k]['attrs'][frame_id - 1] + attr_res.append(t_attr) + image = visualize_attr( + image, attr_res, boxes, is_mtmct=True) + writer.write(image) writer.release() @@ -349,4 +373,8 @@ def mtmct_process(multi_res, captures, mtmct_vis=True, output_dir="output"): camera_results, cid_tid_fid_res = get_mtmct_matching_results( pred_mtmct_file) - save_mtmct_vis_results(camera_results, captures, output_dir=output_dir) + save_mtmct_vis_results( + camera_results, + captures, + output_dir=output_dir, + multi_res=cid_tid_dict) diff --git a/deploy/python/visualize.py b/deploy/python/visualize.py index 8f0da233df29ddc617cef99fe415b9ec7ed39efc..24aa40796f7ef9ae4d55df46f9a6223119183bec 100644 --- a/deploy/python/visualize.py +++ b/deploy/python/visualize.py @@ -331,7 +331,7 @@ def visualize_pose(imgfile, plt.close() -def visualize_attr(im, results, boxes=None): +def visualize_attr(im, results, boxes=None, is_mtmct=False): if isinstance(im, str): im = Image.open(im) im = np.ascontiguousarray(np.copy(im)) @@ -348,8 +348,12 @@ def visualize_attr(im, results, boxes=None): if boxes is None: text_w = 3 text_h = 1 + elif is_mtmct: + box = boxes[i] # multi camera, bbox shape is x,y, w,h + text_w = int(box[0]) + 3 + text_h = int(box[1]) else: - box = boxes[i] + box = boxes[i] # single camera, bbox shape is 0, 0, x,y, w,h text_w = int(box[2]) + 3 text_h = int(box[3]) for text in res: