92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
# %% Импорт библиотек
|
|
|
|
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)})
|