update mtmct.py and visualize.py for multi-camera attrs visualize (#6580)

* update mtmct.py and visualize.py for multi-camera attrs visualize

* update mtmct.py

* update mcmct_attrvisualzie logic, without enable_attr will not goes to attrs visualize

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)

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: