first commit

MyDataExamplePL.py

@@ -0,0 +1,91 @@
+# %% Импорт библиотек
+
+import kagglehub
+import splitfolders
+import pytorch_lightning as pl
+
+from torch.utils.data import DataLoader
+from torchvision.datasets import ImageFolder
+from torchvision.transforms import v2 as transforms
+
+from pathlib import Path
+from typing import Optional, Callable, Any
+
+# %% Скачивание набора данных
+
+path = kagglehub.dataset_download("jehanbhathena/weather-dataset",
+                                  output_dir="./Datasets",
+                                  force_download=False)
+
+# %% Разделение набора данных
+
+splitfolders.ratio("./Datasets/weather-dataset",
+                   output="./Datasets/weather-dataset-split",
+                   ratio=(.8, .1, .1),
+                   group_prefix=None,
+                   move=False)
+
+# %% Пример модуля для загрузки данных
+
+TransformType = Optional[Callable[..., Any]]
+
+class MyImageDataset(pl.LightningDataModule):
+    def __init__(self,
+                 batch_size: int = 1,
+                 train_transform: TransformType = None,
+                 val_transform: TransformType = None,
+                 img_root: str = ""
+                 ):
+        
+        super().__init__()
+        
+        self.batch_size = batch_size
+        self.train_transform = train_transform
+        self.val_transform = val_transform
+        self.img_root = img_root
+        
+    def setup(self, stage=None):
+        self.train_set = ImageFolder(self.img_root / Path("train"), 
+                                     transform=self.train_transform)
+        self.val_set = ImageFolder(self.img_root / Path("val"), 
+                                     transform=self.val_transform)
+        self.test_set = ImageFolder(self.img_root / Path("test"), 
+                                     transform=self.val_transform)
+        
+    def train_dataloader(self):
+        return DataLoader(self.train_set, batch_size=self.batch_size,
+                          shuffle=True, num_workers=2)
+    
+    def val_dataloader(self):
+        return DataLoader(self.val_set, batch_size=self.batch_size,
+                          shuffle=False, num_workers=2)
+    
+    def test_dataloader(self):
+        return DataLoader(self.test_set, batch_size=self.batch_size,
+                          shuffle=False, num_workers=2)
+
+# %% Создание трансформаций
+
+transform_train = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.RandomHorizontalFlip(p=.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])     
+])
+
+transform_test = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])     
+])
+
+# %% Создание экземляра класса и вывод объектов класса
+
+dm = MyImageDataset(train_transform=transform_train,
+                    val_transform=transform_test,
+                    img_root="./Datasets/weather-dataset-split")
+dm.setup()
+
+print({k:v for k, v in enumerate(dm.test_set.classes)})