|
|
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 |
|
|
) |
