Repository: SEA-AI/yolov5
Pull request #18: CV-336: GPU Preprocessing inside Model's Graph
Author: login: GilSeamas
Created: 2025-04-17T13:44:32Z | Merged: 2025-04-18T09:50:25Z

## What?
 
Add possibility to include preprocessing inside model's graph
 
## Why?
 
preprocessing in GPU is considerably faster (check values below)
 
## How?
 
by using torchivision tranforms and applying them only if necessery

## Backout plan

No backout plan should be necessery. For no preprocssing just be sure to export model with --imgsz equal to --infsz

## Testing

```
Average time - no preprocessing:  0.009861427148183187
Average time - gpu preprocessing:  0.010737588008244833
```
(CPU preprocessing takes more or less the same as the models forward pass)


![output](https://github.com/user-attachments/assets/848db753-40b5-46c9-b174-5c1b4ae3b5d7)

PS: did not include rescaling boxes back to original image input size inside graph but we can do it.





Diff:
diff --git a/export_ahoy.py b/export_ahoy.py
index 9c1f4a0c9fd7..dd96880e60f9 100644
--- a/export_ahoy.py
+++ b/export_ahoy.py
@@ -14,11 +14,16 @@
         --det-weights yolov5n.pt \
         --hor-weights yolov5h.pt \
         --imgsz 640 \
+        --infsz 320
         --batch-size 2 \
         --fuse \
         --half \
         --fname ahoy.onnx
 
+# flatten this
+    python export_ahoy.py --det-weights yolov5n.pt --hor-weights yolov5h.pt --imgsz 640 --infsz 320 --batch-size 2 --fuse --half --fname ahoy.onnx
+
+
     # Using W&B artifacts:
     python export_ahoy.py \
         --det-weights YOLOv5n-IR:latest \
@@ -78,7 +83,7 @@ def get_weights_path(weights_path: str) -> str:
         raise e
     
 
-def _transform_imgsz(imgsz: int | List[int] | Tuple[int,int]) -> Tuple[int,int]:
+def _transform_sz(imgsz: int | List[int] | Tuple[int,int]) -> Tuple[int,int]:
     """
     Convert size specifications to (height, width) tuple format.
     
@@ -104,6 +109,7 @@ def main(
     det_weights: str,
     hor_weights: str,
     imgsz: int | Tuple[int, int],
+    infsz: int | Tuple[int, int],
     batch_size: int,
     half: bool,
     fuse: bool,
@@ -112,8 +118,8 @@ def main(
 ):
     """Export the model to TensorRT engine."""
     # Transform image size to (height, width) format
-    imgsz = _transform_imgsz(imgsz)
-
+    imgsz = _transform_sz(imgsz)
+    infsz = _transform_sz(infsz)
     det_weights = get_weights_path(det_weights)
     hor_weights = get_weights_path(hor_weights)
 
@@ -122,6 +128,8 @@ def main(
         hor_det_weights=hor_weights,
         fp16=half,
         fuse=fuse,
+        imgsz=imgsz,
+        infsz=infsz,
     )
 
     if not fname:
@@ -184,7 +192,14 @@ def _parse_args():
         nargs="+", 
         type=int, 
         default=[640, 640], 
-        help="image (h, w)"
+        help="image input shape (h, w)"
+    )
+    parser.add_argument(
+        "--infsz",
+        nargs="+", 
+        type=int, 
+        default=[640, 640], 
+        help="image shape during inference (h, w)"
     )
     parser.add_argument(
         "-bs",
diff --git a/models/custom.py b/models/custom.py
index 6bcf54fd6040..f1b7a40fddc3 100644
--- a/models/custom.py
+++ b/models/custom.py
@@ -4,6 +4,7 @@
 import numpy as np
 import torch
 from torch import nn
+from torchvision import transforms
 
 from models.common import Classify, DetectMultiBackend
 from models.experimental import attempt_load
@@ -255,6 +256,7 @@ def __init__(
         self.save = model.save
 
 
+
 class AHOY(nn.Module):
     """A H-orizon O-bject detection Y-OLOv5."""
 
@@ -266,6 +268,8 @@ def __init__(
         device: Union[str, torch.device] = None,  # automatically select device
         fp16: bool = False,
         fuse: bool = True,  # fuse conv and bn layers
+        imgsz: list = [640,640],
+        infsz: list = [640,640],
         inplace: bool = True,  # inplace modification of models
     ):
         super().__init__()
@@ -275,10 +279,15 @@ def __init__(
         self.fp16 = fp16
         self.stride = self.obj_det.stride
         self.names = self.obj_det.names
+        self.imgsz = imgsz
+        self.infsz = infsz
 
         # keep track of hooks
         self.hooks = {}
 
+        # Scaling in Padding Preprocessing
+        self.transform = self.get_transform(imgsz, infsz)
+
     def forward(self, x, profile=False, visualize=False):
         """Forward pass through models."""
         objects = self.obj_det(x, profile, visualize)
@@ -308,20 +317,64 @@ def remove_hooks(self):
             hook.remove()
         self.hooks.clear()
 
+    @staticmethod
+    def get_transform(imgsz, infsz):
+        """Get the transformation to be applied to the image. Padding and/or resize, if needed."""
+        if imgsz == infsz:
+            return None
+
+        pad_left, pad_right, pad_top, pad_bottom = AHOY.get_padding_for_aspect_ratio(imgsz, infsz)
+        ratio = max(imgsz[0] / infsz[0], imgsz[1] / infsz[1])
+        transform = transforms.Compose([])
+        if ratio != 1:
+            transform.transforms.extend(
+                [transforms.Resize((infsz[0]-pad_top-pad_bottom, infsz[1]-pad_left-pad_right), interpolation=transforms.InterpolationMode.NEAREST, antialias=False)]
+            )
+        if pad_left != 0 or pad_right != 0 or pad_top != 0 or pad_bottom != 0:
+            transform.transforms.extend(
+                [transforms.Pad(padding=(pad_left, pad_top, pad_right, pad_bottom), fill = 0, padding_mode="constant")]
+            )
+        return transform
+
+    @staticmethod
+    def get_padding_for_aspect_ratio(imgsz, infsz):
+        """
+        Calculates padding (left, right, top, bottom) needed after resizing
+        while preserving aspect ratio to fit exactly into (h_out, w_out).
+        """
+        h_in, w_in = imgsz
+        h_out, w_out = infsz
+
+        scale = min(h_out / h_in, w_out / w_in)
+        new_h = int(h_in * scale)
+        new_w = int(w_in * scale)
+
+        pad_h = h_out - new_h
+        pad_w = w_out - new_w
+
+        pad_top = pad_h // 2
+        pad_bottom = pad_h - pad_top
+        pad_left = pad_w // 2
+        pad_right = pad_w - pad_left
+
+        return pad_left, pad_right, pad_top, pad_bottom
+
     @staticmethod
     def _preprocessing_hook(module, inputs):
         """Add preprocessing operations to be part of the model."""
 
         def _preprocess(x):
-            if not isinstance(x, torch.Tensor):
+            if len(x.shape) <1 :
                 return x
+            if module.transform is not None:
+                x = x.float()
+                x = module.transform(x)
             x = x.half() if module.fp16 else x.float()
             x = x / 255.0  # 0-255 to 0.0-1.0
             return x
 
         return tuple(_preprocess(inp) for inp in inputs)
 
-
     @staticmethod
     def _postprocessing_hook(module, inputs, outputs):
         """Convert outputs to float (if needed) and apply softmax to logits."""


Reviews:
user: login: Zwicii
state: DISMISSED
author_association: NONE
submitted_at: 2025-04-18T09:21:11Z
user: login: GilSeamas
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-18T09:44:28Z
user: login: Zwicii
state: APPROVED
author_association: NONE
submitted_at: 2025-04-18T09:49:50Z

Review comments:
diff_hunk: @@ -112,8 +118,9 @@ def main(
 ):
     """Export the model to TensorRT engine."""
     # Transform image size to (height, width) format
-    imgsz = _transform_imgsz(imgsz)
-
+    imgsz = _transform_sz(imgsz)
+    infsz = _transform_sz(infsz)
+    print("main", imgsz, infsz)
path: export_ahoy.py
user: login: Zwicii
body: 👀
created_at: 2025-04-18T09:09:05Z
author_association: NONE
diff_hunk: @@ -308,20 +317,64 @@ def remove_hooks(self):
             hook.remove()
         self.hooks.clear()
 
+    @staticmethod
+    def get_transform(imgsz, infsz):
+        """Get the transformation to be applied to the image. Padding and/or resize, if needed."""
+        if imgsz != infsz:
+            pad_left, pad_right, pad_top, pad_bottom = AHOY.get_padding_for_aspect_ratio(imgsz, infsz)
+            ratio = max(imgsz[0] / infsz[0], imgsz[1] / infsz[1])
+            transform = transforms.Compose([])
+            if ratio != 1:
+                transform.transforms.extend(
+                    [transforms.Resize((infsz[0]-pad_top-pad_bottom, infsz[1]-pad_left-pad_right), interpolation=transforms.InterpolationMode.NEAREST, antialias=False)]
+                )
+            if pad_left != 0 or pad_right != 0 or pad_top != 0 or pad_bottom != 0:
+                transform.transforms.extend(
+                    [transforms.Pad(padding=(pad_left, pad_top, pad_right, pad_bottom), fill = 0, padding_mode="constant")]
+                )
+            return transform
+        else:
+            return None
path: models/custom.py
user: login: Zwicii
body: I would do an early return here because it simplifies the code by reducing nesting and improves readability

```python
# Early return if sizes match
if imgsz == infsz:
    return None
```
created_at: 2025-04-18T09:19:42Z
author_association: NONE
diff_hunk: @@ -308,20 +317,64 @@ def remove_hooks(self):
             hook.remove()
         self.hooks.clear()
 
+    @staticmethod
+    def get_transform(imgsz, infsz):
+        """Get the transformation to be applied to the image. Padding and/or resize, if needed."""
+        if imgsz != infsz:
+            pad_left, pad_right, pad_top, pad_bottom = AHOY.get_padding_for_aspect_ratio(imgsz, infsz)
+            ratio = max(imgsz[0] / infsz[0], imgsz[1] / infsz[1])
+            transform = transforms.Compose([])
+            if ratio != 1:
+                transform.transforms.extend(
+                    [transforms.Resize((infsz[0]-pad_top-pad_bottom, infsz[1]-pad_left-pad_right), interpolation=transforms.InterpolationMode.NEAREST, antialias=False)]
+                )
+            if pad_left != 0 or pad_right != 0 or pad_top != 0 or pad_bottom != 0:
+                transform.transforms.extend(
+                    [transforms.Pad(padding=(pad_left, pad_top, pad_right, pad_bottom), fill = 0, padding_mode="constant")]
+                )
+            return transform
+        else:
+            return None
path: models/custom.py
user: login: GilSeamas
body: yes, looks better that way, thx!

created_at: 2025-04-18T09:44:28Z
author_association: NONE

PR comments:
user: login: Zwicii
created_at: 2025-04-18T09:22:53Z
body: Very cool that you got it working so fast :) 
(just one remaining debug print statement and one suggestion - but otherwise LGTM)
author_association: NONE
user: login: GilSeamas
created_at: 2025-04-18T09:46:26Z
body: > Very cool that you got it working so fast :) (just one remaining debug print statement and one suggestion - but otherwise LGTM)

ahah, thanks you! :D
author_association: NONE