Add FiftyOneMistakenness callback

Author: lumurillo

mart/callbacks/fiftyone.py

@@ -14,10 +14,11 @@
 logger = logging.getLogger(__name__)
 try:
     import fiftyone as fo
+    import fiftyone.brain as fob
 except ImportError:
     logger.debug("fiftyone module is not installed!")
 
-__all__ = ["FiftyOneEvaluateDetections", "FiftyOnePredictionAdder"]
+__all__ = ["FiftyOneEvaluateDetections", "FiftyOneMistakenness", "FiftyOnePredictionAdder"]
 
 
 class FiftyOneEvaluateDetections(Callback):
@@ -48,6 +49,25 @@ def on_predict_end(self, trainer, pl_module):
         results.print_report(classes=classes_top10)
 
 
+class FiftyOneMistakenness(Callback):
+    def __init__(self, run_id: str, gt_field: str = "ground_truth_detections") -> None:
+        self.prediction_field = f"prediction_{run_id}"
+        self.gt_field = gt_field
+
+    def on_predict_start(self, trainer, pl_module):
+        self.predict_dataset = trainer.datamodule.predict_dataset
+        assert isinstance(self.predict_dataset, FiftyOneDataset)
+
+        # reset mistakenness fields
+        if self.predict_dataset.dataset.has_brain_run("mistakenness"):
+            self.predict_dataset.dataset.delete_brain_run("mistakenness")
+
+    def on_predict_end(self, trainer, pl_module):
+        fob.compute_mistakenness(
+            self.predict_dataset.filtered_dataset, self.prediction_field, label_field=self.gt_field
+        )
+
+
 class FiftyOnePredictionAdder(BasePredictionWriter):
     def __init__(self, output_dir: str, write_interval: List[str]) -> None:
         super().__init__(write_interval)

mart/configs/callbacks/fiftyone_mistakenness.yaml

@@ -0,0 +1,3 @@
+fiftyone_mistakenness:
+  _target_: mart.callbacks.FiftyOneMistakenness
+  run_id: ${now:%Y-%m-%d}_${now:%H-%M-%S}

mart/configs/experiment/FiftyOne_TorchvisionFasterRCNN.yaml

@@ -9,6 +9,7 @@ defaults:
         lr_monitor,
         fiftyone_prediction_adder,
         fiftyone_evaluate_detections,
+        fiftyone_mistakenness,
       ]
 
 task_name: "FiftyOne_TorchvisionFasterRCNN"