File size: 3,230 Bytes
8bf25c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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
) |