tqdm temp

user-friendly-metrics.py

@@ -15,6 +15,8 @@
 import datetime
 import os
 
+from tqdm import tqdm
+
 import datasets
 import evaluate
 from seametrics.user_friendly.utils import payload_to_uf_metrics, UFM
@@ -71,6 +73,7 @@ class UserFriendlyMetrics(evaluate.Metric):
         self.iou_threshold = iou_threshold
         self.filter_dict = filter_dict
         self.recognition_thresholds = recognition_thresholds
+        self.metric = UFM(iou_threshold, recognition_thresholds)
 
 
     def _info(self):
@@ -116,19 +119,14 @@ class UserFriendlyMetrics(evaluate.Metric):
 
             for sequence_predictions, sequence_references in zip(predictions, references):
                 
-                ufm = UFM(
-                    iou_threshold=self.iou_threshold,
-                    recognition_thresholds=self.recognition_thresholds
-                )
-                
-                sequence_range_results = ufm.calculate(
+                sequence_range_results = self.metric.calculate(
                     sequence_predictions,
                     sequence_references[filter_range_name],
                 )
 
                 range_results = sum_dicts(range_results, sequence_range_results)
 
-            results[filter_range_name] = ufm.realize_metrics(range_results, self.recognition_thresholds)
+            results[filter_range_name] = self.metric.realize_metrics(range_results, self.recognition_thresholds)
         
         return results
 
@@ -142,7 +140,9 @@ class UserFriendlyMetrics(evaluate.Metric):
             results[model_name] = {"overall": {}, "per_sequence": {}}
             
             # per-sequence loop
-            for seq_name, sequence in payload.sequences.items():
+            progress_bar = tqdm(payload.sequences.items())
+            for seq_name, sequence in progress_bar:
+                progress_bar.set_description(f"Getting sequence payload: {seq_name}")
                 # create new payload only with specific sequence and model
                 sequence_payload = Payload(
                     dataset=payload.dataset,
@@ -150,7 +150,7 @@ class UserFriendlyMetrics(evaluate.Metric):
                     models=[model_name],
                     sequences={seq_name: sequence}
                 )
-                
+                progress_bar.set_description(f"Processing sequence: {seq_name}")
                 predictions, references = payload_to_uf_metrics(payload, model_name=model_name, filter_dict=self.filter_dict)
 
                 results[model_name]["per_sequence"][seq_name] = self._compute(predictions=predictions, references=references)