init

data/__init__.py

@@ -0,0 +1,3 @@
+from .base import BaseDataModule
+from .camelyon import CamelyonDataModule
+from .breakhis import BreakhisDataModule

data/base.py

@@ -0,0 +1,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
+        )

data/breakhis.py

@@ -0,0 +1,146 @@
+from torch.utils.data import Dataset, DataLoader
+from typing import *
+from dataclasses import dataclass, field
+from PIL import Image
+from utils import parse_structure
+from glob import glob
+from random import shuffle
+from torchvision.transforms import v2
+
+import os
+import lightning.pytorch as pl
+import numpy as np
+import torch
+import random
+
+
+class BreakhisDataset(Dataset):
+    def __init__(self, root_dir: str, image_size: Tuple[int, int], subset: str, aug: dict = None) -> None:
+        self.root_dir = root_dir
+        self.image_size = image_size
+        self.classes = {
+            'benign' : 0,
+            'malignant' : 1
+        }
+        self.ratio = [0.8, 0.1]
+        self.subset = subset
+        self.aug = aug
+
+        self.benign_subclasses = ['adenosis', 'fibroadenoma', 'phyllodes_tumor', 'tubular_adenoma']
+        self.malignant_subclasses = ['ductal_carcinoma', 'lobular_carcinoma', 'mucinous_carcinoma', 'papillary_carcinoma']
+        self.cls2sublst = {
+            'benign' : self.benign_subclasses,
+            'malignant' : self.malignant_subclasses
+        }
+        self.factors = ['100X', '200X', '400X', '40X']
+
+        self.sample_paths = []
+        self.sample_labels = []
+
+        random.seed(42)
+        for cate in ['benign', 'malignant']:
+            for subcls in self.cls2sublst[cate]:
+                for factor in self.factors:
+                    lst = glob(os.path.join(self.root_dir, f'{cate}/*/{subcls}/*/{factor}/*.png'))
+                    random.shuffle(lst)
+
+                    sublst = self.get_subset(lst)
+                    self.sample_paths += sublst
+                    self.sample_labels += [self.classes[cate]] * len(sublst)
+        
+        if self.aug is not None:
+            self.transforms = [v2.Resize(self.image_size, antialias=True)] + \
+                [getattr(v2, x)(**self.aug[x]) for x in self.aug] + \
+                [
+                    v2.ToDtype(torch.float32, scale=True),
+                    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ]
+        else:
+            self.transforms = [
+                v2.Resize(self.image_size, antialias=True),
+                v2.ToDtype(torch.float32, scale=True),
+                v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+            ]
+        
+        self.transform = v2.Compose(self.transforms)
+    
+    def get_subset(self, x: list):
+        if self.subset == 'train':
+            return x[ : int(self.ratio[0] * len(x))]
+        elif self.subset == 'valid':
+            return x[int(self.ratio[0] * len(x)) : int((self.ratio[0] + self.ratio[1]) * len(x))]
+        elif self.subset == 'test':
+            return x[int((self.ratio[0] + self.ratio[1]) * len(x)) : ]
+        else:
+            return ValueError('Unknown subset')
+
+    def __len__(self) -> int:
+        return len(self.sample_paths)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
+        img_path = self.sample_paths[idx]
+        label = self.sample_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)
+        image = self.transform(image)
+
+        return image, label
+
+
+@dataclass
+class BaseDatasetConfig:
+    data_source: str = ''
+    batch_size:int = 32
+    shuffle:bool = True
+    num_workers:int = 24
+    image_size:Tuple[int, int] = (224, 224)
+    aug: dict = field(default_factory=dict)
+
+class BreakhisDataModule(pl.LightningDataModule):
+    cfg: BaseDatasetConfig
+
+    def __init__(self, cfg: BaseDatasetConfig) -> None:
+        super().__init__()
+        self.cfg:BaseDatasetConfig = parse_structure(BaseDatasetConfig, cfg)
+        self.data_source = self.cfg.data_source
+        self.img_size = self.cfg.image_size
+        self.aug = self.cfg.aug
+
+    def setup(self, stage=None) -> None:
+        if stage in [None, "fit"]:
+            self.train_dataset = BreakhisDataset(self.data_source, self.img_size, 'train', self.aug)
+        if stage in [None, "fit", "validate"]:
+            self.val_dataset = BreakhisDataset(self.data_source, self.img_size, 'valid', self.aug)
+        if stage in [None, "test", "predict"]:
+            self.test_dataset = BreakhisDataset(self.data_source, self.img_size, 'test', self.aug)
+
+    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
+        )

