✨ [AAAI|Add] meanBoxCoverScore validation metric

Author: nathalie000

yolo/aaai.py

@@ -18,7 +18,7 @@
 
 from yolo import Config, PostProcess, create_converter, create_model, draw_bboxes
 from yolo.tools.data_loader import AAAIDataset
-from yolo.tools.loss_functions import AAAILoss, NT_Xent
+from yolo.tools.loss_functions import AAAILoss, NT_Xent, meanBoxCoverScore
 from yolo.utils.logging_utils import YOLORichModelSummary, setup
 from yolo.utils.model_utils import create_optimizer, create_scheduler
 
@@ -30,6 +30,7 @@ def __init__(self, cfg: Config, model):
         self.construct_loss = nn.MSELoss()
         self.contrastive_loss = NT_Xent
         self.cfg = cfg
+        self.metric = meanBoxCoverScore()
 
     def set_task(self, task):
         self.task = task
@@ -140,7 +141,21 @@ def validation_step(self, batch, batch_idx):
 
             origin_outputs = self(images, dict(target=picked_vector.permute(0, 2, 1)))
             H, W = images.shape[2:]
-            predicts = self.post_process(origin_outputs, image_size=[W, H], target=picked_vector)
+
+            unpacked_outputs = []
+            num_target = self.cfg.task.data.num_target
+            repeat_num = [(num_target, 1, 1, 1), (num_target, 1, 1, 1, 1)]
+            for batch_idx in range(batch_size):
+                batch_sample = [
+                    [x[[batch_idx]].repeat(*repeat_num[idx % 2]) for idx, x in enumerate(res_list)]
+                    for res_list in origin_outputs["Main"]
+                ]
+                unpacked_outputs.append(batch_sample)
+
+            for outputs, tar_bbox, tar_vec in zip(unpacked_outputs, bbox, picked_vector):
+                outputs = {"Main": outputs}
+                predicts = self.post_process(outputs, image_size=[W, H], target=tar_vec[:, None])
+                self.metric(predicts, tar_bbox)
 
         return predicts
 
@@ -165,6 +180,10 @@ def on_validation_epoch_end(self):
             logger = self.logger.experiment
             logger.add_histogram("Validation Error Distribution", Tensor(self.FIDs), self.current_epoch)
 
+        epoch_metric = self.metric.compute()
+        self.log("detection_metric", epoch_metric, prog_bar=True, sync_dist=True, rank_zero_only=True)
+        self.metric.reset()
+
     def configure_optimizers(self):
         optimizer = create_optimizer(self.model, self.cfg.task.optimizer)
         scheduler = create_scheduler(optimizer, self.cfg.task.scheduler)

yolo/tools/loss_functions.py

@@ -6,6 +6,7 @@
 from torch import Tensor, nn
 from torch.nn import BCEWithLogitsLoss, MSELoss
 from torch.nn.functional import cosine_similarity, cross_entropy
+from torchmetrics import MeanMetric
 
 from yolo.config.config import Config, LossConfig, MatcherConfig
 from yolo.utils.bounding_box_utils import BoxMatcher, Vec2Box, calculate_iou
@@ -354,3 +355,24 @@ def create_loss_function(cfg: Config, vec2box) -> DualLoss:
     loss_function = DualLoss(cfg, vec2box)
     logger.info(":white_check_mark: Success load loss function")
     return loss_function
+
+
+class meanBoxCoverScore(MeanMetric):
+    def __init__(self, hit_iou=0.7, **kwargs):
+        super().__init__(**kwargs)
+        self.hit_iou = hit_iou
+
+    def update(self, predicts: List[Tensor], bboxes: Tensor) -> None:
+        for idx, predict in enumerate(predicts):
+            if predict.size(0) == 0:
+                super().update(0)
+                continue
+            predict = predict[predict[:, 5].argsort(descending=True)]
+            iou_matrix = calculate_iou(predict[:, 1:5], bboxes[idx])
+            best_hit = (iou_matrix > self.hit_iou).float()
+            rank_score = best_hit.argmax().where(best_hit.any(), float("inf"))
+            score = (1 / (rank_score + 1)).mean()
+            super().update(score)
+
+    def compute(self) -> Tensor:
+        return super().compute()