fix: Fixed dataloader issues

configs/train.yaml

@@ -8,9 +8,9 @@ defaults: # these are the sub-config files for specific details, noee for the mo
 
 # Training Configuration
 training_config:
-  batch_size: 8
-  num_workers: 16
-  epochs: 200
+  batch_size: 4
+  num_workers: 8
+  epochs: 1000
   accum_steps: 1 # Simulate batch size * accum_steps
   optimizer_momentum: 0.9
   optimizer_weight_decay: 5e-4

src/datasets/aeroscapes.py

@@ -25,6 +25,7 @@
     RandomScale,
     RandomTranslate,
     RandomVerticalFlip,
+    ResizeIfLarger,
 )
 
 # Same knobs/semantics as src/datasets/uavid.py's DEFAULT_AUGMENTATION —
@@ -174,6 +175,12 @@ def __init__(
         self.trans_train = (
             Compose(
                 [
+                    # No-op at AeroScapes' native 1280x720 with the default
+                    # cropsize, but bounds the working resolution before the
+                    # expensive geometric ops below regardless of cropsize —
+                    # see ResizeIfLarger's docstring (uavid.py/vdd.py hit
+                    # this for real, at their much larger native resolutions).
+                    ResizeIfLarger(max_size=2 * max(self.cropsize)),
                     RandomHorizontalFlip(p=float(self.aug["fliplr"])),
                     RandomVerticalFlip(p=float(self.aug["flipud"])),
                     RandomTranslate(

src/datasets/transform.py

@@ -26,6 +26,42 @@ def __call__(self, im_lb):
         return im_lb
 
 
+class ResizeIfLarger(object):
+    """Downsize (never upsize) so the image's longer side is at most
+    ``max_size``, preserving aspect ratio. A no-op when the image already fits.
+
+    Purpose: the geometric transforms that typically run before ``RandomCrop``
+    in this pipeline (``RandomTranslate``, ``RandomRotate``, ``RandomScale``)
+    are CPU-bound PIL/numpy operations whose cost scales with the *native*
+    image resolution, not the eventual training crop size. For datasets with
+    very large native images (e.g. VDD's 4000x3000, ~11x a 1024x1024 crop),
+    this makes CPU-side augmentation the training bottleneck — the GPU sits
+    idle waiting for augmented batches (low, bursty utilization). Placing this
+    as the *first* step in the pipeline bounds per-sample CPU cost regardless
+    of native resolution, with no effect on the final crop (still exactly
+    ``cropsize``) or on crop-position randomness, as long as ``max_size`` is
+    kept comfortably larger than ``cropsize``.
+    """
+
+    def __init__(self, max_size, interp_image=Image.BILINEAR, interp_label=Image.NEAREST):
+        self.max_size = max_size
+        self.interp_image = interp_image
+        self.interp_label = interp_label
+
+    def __call__(self, im_lb):
+        im = im_lb["im"]
+        lb = im_lb["lb"]
+        w, h = im.size
+        longest = max(w, h)
+        if longest <= self.max_size:
+            return im_lb
+        scale = self.max_size / longest
+        new_w, new_h = max(1, round(w * scale)), max(1, round(h * scale))
+        im_lb["im"] = im.resize((new_w, new_h), self.interp_image)
+        im_lb["lb"] = lb.resize((new_w, new_h), self.interp_label)
+        return im_lb
+
+
 class RandomScale(object):
     """Random resize, either from a discrete list of scale factors
     (default) or a continuous range — pass ``continuous=True`` with

src/datasets/uavid.py

@@ -25,6 +25,7 @@
     RandomScale,
     RandomTranslate,
     RandomVerticalFlip,
+    ResizeIfLarger,
 )
 
 # Mirrors configs/yolo/uavid_train.yaml / configs/train_yolo.yaml's
@@ -191,6 +192,12 @@ def __init__(
         self.trans_train = (
             Compose(
                 [
+                    # UAVid's native images run up to 4096x2160 (~8.85MP) —
+                    # cap the working resolution before the expensive
+                    # geometric ops below, or CPU-side augmentation becomes
+                    # the training bottleneck (GPU starves waiting on the
+                    # DataLoader). See ResizeIfLarger's docstring.
+                    ResizeIfLarger(max_size=2 * max(self.cropsize)),
                     RandomHorizontalFlip(p=float(self.aug["fliplr"])),
                     RandomVerticalFlip(p=float(self.aug["flipud"])),
                     RandomTranslate(

src/datasets/vdd.py

@@ -25,6 +25,7 @@
     RandomScale,
     RandomTranslate,
     RandomVerticalFlip,
+    ResizeIfLarger,
 )
 
 # Same knobs/semantics as src/datasets/uavid.py's DEFAULT_AUGMENTATION — VDD
@@ -174,6 +175,13 @@ def __init__(
         self.trans_train = (
             Compose(
                 [
+                    # VDD's native images are 4000x3000 (~11x a 1024x1024
+                    # crop) — cap the working resolution before the
+                    # expensive geometric ops below, or CPU-side
+                    # augmentation becomes the training bottleneck (GPU
+                    # starves waiting on the DataLoader). See
+                    # ResizeIfLarger's docstring.
+                    ResizeIfLarger(max_size=2 * max(self.cropsize)),
                     RandomHorizontalFlip(p=float(self.aug["fliplr"])),
                     RandomVerticalFlip(p=float(self.aug["flipud"])),
                     RandomTranslate(

tests/unit/test_transforms.py

@@ -17,6 +17,7 @@
     RandomScale,
     RandomTranslate,
     RandomVerticalFlip,
+    ResizeIfLarger,
 )
 
 
@@ -384,3 +385,49 @@ def test_large_gains_change_pixel_values(self):
         after = np.array(result["im"])
 
         assert not np.array_equal(before, after)
+
+
+class TestResizeIfLarger:
+    def test_noop_when_already_within_max_size(self):
+        img = Image.new("RGB", (256, 256), color=(1, 2, 3))
+        lb = Image.new("L", (256, 256), color=1)
+        transform = ResizeIfLarger(max_size=1024)
+        result = transform({"im": img, "lb": lb})
+        assert result["im"].size == (256, 256)
+        assert result["im"] is img  # untouched, not even a copy
+
+    def test_noop_when_exactly_at_max_size(self):
+        img = Image.new("RGB", (1024, 512), color=(1, 2, 3))
+        lb = Image.new("L", (1024, 512), color=1)
+        transform = ResizeIfLarger(max_size=1024)
+        result = transform({"im": img, "lb": lb})
+        assert result["im"].size == (1024, 512)
+
+    def test_downsizes_larger_image_preserving_aspect_ratio(self):
+        img = Image.new("RGB", (4000, 3000), color=(1, 2, 3))
+        lb = Image.new("L", (4000, 3000), color=1)
+        transform = ResizeIfLarger(max_size=2000)
+        result = transform({"im": img, "lb": lb})
+        assert result["im"].size == (2000, 1500)
+        assert result["lb"].size == (2000, 1500)
+
+    def test_never_upsizes(self):
+        img = Image.new("RGB", (100, 80), color=(1, 2, 3))
+        lb = Image.new("L", (100, 80), color=1)
+        transform = ResizeIfLarger(max_size=2000)
+        result = transform({"im": img, "lb": lb})
+        assert result["im"].size == (100, 80)
+
+    def test_caps_on_the_longer_dimension_for_portrait_images(self):
+        img = Image.new("RGB", (1000, 4000), color=(1, 2, 3))
+        lb = Image.new("L", (1000, 4000), color=1)
+        transform = ResizeIfLarger(max_size=2000)
+        result = transform({"im": img, "lb": lb})
+        assert result["im"].size == (500, 2000)
+
+    def test_label_mode_preserved_after_downsize(self):
+        img = Image.new("RGB", (4000, 3000), color=(1, 2, 3))
+        lb = Image.new("L", (4000, 3000), color=1)
+        transform = ResizeIfLarger(max_size=2000)
+        result = transform({"im": img, "lb": lb})
+        assert result["lb"].mode == "L"