add Gait3D support

opengait/evaluation/evaluator.py

@@ -3,7 +3,7 @@ from time import strftime, localtime
 import numpy as np
 from utils import get_msg_mgr, mkdir
 
-from .metric import mean_iou, cuda_dist, compute_ACC_mAP
+from .metric import mean_iou, cuda_dist, compute_ACC_mAP, evaluate_rank
 from .re_rank import re_ranking
 
 
@@ -225,3 +225,43 @@ def evaluate_segmentation(data, dataset):
     miou = mean_iou(pred, labels)
     get_msg_mgr().log_info('mIOU: %.3f' % (miou.mean()))
     return {"scalar/test_accuracy/mIOU": miou}
+
+def evaluate_Gait3D(data, conf, metric='euc'):
+    msg_mgr = get_msg_mgr()
+
+    features, labels, cams, time_seqs = data['embeddings'], data['labels'], data['types'], data['views']
+    import json
+    probe_sets = json.load(
+        open('./datasets/Gait3D/Gait3D.json', 'rb'))['PROBE_SET']
+    probe_mask = []
+    for id, ty, sq in zip(labels, cams, time_seqs):
+        if '-'.join([id, ty, sq]) in probe_sets:
+            probe_mask.append(True)
+        else:
+            probe_mask.append(False)
+    probe_mask = np.array(probe_mask)
+
+    # probe_features = features[:probe_num]
+    probe_features = features[probe_mask]
+    # gallery_features = features[probe_num:]
+    gallery_features = features[~probe_mask]
+    # probe_lbls = np.asarray(labels[:probe_num])
+    # gallery_lbls = np.asarray(labels[probe_num:])
+    probe_lbls = np.asarray(labels)[probe_mask]
+    gallery_lbls = np.asarray(labels)[~probe_mask]
+
+    results = {}
+    msg_mgr.log_info(f"The test metric you choose is {metric}.")
+    dist = cuda_dist(probe_features, gallery_features, metric).cpu().numpy()
+    cmc, all_AP, all_INP = evaluate_rank(dist, probe_lbls, gallery_lbls)
+
+    mAP = np.mean(all_AP)
+    mINP = np.mean(all_INP)
+    for r in [1, 5, 10]:
+        results['scalar/test_accuracy/Rank-{}'.format(r)] = cmc[r - 1] * 100
+    results['scalar/test_accuracy/mAP'] = mAP * 100
+    results['scalar/test_accuracy/mINP'] = mINP * 100
+
+    # print_csv_format(dataset_name, results)
+    msg_mgr.log_info(results)
+    return results