feat: max_eval_orig_size to restrain eval images loaded to a maximum size on the longest side

src/rfdetr/config.py

@@ -503,6 +503,7 @@ class TrainConfig(BaseModel):
     eval_max_dets: int = 500
     eval_interval: int = 1
     log_per_class_metrics: bool = True
+    max_eval_orig_size: Optional[int] = None
     aug_config: Optional[Dict[str, Any]] = None
     augmentation_backend: Literal["cpu", "auto", "gpu"] = "cpu"
     save_dataset_grids: bool = False

src/rfdetr/training/callbacks/coco_eval.py

@@ -48,6 +48,12 @@ class COCOEvalCallback(Callback):
         eval_interval: Run validation metrics every N epochs. Test metrics are
             always computed when ``trainer.test()`` is called.
         log_per_class_metrics: When ``False``, skip per-class AP logging/table.
+        max_eval_orig_size: Cap the longer side of each image's original size to
+            this value (in pixels) before upsampling masks for COCO evaluation.
+            Reduces mask buffer memory proportionally to ``(cap/orig)²`` with
+            negligible impact on mAP since the model input is already much
+            smaller (e.g. SegNano trains at 312 px). ``None`` disables the cap
+            and evaluates at full original resolution. Defaults to ``None``.
     """
 
     def __init__(
@@ -57,12 +63,14 @@ def __init__(
         eval_interval: int = 1,
         log_per_class_metrics: bool = True,
         in_notebook: bool | None = None,
+        max_eval_orig_size: int | None = None,
     ) -> None:
         super().__init__()
         self._max_dets = max_dets
         self._segmentation = segmentation
         self._eval_interval = max(1, int(eval_interval))
         self._log_per_class_metrics = bool(log_per_class_metrics)
+        self._max_eval_orig_size = max_eval_orig_size
         self._class_names: list[str] = []
         self._cat_id_to_name: dict[int, str] = {}
         self._f1_local: dict[int, dict[str, Any]] = init_matching_accumulator()
@@ -183,6 +191,12 @@ def on_validation_batch_end(
                 ).to(pl_module.device)
             samples, _ = batch
             orig_sizes = torch.stack([t["orig_size"] for t in outputs["targets"]]).to(pl_module.device)
+            if self._max_eval_orig_size is not None:
+                # Cap each (H, W) so the longer side ≤ max_eval_orig_size.
+                # Masks are upsampled to orig_size inside postprocess; capping here
+                # keeps mask buffers small without affecting training.
+                scale = (self._max_eval_orig_size / orig_sizes.float().amax(dim=1, keepdim=True)).clamp(max=1.0)
+                orig_sizes = (orig_sizes.float() * scale).long()
             ema_underlying = ema_cb._average_model.module.model
             with torch.no_grad():
                 ema_underlying.eval()  # AveragedModel deepcopy is not managed by PTL
@@ -708,8 +722,24 @@ def _convert_preds(self, preds: list[dict[str, torch.Tensor]]) -> list[dict[str,
         out = []
         for p in preds:
             entry = dict(p)
-            if "masks" in entry and entry["masks"].ndim == 4 and entry["masks"].shape[1] == 1:
-                entry["masks"] = entry["masks"].squeeze(1)
+            if "masks" in entry:
+                masks = entry["masks"]
+                if masks.ndim == 4 and masks.shape[1] == 1:
+                    masks = masks.squeeze(1)
+                if self._max_eval_orig_size is not None:
+                    h, w = masks.shape[-2:]
+                    if max(h, w) > self._max_eval_orig_size:
+                        scale = self._max_eval_orig_size / max(h, w)
+                        new_h = max(1, int(h * scale))
+                        new_w = max(1, int(w * scale))
+                        masks = F.interpolate(
+                            masks.float().unsqueeze(1),
+                            size=(new_h, new_w),
+                            mode="nearest",
+                        ).squeeze(1)
+                        if "boxes" in entry:
+                            entry["boxes"] = entry["boxes"] * scale
+                entry["masks"] = masks
             out.append(entry)
         return out
 
@@ -729,13 +759,17 @@ def _convert_targets(self, targets: list[dict[str, torch.Tensor]]) -> list[dict[
         out = []
         for t in targets:
             h, w = t["orig_size"].tolist()
+            if self._max_eval_orig_size is not None and max(h, w) > self._max_eval_orig_size:
+                cap_scale = self._max_eval_orig_size / max(h, w)
+                h = max(1, int(h * cap_scale))
+                w = max(1, int(w * cap_scale))
             scale = t["boxes"].new_tensor([w, h, w, h])
             boxes = box_cxcywh_to_xyxy(t["boxes"]) * scale
             entry: dict[str, torch.Tensor] = {"boxes": boxes, "labels": t["labels"]}
             if "masks" in t:
                 masks = t["masks"].bool()
-                # PostProcess resizes predicted masks to orig_size; resize GT
-                # masks to match so that mask-IoU comparisons are size-consistent.
+                # PostProcess resizes predicted masks to (h, w); resize GT masks
+                # to match so that mask-IoU comparisons are size-consistent.
                 if masks.shape[-2:] != (int(h), int(w)):
                     masks = (
                         F.interpolate(

src/rfdetr/training/module_model.py

@@ -220,6 +220,10 @@ def validation_step(self, batch: Tuple, batch_idx: int) -> Dict[str, Any]:
             self.log("val/loss", loss, prog_bar=True, on_epoch=True, sync_dist=True, batch_size=len(targets))
 
         orig_sizes = torch.stack([t["orig_size"] for t in targets])
+        cap = getattr(self.train_config, "max_eval_orig_size", None)
+        if cap is not None:
+            scale = (cap / orig_sizes.float().amax(dim=1, keepdim=True)).clamp(max=1.0)
+            orig_sizes = (orig_sizes.float() * scale).long()
         results = self.postprocess(outputs, orig_sizes)
         return {"results": results, "targets": targets}
 
@@ -355,6 +359,10 @@ def test_step(self, batch: Tuple, batch_idx: int) -> Dict[str, Any]:
             self.log("test/loss", loss, sync_dist=True, batch_size=len(targets))
 
         orig_sizes = torch.stack([t["orig_size"] for t in targets])
+        cap = getattr(self.train_config, "max_eval_orig_size", None)
+        if cap is not None:
+            scale = (cap / orig_sizes.float().amax(dim=1, keepdim=True)).clamp(max=1.0)
+            orig_sizes = (orig_sizes.float() * scale).long()
         results = self.postprocess(outputs, orig_sizes)
         return {"results": results, "targets": targets}
 

src/rfdetr/training/trainer.py

@@ -211,6 +211,7 @@ def _resolve_precision() -> str:
             segmentation=model_config.segmentation_head,
             eval_interval=tc.eval_interval,
             log_per_class_metrics=tc.log_per_class_metrics,
+            max_eval_orig_size=tc.max_eval_orig_size,
         )
     )
 

tests/models/test_config.py

@@ -232,6 +232,7 @@ def test_compute_test_loss_default_is_true(self, tmp_path):
             pytest.param("train_log_sync_dist", True, id="train_log_sync_dist"),
             pytest.param("train_log_on_step", True, id="train_log_on_step"),
             pytest.param("log_per_class_metrics", False, id="log_per_class_metrics"),
+            pytest.param("max_eval_orig_size", 640, id="max_eval_orig_size"),
             pytest.param("prefetch_factor", 4, id="prefetch_factor"),
             pytest.param("pin_memory", False, id="pin_memory"),
             pytest.param("persistent_workers", False, id="persistent_workers"),
@@ -261,6 +262,15 @@ def test_interval_and_prefetch_reject_non_positive_values(self, tmp_path, field,
         with pytest.raises((ValueError, ValidationError)):
             self._tc(tmp_path, **{field: value})
 
+    def test_max_eval_orig_size_defaults_to_none(self, tmp_path):
+        """max_eval_orig_size defaults to None (no cap — evaluates at full original resolution)."""
+        assert self._tc(tmp_path).max_eval_orig_size is None
+
+    def test_max_eval_orig_size_inherited_by_segmentation_train_config(self):
+        """SegmentationTrainConfig inherits max_eval_orig_size from TrainConfig."""
+        tc = SegmentationTrainConfig(dataset_dir="/tmp", max_eval_orig_size=640)
+        assert tc.max_eval_orig_size == 640
+
     def test_batch_size_auto_is_accepted(self, tmp_path):
         """batch_size accepts the special 'auto' value."""
         tc = self._tc(tmp_path, batch_size="auto")

tests/training/test_coco_eval_callback.py

@@ -747,3 +747,152 @@ def test_no_masks_no_iscrowd_keys_absent(self) -> None:
         ]
         out = cb._convert_targets(targets)
         assert set(out[0].keys()) == {"boxes", "labels"}
+
+
+class TestMaxEvalOrigSize:
+    """max_eval_orig_size caps mask resolution and box coordinates for COCO eval."""
+
+    # ------------------------------------------------------------------
+    # _convert_preds: mask downsampling
+    # ------------------------------------------------------------------
+
+    def test_convert_preds_downsamples_masks_when_above_cap(self) -> None:
+        """Masks larger than cap are downsampled to fit the longer side."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        preds = [
+            {
+                "masks": torch.ones(2, 1, 1080, 1920),
+                "boxes": torch.zeros(2, 4),
+                "scores": torch.ones(2),
+                "labels": torch.zeros(2, dtype=torch.long),
+            }
+        ]
+        out = cb._convert_preds(preds)
+        h, w = out[0]["masks"].shape[-2:]
+        assert max(h, w) == 640
+
+    def test_convert_preds_preserves_aspect_ratio_on_downsample(self) -> None:
+        """Downsampled masks preserve the original aspect ratio."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        preds = [
+            {
+                "masks": torch.ones(1, 1, 1080, 1920),
+                "boxes": torch.zeros(1, 4),
+                "scores": torch.ones(1),
+                "labels": torch.zeros(1, dtype=torch.long),
+            }
+        ]
+        out = cb._convert_preds(preds)
+        h, w = out[0]["masks"].shape[-2:]
+        assert abs(w / h - 1920 / 1080) < 0.02
+
+    def test_convert_preds_scales_boxes_with_masks(self) -> None:
+        """Box coordinates are scaled by the same factor as the mask downsample."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        boxes = torch.tensor([[100.0, 200.0, 300.0, 400.0]])
+        preds = [
+            {
+                "masks": torch.ones(1, 1, 1080, 1920),
+                "boxes": boxes,
+                "scores": torch.ones(1),
+                "labels": torch.zeros(1, dtype=torch.long),
+            }
+        ]
+        out = cb._convert_preds(preds)
+        scale = 640 / 1920
+        expected = boxes * scale
+        assert torch.allclose(out[0]["boxes"], expected, atol=1e-4)
+
+    def test_convert_preds_no_op_when_within_cap(self) -> None:
+        """Masks already within the cap are not resized."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        preds = [
+            {
+                "masks": torch.ones(1, 400, 600),
+                "boxes": torch.zeros(1, 4),
+                "scores": torch.ones(1),
+                "labels": torch.zeros(1, dtype=torch.long),
+            }
+        ]
+        out = cb._convert_preds(preds)
+        assert out[0]["masks"].shape[-2:] == (400, 600)
+
+    def test_convert_preds_no_op_when_cap_is_none(self) -> None:
+        """When max_eval_orig_size is None no downsampling occurs regardless of mask size."""
+        cb = COCOEvalCallback(max_eval_orig_size=None)
+        preds = [
+            {
+                "masks": torch.ones(1, 1, 1080, 1920),
+                "boxes": torch.zeros(1, 4),
+                "scores": torch.ones(1),
+                "labels": torch.zeros(1, dtype=torch.long),
+            }
+        ]
+        out = cb._convert_preds(preds)
+        assert out[0]["masks"].shape[-2:] == (1080, 1920)
+
+    # ------------------------------------------------------------------
+    # _convert_targets: orig_size capping
+    # ------------------------------------------------------------------
+
+    def test_convert_targets_caps_gt_mask_size(self) -> None:
+        """GT masks are resized to the capped (h, w), not the original orig_size."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        targets = [
+            {
+                "boxes": torch.tensor([[0.5, 0.5, 0.1, 0.1]]),
+                "labels": torch.tensor([0]),
+                "orig_size": torch.tensor([1080, 1920]),
+                "masks": torch.ones(1, 1080, 1920, dtype=torch.bool),
+            }
+        ]
+        out = cb._convert_targets(targets)
+        h, w = out[0]["masks"].shape[-2:]
+        assert max(h, w) == 640
+
+    def test_convert_targets_scales_boxes_to_capped_size(self) -> None:
+        """GT boxes are scaled to the capped (h, w) coordinate space."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        # Normalised box centred at image centre, half the image size
+        targets = [
+            {
+                "boxes": torch.tensor([[0.5, 0.5, 1.0, 1.0]]),
+                "labels": torch.tensor([0]),
+                "orig_size": torch.tensor([1080, 1920]),
+            }
+        ]
+        out = cb._convert_targets(targets)
+        scale = 640 / 1920
+        capped_h = int(1080 * scale)
+        capped_w = 640
+        # box_cxcywh_to_xyxy([0.5,0.5,1.0,1.0]) * [w,h,w,h] → [0,0,w,h]
+        expected = torch.tensor([[0.0, 0.0, float(capped_w), float(capped_h)]])
+        assert torch.allclose(out[0]["boxes"], expected, atol=1.0)
+
+    def test_convert_targets_no_op_when_cap_is_none(self) -> None:
+        """With max_eval_orig_size=None, GT boxes use full orig_size as scale."""
+        cb = COCOEvalCallback(max_eval_orig_size=None)
+        targets = [
+            {
+                "boxes": torch.tensor([[0.5, 0.5, 1.0, 1.0]]),
+                "labels": torch.tensor([0]),
+                "orig_size": torch.tensor([1080, 1920]),
+            }
+        ]
+        out = cb._convert_targets(targets)
+        expected = torch.tensor([[0.0, 0.0, 1920.0, 1080.0]])
+        assert torch.allclose(out[0]["boxes"], expected, atol=1e-4)
+
+    def test_convert_targets_no_op_when_within_cap(self) -> None:
+        """Images already within the cap are not resized."""
+        cb = COCOEvalCallback(max_eval_orig_size=640)
+        targets = [
+            {
+                "boxes": torch.tensor([[0.5, 0.5, 1.0, 1.0]]),
+                "labels": torch.tensor([0]),
+                "orig_size": torch.tensor([480, 640]),
+                "masks": torch.ones(1, 480, 640, dtype=torch.bool),
+            }
+        ]
+        out = cb._convert_targets(targets)
+        assert out[0]["masks"].shape[-2:] == (480, 640)