data/camelyon.py

@@ -0,0 +1,73 @@
+from wilds.datasets.camelyon17_dataset import Camelyon17Dataset
+from .base import BaseDatasetConfig, BaseDataModule
+
+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 numpy as np
+import torch
+import albumentations as A
+
+class CamelyonDataset(Dataset):
+    def __init__(self, root_dir: str, subset: str, image_size: Tuple[int, int]) -> None:
+        self.root_dir   = root_dir
+        self.dataset    = Camelyon17Dataset(root_dir=root_dir, download=True).get_subset(subset)
+        self.transform  = {
+            "train" : A.Compose([
+                A.HorizontalFlip(),
+                A.Affine(scale=(-0.2, 0.2),
+                    rotate=(-10, 10),
+                    # shear=(-5, 5), 
+                    keep_ratio=True, 
+                    p=0.5),
+                A.OneOf([
+                    A.MotionBlur(p=0.2),
+                    A.MedianBlur(blur_limit=3, p=0.1),
+                    A.Blur(blur_limit=3, p=0.1),
+                ], p=0.5),
+                A.OneOf([
+                    A.CLAHE(clip_limit=2),
+                    A.RandomBrightnessContrast(),
+                ], p=0.5),
+                A.HueSaturationValue(p=0.25),
+                A.Resize(image_size[0], image_size[1])
+            ], p=1.0),
+            "val" : A.Compose([
+                A.Resize(image_size[0], image_size[1])
+            ], p=1.0),
+            "test" : A.Compose([
+                A.Resize(image_size[0], image_size[1])
+            ], p=1.0)
+        }[subset]
+        self.image_size = image_size
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
+        (image, label, _) = self.dataset.__getitem__(idx)
+        # image = image.resize(self.image_size)
+        image = np.array(image)
+        image = self.transform(image=image)["image"]
+        image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
+        return image, label
+
+class CamelyonDataModule(BaseDataModule):
+    cfg: BaseDatasetConfig
+
+    def __init__(self, cfg: BaseDatasetConfig) -> None:
+        super().__init__(cfg)
+        self.cfg:DatasetConfig = parse_structure(BaseDatasetConfig, cfg)
+        self.img_size = cfg.image_size
+
+    def setup(self, stage=None) -> None:
+        if stage in [None, "fit"]:
+            self.train_dataset = CamelyonDataset(self.cfg.data_source, "train", self.img_size)
+        if stage in [None, "fit", "validate"]:
+            self.val_dataset = CamelyonDataset(self.cfg.data_source, "val", self.img_size)
+        if stage in [None, "test", "predict"]:
+            self.test_dataset = CamelyonDataset(self.cfg.data_source, "test", self.img_size)

metrics/__init__.py

@@ -0,0 +1 @@
+from .base import BaseMetrics, BaseMetricsConfig

metrics/base.py

@@ -0,0 +1,33 @@
+from torchmetrics import classification
+from dataclasses import dataclass, field
+from typing import Any, Dict, Mapping
+from utils import parse_structure
+from torch import Tensor, nn
+
+import lightning.pytorch as pl
+import torch
+
+@dataclass
+class BaseMetricsConfig:
+    metrics_names:list = field(default_factory=list)
+    metrics_short_names:list = field(default_factory=list)
+
+class BaseMetrics(pl.LightningModule):
+    def __init__(self, cfg: Dict, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.cfg: BaseMetricsConfig = parse_structure(BaseMetricsConfig, cfg)
+        self.metrics_names = self.cfg.metrics_names
+        self.metrics_short_names = self.cfg.metrics_short_names
+
+        self.metrics = nn.ModuleDict()
+        for name, short_name in zip(self.metrics_names, self.metrics_short_names):
+            obj = getattr(classification, name)
+            metric = obj()
+            self.metrics[short_name] = metric
+
+        print(f"[INFO]: Metrics: {self.metrics}")
+    
+    def __call__(self, pred: Tensor, target: Tensor, prefix:str) -> Dict[str, float]:
+        pred = torch.sigmoid(pred).round()
+        return {f'{prefix}/{name}': metric(pred, target) for name, metric in self.metrics.items()}