Repository: airctic/icevision

Issue:
Fix plot_top_losses
`plot_top_losses` got broken because of the new refactor, the type of the data passed to it now changed, we just need to make the tests pass.

What were `samples` before is not just `Record`s and don't have a `dict` like interface anymore.

@FraPochetti would you be interested in this one?


Discussion:
Sure, but not immediately. I think it might take a couple of months or so, given I am on yolov5 now. 
> Sure, but not immediately. I think it might take a couple of months or so, given I am on yolov5 now.

Okay, no problem!!
@lgvaz 
Those "couple of months" that I mention are completely indicative btw :D (just wanted to be conservative with yolo).
I really hope to be done sooner than that! and in any case, this fix should not take much time so I might just slide it somewhere when I have a couple of free hours ;)
Is this just the tests that are broken or the entire functionality?
> Is this just the tests that are broken or the entire functionality?

What broke is that before `sample` was a dict and now it's a `Record`, so getting attributes with  `["labels"]` and calling `.update` doesn't work anymore

But don't worry, it should be an straight forward fix on my side, I'll take a look in a few days
I prefer doing that actually 😉. It will be the occasion to check out the new API.
I will take a look and keep you posted! 
I have made up my mind and decided to look into this one tonight 😁. I think it is important to have all tests passing ASAP. Yolo can wait another couple of days.
Coming back from a weekend out of town. Will check this out as soon as I step at home 😉

Resolving patch:
diff --git a/icevision/core/record_components.py b/icevision/core/record_components.py
--- a/icevision/core/record_components.py
+++ b/icevision/core/record_components.py
@@ -14,6 +14,7 @@
     "IsCrowdsRecordComponent",
     "KeyPointsRecordComponent",
     "ScoresRecordComponent",
+    "LossesRecordComponent",
 ]
 
 from icevision.imports import *
@@ -470,3 +471,12 @@ def _repr(self) -> List[str]:
 
     def as_dict(self) -> dict:
         return {"scores": self.scores}
+
+
+class LossesRecordComponent(RecordComponent):
+    def __init__(self, task=tasks.common):
+        super().__init__(task=task)
+        self.losses = None
+
+    def set_losses(self, losses: Dict):
+        self.losses = losses
diff --git a/icevision/models/interpretation.py b/icevision/models/interpretation.py
--- a/icevision/models/interpretation.py
+++ b/icevision/models/interpretation.py
@@ -12,13 +12,14 @@
 from icevision.core import *
 from icevision.data import *
 from icevision.visualize.show_data import show_preds
+from icevision.core.record_components import LossesRecordComponent
 
 
 def get_weighted_sum(sample, weights):
     loss_weighted = 0
     for loss, weight in weights.items():
-        loss_weighted += sample[loss] * weight
-    sample["loss_weighted"] = loss_weighted
+        loss_weighted += sample.losses[loss] * weight
+    sample.losses["loss_weighted"] = loss_weighted
     return sample
 
 
@@ -27,7 +28,7 @@ def sort_losses(
 ) -> Tuple[List[dict], List[dict], List[str]]:
     by_copy = deepcopy(by)
     losses_expected = [
-        k for k in samples[0].keys() if "loss" in k and k != "loss_total"
+        k for k in samples[0].losses.keys() if "loss" in k and k != "loss_total"
     ]
     if "effdet_total_loss" in losses_expected:
         losses_expected.remove("effdet_total_loss")
