feat: add class ImageTransforms in transforms module #1

src/lightning_classify/transforms/__init__.py

@@ -0,0 +1 @@
+from .transforms import *

src/lightning_classify/transforms/transforms.py

@@ -0,0 +1,33 @@
+from PIL import Image
+
+import torch
+from torchvision import transforms
+
+
+class ImageTransform:
+    def __init__(self, is_train: bool) -> None:
+        if is_train:
+            self.transform = transforms.Compose([
+                transforms.Resize(256),
+                transforms.RandomResizedCrop(224),
+                transforms.RandomHorizontalFlip(),
+                transforms.RandomRotation(10),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406],
+                    std=[0.229, 0.224, 0.225],
+                ),
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.Resize(256),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406], 
+                    std=[0.229, 0.224, 0.225]
+                ),
+            ])
+
+    def __call__(self, img: Image.Image) -> torch.Tensor:
+        return self.transform(img)