add feature to calculate pq-score per area

PanopticQuality.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """TODO: Add a description here."""
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import evaluate
 import datasets
@@ -82,6 +82,7 @@ class PQMetric(evaluate.Metric):
         stuff: List[str] = None,
         per_class: bool = True,
         split_sq_rq: bool = True,
+        area_rng: List[Tuple[float]] = [(0, 1e8)],
         **kwargs
     ):
         super().__init__(**kwargs)
@@ -117,7 +118,8 @@ class PQMetric(evaluate.Metric):
             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]),
             return_per_class=per_class,
-            return_sq_and_rq=split_sq_rq
+            return_sq_and_rq=split_sq_rq,
+            areas=area_rng
         )
         self.things_stuffs = sorted(self.pq_metric.things) + sorted(self.pq_metric.stuffs)
 
@@ -174,22 +176,22 @@ class PQMetric(evaluate.Metric):
         iou = self.pq_metric.metric.iou_sum.clone()
 
         # compute scores
-        result = self.pq_metric.compute() # shape : (n_classes (sorted things + sorted stuffs), scores (pq, sq, rq))
+        result = self.pq_metric.compute() # shape : (area_rngs, n_classes (sorted things + sorted stuffs), scores (pq, sq, rq))
 
         result_dict = dict()
 
         if self.per_class:
             if not self.split_sq_rq:
-                result = result.T
-            result_dict["scores"] = {self.id2label[numeric_label]: result[i].tolist() \
+                result = result.unsqueeze(0)
+            result_dict["scores"] = {self.id2label[numeric_label]: result[:,:, i].tolist() \
                                      for i, numeric_label in enumerate(self.things_stuffs)}
-            result_dict["scores"].update({"ALL": result.mean(axis=0).tolist()})
-            result_dict["numbers"] =  {self.id2label[numeric_label]: [tp[i].item(), fp[i].item(), fn[i].item(), iou[i].item()] \
+            result_dict["scores"].update({"ALL": result.mean(dim=-1).tolist()})
+            result_dict["numbers"] =  {self.id2label[numeric_label]: [tp[:, i].tolist(), fp[:, i].tolist(), fn[:, i].tolist(), iou[:, i].tolist()] \
                                        for i, numeric_label in enumerate(self.things_stuffs)}
-            result_dict["numbers"].update({"ALL": [tp.sum().item(), fp.sum().item(), fn.sum().item(), iou.sum().item()]})
+            result_dict["numbers"].update({"ALL": [tp.sum(dim=1).tolist(), fp.sum(dim=1).tolist(), fn.sum(dim=1).tolist(), iou.sum(dim=1).tolist()]})
         else:
-            result_dict["scores"] = {"ALL": result.tolist() if self.split_sq_rq else [result.tolist()]}
-            result_dict["numbers"] = {"ALL": [tp.sum().item(), fp.sum().item(), fn.sum().item(), iou.sum().item()]}
+            result_dict["scores"] = {"ALL": result.tolist() if self.split_sq_rq else ([result.tolist()] if len(self.pq_metric.get_areas())>1 else [[result.tolist()]])}
+            result_dict["numbers"] = {"ALL": [tp.sum(dim=-1).tolist(), fp.sum(dim=-1).tolist(), fn.sum(dim=-1).tolist(), iou.sum(dim=-1).tolist()]}
         
         return result_dict