fix bugs and add improvements

PanopticQuality.py

@@ -71,7 +71,6 @@ Examples:
     Added data ...
     Start computing ...
     Finished!
-    tensor(0.2082, dtype=torch.float64)
 """
 
 
@@ -81,6 +80,8 @@ class PQMetric(evaluate.Metric):
         self,
         label2id: Dict[str, int] = None,
         stuff: List[str] = None,
+        per_class: bool = True,
+        split_sq_rq: bool = True,
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -109,9 +110,13 @@ class PQMetric(evaluate.Metric):
 
         self.label2id = label2id if label2id is not None else DEFAULT_LABEL2ID
         self.stuff =  stuff if stuff is not None else DEFAULT_STUFF
+        self.per_class = per_class
+        self.split_sq_rq = split_sq_rq
         self.pq_metric = PanopticQuality(
             things=set([self.label2id[label] for label in self.label2id.keys() if label not in self.stuff]), 
-            stuffs=set([self.label2id[label] for label in self.label2id.keys() if label in self.stuff])
+            stuffs=set([self.label2id[label] for label in self.label2id.keys() if label in self.stuff]),
+            return_per_class=per_class,
+            return_sq_and_rq=split_sq_rq
         )
 
     def _info(self):
@@ -151,9 +156,6 @@ class PQMetric(evaluate.Metric):
         # in case the inputs are lists, convert them to numpy arrays
 
         self.pq_metric.update(prediction, reference)
-        print("TP:", self.pq_metric.metric.true_positives)
-        print("FP:", self.pq_metric.metric.false_positives)
-        print("FN:", self.pq_metric.metric.false_negatives)
 
         # does not impact the metric, but is required for the interface x_x
         super(evaluate.Metric, self).add(
@@ -164,12 +166,30 @@ class PQMetric(evaluate.Metric):
 
     def _compute(self, *, predictions, references, **kwargs):
         """Called within the evaluate.Metric.compute() method"""
-        result = self.pq_metric.compute() # n_classes (sorted things + sorted stuffs), (pq, sq, rq)
+        tp = self.pq_metric.metric.true_positives.clone()
+        fp = self.pq_metric.metric.false_positives.clone()
+        fn = self.pq_metric.metric.false_negatives.clone()
+        iou = self.pq_metric.metric.iou_sum.clone()
+
         id2label = {id: label for label, id in self.label2id.items()}
-        return {
-            id2label[numeric_label]: result[i] for i, numeric_label in enumerate(self.pq_metric.things + self.pq_metric.stuffs)
+        things_stuffs = sorted(self.pq_metric.things) + sorted(self.pq_metric.stuffs)
+
+        # 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
         }
 
+        if self.per_class:
+            result_dict["scores"] = {id2label[numeric_label]: result[i].tolist() for i, numeric_label in enumerate(things_stuffs)}
+        else:
+            result_dict["scores"] = result.tolist()
+        
+        return result_dict
+
     def add_payload(self, payload: Payload, model_name: str = None):
         """Converts the payload to the format expected by the metric"""
         # import only if needed since fiftyone is not a direct dependency