improve return structure

PanopticQuality.py

@@ -177,16 +177,20 @@ class PQMetric(evaluate.Metric):
         # compute scores
         result = self.pq_metric.compute() # shape : (n_classes (sorted things + sorted stuffs), scores (pq, sq, rq))
 
-        result_dict = {
-            "numbers": {id2label[numeric_label]: [tp[i].item(), fp[i].item(), fn[i].item(), iou[i].item()] \
-                        for i, numeric_label in enumerate(things_stuffs)},
-            "scores": None
-        }
+        result_dict = dict()
 
         if self.per_class:
-            result_dict["scores"] = {id2label[numeric_label]: result[i].tolist() for i, numeric_label in enumerate(things_stuffs)}
+            if not self.split_sq_rq:
+                result = result.T
+            result_dict["scores"] = {id2label[numeric_label]: result[i].tolist() \
+                                     for i, numeric_label in enumerate(things_stuffs)}
+            result_dict["scores"].update({"overall": result.mean(axis=0).tolist()})
+            result_dict["numbers"] =  {id2label[numeric_label]: [tp[i].item(), fp[i].item(), fn[i].item(), iou[i].item()] \
+                                       for i, numeric_label in enumerate(things_stuffs)}
+            result_dict["numbers"].update({"overall": [tp.sum().item(), fp.sum().item(), fn.sum().item(), iou.sum().item()]})
         else:
-            result_dict["scores"] = result.tolist()
+            result_dict["scores"] = result.tolist() if self.split_sq_rq else [result.tolist()]
+            result_dict["numbers"] = [tp.sum().item(), fp.sum().item(), fn.sum().item(), iou.sum().item()]
         
         return result_dict