@@ -54,14 +55,14 @@ def sort_losses(
             by = "loss_weighted"
 
     l = list(zip(samples, preds))
-    l = sorted(l, key=lambda i: i[0][by], reverse=True)
+    l = sorted(l, key=lambda i: i[0].losses[by], reverse=True)
     sorted_samples, sorted_preds = zip(*l)
-    annotations = [el["text"] for el in sorted_samples]
+    annotations = [el.losses["text"] for el in sorted_samples]
 
     if isinstance(by_copy, dict):
         if by_copy["method"] == "weighted":
             annotations = [
-                f"loss_weighted: {round(s['loss_weighted'], 5)}\n" + a
+                f"loss_weighted: {round(s.losses['loss_weighted'], 5)}\n" + a
                 for a, s in zip(annotations, sorted_samples)
             ]
 
@@ -143,7 +144,9 @@ def _loop(self, dl, model, losses_stats, device):
                     losses_stats[l].append(loss[l])
 
                 loss = _prepend_str(loss, "loss")
-                sample[0].update(loss)
+                loss_comp = LossesRecordComponent()
+                loss_comp.set_losses(loss)
+                sample[0].add_component(loss_comp)
                 samples_plus_losses.append(sample[0])
         return samples_plus_losses, losses_stats
 
@@ -219,6 +222,7 @@ def plot_top_losses(
 
         dl = self.infer_dl(dataset, batch_size=batch_size)
         preds = self.predict_dl(model=model, infer_dl=dl)
+        preds = [p.pred for p in preds]
 
         sorted_samples, sorted_preds, annotations = sort_losses(
             samples, preds, by=sort_by
@@ -232,8 +236,12 @@ def plot_top_losses(
             ann2 = "\n".join(ann[4:])
             anns.append((ann1, ann2))
 
+        sorted_preds = [
+            Prediction(pred=p, ground_truth=s)
+            for s, p in zip(sorted_samples, sorted_preds)
+        ]
+
         show_preds(
-            samples=sorted_samples[:n_samples],
             preds=sorted_preds[:n_samples],
             annotations=anns[:n_samples],
         )
@@ -247,11 +255,11 @@ def add_annotations(samples: List[dict]) -> List[dict]:
     """
     for sample in samples:
         text = ""
-        for key in sample.keys():
+        for key in sample.losses.keys():
             if "loss" in key:
-                text += f"{key}: {round(sample[key], 5)}\n"
-        text += f"IMG: {sample['filepath'].name}"
-        sample["text"] = text
+                text += f"{key}: {round(sample.losses[key], 5)}\n"
+        text += f"IMG: {sample.filepath.name}"
+        sample.losses["text"] = text
     return samples
 
 
diff --git a/icevision/models/mmdet/common/interpretation_utils.py b/icevision/models/mmdet/common/interpretation_utils.py
--- a/icevision/models/mmdet/common/interpretation_utils.py
+++ b/icevision/models/mmdet/common/interpretation_utils.py
@@ -5,6 +5,7 @@
 from icevision.core import *
 from icevision.data import *
 from icevision.models.interpretation import _move_to_device
+from icevision.core.record_components import LossesRecordComponent
 
 
 def sum_losses_mmdet(losses_dict):
@@ -33,7 +34,9 @@ def loop_mmdet(dl, model, losses_stats, device):
             for l in losses_stats.keys():
                 losses_stats[l].append(loss[l])
 
-            sample[0].update(loss)
+            loss_comp = LossesRecordComponent()
+            loss_comp.set_losses(loss)
+            sample[0].add_component(loss_comp)
             samples_plus_losses.append(sample[0])
 
     return samples_plus_losses, losses_stats


Test patch:
diff --git a/tests/models/torchvision_models/faster_rcnn/test_show_results.py b/tests/models/torchvision_models/faster_rcnn/test_show_results.py
--- a/tests/models/torchvision_models/faster_rcnn/test_show_results.py
+++ b/tests/models/torchvision_models/faster_rcnn/test_show_results.py
@@ -64,8 +64,8 @@ def test_get_losses(fridge_faster_rcnn_model, fridge_ds):
         "loss_rpn_box_reg",
         "loss_total",
     }
-    assert "loss_box_reg" in samples[0].keys()
-    assert "text" not in samples[0].keys()
+    assert "loss_box_reg" in samples[0].losses.keys()
+    assert "text" not in samples[0].losses.keys()
 
 
 def test_add_annotations(fridge_faster_rcnn_model, fridge_ds):
@@ -74,5 +74,5 @@ def test_add_annotations(fridge_faster_rcnn_model, fridge_ds):
 
     samples, _ = faster_rcnn.interp.get_losses(model, ds)
     samples = add_annotations(samples)
-    assert "loss_classifier" in samples[0]["text"]
-    assert "IMG" in samples[0]["text"]
+    assert "loss_classifier" in samples[0].losses["text"]
+    assert "IMG" in samples[0].losses["text"]
diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py
--- a/tests/utils/test_utils.py
+++ b/tests/utils/test_utils.py
@@ -1,4 +1,5 @@
 from icevision.all import *
+from icevision.core.record_components import LossesRecordComponent
 
 
 def test_notnone():
@@ -61,12 +62,17 @@ def test_get_stats():
     assert result == expected
 
 
-def test_sort_losses():
+def _set_losses_lrc(s):
+    br = LossesRecordComponent()
+    br.set_losses(s)
+    return br
+
 
+def test_sort_losses():
     samples = [
-        {"stuff": 0.2, "loss_total": 2, "text": "text2"},
-        {"stuff": 0.1, "loss_total": 1, "text": "text1"},
-        {"stuff": 0.3, "loss_total": 3, "text": "text3"},
+        _set_losses_lrc({"stuff": 0.2, "loss_total": 2, "text": "text2"}),
+        _set_losses_lrc({"stuff": 0.1, "loss_total": 1, "text": "text1"}),
+        _set_losses_lrc({"stuff": 0.3, "loss_total": 3, "text": "text3"}),
     ]
 
     preds = [
@@ -89,7 +95,7 @@ def test_sort_losses():
 
     sorted_samples, sorted_preds, annotations = sort_losses(samples, preds)
 
-    assert sorted_samples == sorted_samples_ex
+    assert [s.losses for s in sorted_samples] == sorted_samples_ex
     assert sorted_preds == sorted_preds_ex
     assert annotations == ["text3", "text2", "text1"]
 
@@ -107,8 +113,9 @@ def test_get_weighted_sum():
         "loss3": 0.25,
         "loss4": 0.25,
     }
+    br = _set_losses_lrc(s)
     expected = {"loss1": 1, "loss2": 1, "loss3": 1, "loss4": 1, "loss_weighted": 1.0}
 
-    result = get_weighted_sum(s, weights)
+    result = get_weighted_sum(br, weights)
 
-    assert result["loss_weighted"] == expected["loss_weighted"]
+    assert result.losses["loss_weighted"] == expected["loss_weighted"]