From aa5292e8d62d024236c3b234b7b8352efe6e96e2 Mon Sep 17 00:00:00 2001
From: Evening <dev_evening@hotmail.com>
Date: Tue, 11 Jun 2024 16:08:11 +0800
Subject: [PATCH] Fix issue that slows down preprocessing

---
 tests/model_tests/utils.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/model_tests/utils.py b/tests/model_tests/utils.py
index a596964..f9d7917 100644
--- a/tests/model_tests/utils.py
+++ b/tests/model_tests/utils.py
@@ -14,6 +14,7 @@
     RandomRotation,
     RandomApply,
     Resize,
+    RandomResizedCrop,
 )
 from torchvision.transforms.v2 import RandomHorizontalFlip
 
@@ -53,11 +54,12 @@ def rand_weak_aug(size: int):
         [
             ToImage(),
             ToDtype(torch.float32, scale=True),
+            RandomResizedCrop(
+                size, scale=(0.08, 1.0), ratio=(0.95, 1.05), antialias=True
+            ),
             RandomHorizontalFlip(),
             RandomVerticalFlip(),
             RandomApply([RandomRotation((90, 90))], p=0.5),
-            Resize(size, antialias=True),
-            CenterCrop(size),
         ]
     )
 
@@ -78,10 +80,11 @@ def rand_strong_aug(size: int):
         [
             ToImage(),
             ToDtype(torch.float32, scale=True),
+            RandomResizedCrop(
+                size, scale=(0.08, 1.0), ratio=(0.9, 1.1), antialias=True
+            ),
             RandomHorizontalFlip(),
             RandomVerticalFlip(),
             RandomApply([RandomRotation((90, 90))], p=0.5),
-            Resize(size, antialias=True),
-            RandomCrop(size, pad_if_needed=False),  # Strong
         ]
     )