| from torch.utils.data import Dataset, DataLoader |
| from typing import * |
| from dataclasses import dataclass, field |
| from PIL import Image |
| from utils import parse_structure |
|
|
| import os |
| import lightning.pytorch as pl |
| import numpy as np |
| import torch |
|
|
|
|
| class BaseDataset(Dataset): |
| def __init__(self, root_dir: str, image_size: Tuple[int, int]) -> None: |
| self.root_dir = root_dir |
| self.image_size = image_size |
| self.classes = {folder: idx for idx, folder in enumerate(os.listdir(root_dir))} |
| self.image_paths = [] |
| self.labels = [] |
|
|
| for class_name, class_idx in self.classes.items(): |
| class_dir = os.path.join(root_dir, class_name) |
| for img_name in os.listdir(class_dir): |
| img_path = os.path.join(class_dir, img_name) |
| self.image_paths.append(img_path) |
| self.labels.append(class_idx) |
|
|
| def __len__(self) -> int: |
| return len(self.image_paths) |
|
|
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]: |
| img_path = self.image_paths[idx] |
| label = self.labels[idx] |
| image = Image.open(img_path).convert("RGB") |
| image = image.resize(self.image_size) |
| image = np.array(image) |
| image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 |
| return image, label |
|
|
|
|
| @dataclass |
| class BaseDatasetConfig: |
| data_source: str = '' |
| train_path:str = '' |
| valid_path:str = '' |
| test_path:str = '' |
| batch_size:int = 32 |
| shuffle:bool = True |
| num_workers:int = 24 |
| image_size:Tuple[int, int] = (224, 224) |
|
|
| class BaseDataModule(pl.LightningDataModule): |
| cfg: BaseDatasetConfig |
|
|
| def __init__(self, cfg: BaseDatasetConfig) -> None: |
| super().__init__() |
| self.cfg:BaseDatasetConfig = parse_structure(BaseDatasetConfig, cfg) |
| self.train_path = cfg.train_path |
| self.valid_path = cfg.valid_path |
| self.test_path = cfg.test_path |
| self.img_size = cfg.image_size |
|
|
| def setup(self, stage=None) -> None: |
| if stage in [None, "fit"]: |
| self.train_dataset = BaseDataset(self.train_path, self.img_size) |
| if stage in [None, "fit", "validate"]: |
| self.val_dataset = BaseDataset(self.valid_path, self.img_size) |
| if stage in [None, "test", "predict"]: |
| self.test_dataset = BaseDataset(self.test_path, self.img_size) |
|
|
| def general_loader(self, dataset, batch_size) -> DataLoader: |
| return DataLoader( |
| dataset, |
| num_workers=self.cfg.num_workers, |
| batch_size=batch_size |
| ) |
|
|
| def train_dataloader(self) -> DataLoader: |
| return DataLoader( |
| self.train_dataset, |
| num_workers=self.cfg.num_workers, |
| batch_size=self.cfg.batch_size, |
| shuffle=self.cfg.shuffle |
| ) |
|
|
| def val_dataloader(self) -> DataLoader: |
| return DataLoader( |
| self.val_dataset, |
| num_workers=self.cfg.num_workers, |
| batch_size=self.cfg.batch_size |
| ) |
|
|
| def test_dataloader(self) -> DataLoader: |
| return DataLoader( |
| self.test_dataset, |
| num_workers=self.cfg.num_workers, |
| batch_size=self.cfg.batch_size |
| ) |
