From 3e9c8a37e7b9086f4bc4d500ae2a5716573ae0ec Mon Sep 17 00:00:00 2001
From: Henry Ruhs <info@henryruhs.com>
Date: Mon, 9 Jun 2025 07:49:00 +0200
Subject: [PATCH] Some polishing on augmentation (#82)

---
 hyperswap/src/dataset.py | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/hyperswap/src/dataset.py b/hyperswap/src/dataset.py
index 3d9bdb4..3243717 100644
--- a/hyperswap/src/dataset.py
+++ b/hyperswap/src/dataset.py
@@ -146,15 +146,29 @@ class AugmentTransform:
 	def compose_transforms() -> albumentations.Compose:
 		return albumentations.Compose(
 		[
-			albumentations.HorizontalFlip(),
+			albumentations.HorizontalFlip(p = 0.5),
 			albumentations.OneOf(
 			[
-				albumentations.MotionBlur(p = 0.1),
-				albumentations.ZoomBlur(max_factor = (1.0, 1.1), p = 0.1)
+				albumentations.MotionBlur(),
+				albumentations.ZoomBlur(max_factor = (1.0, 1.2))
+			], p = 0.1),
+			albumentations.OneOf(
+			[
+				albumentations.RandomGamma(),
+				albumentations.RandomBrightnessContrast(),
+				albumentations.Illumination()
 			], p = 0.2),
-			albumentations.RandomBrightnessContrast(p = 0.7),
-			albumentations.ColorJitter(p = 0.2),
-			albumentations.RGBShift(p = 0.7),
-			albumentations.Illumination(p = 0.2),
-			albumentations.Affine(translate_percent = (-0.03, 0.03), scale = (0.98, 1.02), rotate = (-2, 2), border_mode = 1, p = 0.3)
+			albumentations.OneOf(
+			[
+				albumentations.ColorJitter(),
+				albumentations.RGBShift(),
+				albumentations.HueSaturationValue()
+			], p = 0.2),
+			albumentations.Affine(
+				translate_percent = (-0.05, 0.05),
+				scale = (0.95, 1.05),
+				rotate = (-2, 2),
+				border_mode = 1,
+				p = 0.2
+			)
 		])