From 5869d0248fa958acd3447e6bffa8761b91e8e921 Mon Sep 17 00:00:00 2001
From: Jordan Gong <jordan.gong@protonmail.com>
Date: Wed, 16 Mar 2022 17:49:51 +0800
Subject: Regular refactoring

---
 supervised/datautils.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)
 create mode 100644 supervised/datautils.py

(limited to 'supervised/datautils.py')

diff --git a/supervised/datautils.py b/supervised/datautils.py
new file mode 100644
index 0000000..196fca7
--- /dev/null
+++ b/supervised/datautils.py
@@ -0,0 +1,13 @@
+from torchvision.transforms import transforms
+
+
+def color_distortion(s=1.0):
+    # s is the strength of color distortion.
+    color_jitter = transforms.ColorJitter(0.8 * s, 0.8 * s, 0.8 * s, 0.2 * s)
+    rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)
+    rnd_gray = transforms.RandomGrayscale(p=0.2)
+    color_distort = transforms.Compose([
+        rnd_color_jitter,
+        rnd_gray
+    ])
+    return color_distort
-- 
cgit v1.2.3