add filters

user-friendly-metrics.py

@@ -89,11 +89,12 @@ class UserFriendlyMetrics(evaluate.Metric):
     def _compute(self, 
                  payload, 
                  max_iou: float = 0.5, 
+                 filters = {},
                  recognition_thresholds = [0.3, 0.5, 0.8], 
                  debug: bool = False):
         """Returns the scores"""
         # TODO: Compute the different scores of the module
-        return calculate_from_payload(payload, max_iou, recognition_thresholds, debug) 
+        return calculate_from_payload(payload, max_iou, filters, recognition_thresholds, debug) 
         #return calculate(predictions, references, max_iou)
 
 def recognition(track_ratios, th = 0.5):
@@ -161,6 +162,7 @@ def calculate(predictions,
 
 def calculate_from_payload(payload: dict, 
                            max_iou: float = 0.5,
+                           filters = {},
                            recognition_thresholds = [0.3, 0.5, 0.8], 
                            debug: bool = False):
     if not isinstance(payload, dict):
@@ -184,12 +186,27 @@ def calculate_from_payload(payload: dict,
     for sequence in sequence_list:
         output[sequence] = {}
         frames = payload['sequences'][sequence][gt_field_name]
-        formated_references = []
+        
+        all_formated_references = {"all": []}
+        for filter, filter_ranges in filters:
+            all_formated_references[filter] = {}
+            for filter_range in filter_ranges:
+                filter_range_name = filter_range[0]
+                all_formated_references[filter][filter_range_name] = []
+
         for frame_id, frame in enumerate(frames):
             for detection in frame:
-                id = detection['index']
+                index = detection['index']
                 x, y, w, h = detection['bounding_box']
-                formated_references.append([frame_id+1, id, x, y, w, h])  
+                all_formated_references["all"].append([frame_id+1, index, x, y, w, h])  
+                for filter, filter_ranges in filters:
+                    filter_value = detection[filter]
+                    for filter_range in filter_ranges:
+                        filter_range_name = filter_range[0]
+                        filter_range_limits = filter_range[1]
+                        if filter_value >= filter_range_limits[0] and filter_value <= filter_range_limits[1]:
+                            all_formated_references[filter][filter_range_name].append([frame_id+1, index, x, y, w, h])
+
 
         for model in models:
             frames = payload['sequences'][sequence][model]
@@ -197,21 +214,28 @@ def calculate_from_payload(payload: dict,
             
             for frame_id, frame in enumerate(frames):
                 for detection in frame:
-                    id = detection['index']
+                    index = detection['index']
                     x, y, w, h = detection['bounding_box']
                     confidence = detection['confidence']
                     confidence = 1  #TODO: remove this line
-                    formated_predictions.append([frame_id+1, id, x, y, w, h, confidence])
+                    formated_predictions.append([frame_id+1, index, x, y, w, h, confidence])
             if debug:
                 print("sequence/model: ", sequence, model)
                 print("formated_predictions: ", formated_predictions)
-                print("formated_references: ", formated_references)
+                print("formated_references: ", all_formated_references)
             if len(formated_predictions) == 0:
                 output[sequence][model] = "Model had no predictions."
-            elif len(formated_references) == 0:
+            elif len(all_formated_references["all"]) == 0:
                 output[sequence][model] = "No ground truth."
             else:
-                output[sequence][model] = calculate(formated_predictions, formated_references, max_iou=max_iou, recognition_thresholds = recognition_thresholds)
+                output[sequence][model] = {}
+                output[sequence][model]["all"] = calculate(formated_predictions, all_formated_references["all"], max_iou=max_iou, recognition_thresholds = recognition_thresholds)
+                for filter, filter_ranges in filters:
+                    output[sequence][model][filter] = {}
+                    for filter_range in filter_ranges:
+                        filter_range_name = filter_range[0]
+                        output[sequence][model][filter][filter_range_name] = calculate(formated_predictions, all_formated_references[filter][filter_range_name], max_iou=max_iou, recognition_thresholds = recognition_thresholds)
+
     return output