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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
import os
from PIL import Image
import numpy as np
import cv2
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from sklearn.model_selection import train_test_split
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
class CLAHETransform:
"""Apply CLAHE to the L channel of LAB color space.
Operates on PIL Images; returns a PIL Image.
"""
def __init__(self, clip_limit: float = 2.0, tile_grid_size: tuple = (8, 8)):
self.clip_limit = clip_limit
self.tile_grid_size = tile_grid_size
def __call__(self, img: Image.Image) -> Image.Image:
img_np = np.array(img) # RGB uint8
lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
clahe = cv2.createCLAHE(clipLimit=self.clip_limit, tileGridSize=self.tile_grid_size)
lab[:, :, 0] = clahe.apply(lab[:, :, 0])
result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
return Image.fromarray(result)
class DRDataset(Dataset):
"""Diabetic Retinopathy dataset from flat class-folder layout."""
def __init__(self, image_paths: list, labels: list, transform=None):
self.image_paths = image_paths
self.labels = labels
self.transform = transform
def __len__(self) -> int:
return len(self.image_paths)
def __getitem__(self, idx: int):
img = Image.open(self.image_paths[idx]).convert('RGB')
label = self.labels[idx]
if self.transform is not None:
img = self.transform(img)
return img, label
def split_dataset(data_root: str, seed: int = 42):
"""Scan class folders 0–4 and return stratified 70/15/15 splits.
Returns:
(train_paths, train_labels, val_paths, val_labels, test_paths, test_labels)
"""
all_paths, all_labels = [], []
for class_idx in range(5):
class_dir = os.path.join(data_root, str(class_idx))
for fname in sorted(os.listdir(class_dir)):
if fname.lower().endswith(('.jpg', '.jpeg', '.png')):
all_paths.append(os.path.join(class_dir, fname))
all_labels.append(class_idx)
# 70% train, 30% temp
train_paths, temp_paths, train_labels, temp_labels = train_test_split(
all_paths, all_labels, test_size=0.30, stratify=all_labels, random_state=seed
)
# 50/50 of temp -> 15% val, 15% test
val_paths, test_paths, val_labels, test_labels = train_test_split(
temp_paths, temp_labels, test_size=0.50, stratify=temp_labels, random_state=seed
)
print(f"Dataset split — train: {len(train_paths)}, val: {len(val_paths)}, test: {len(test_paths)}")
return train_paths, train_labels, val_paths, val_labels, test_paths, test_labels
def build_transforms(use_clahe: bool, training: bool):
"""Build a transform pipeline for training or eval."""
ops = [transforms.Resize((224, 224))]
if use_clahe:
ops.append(CLAHETransform(clip_limit=2.0, tile_grid_size=(8, 8)))
if training:
ops += [
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomVerticalFlip(p=0.5),
transforms.RandomRotation(degrees=15),
]
ops += [
transforms.ToTensor(),
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
]
return transforms.Compose(ops)
def build_dataloaders(
train_paths, train_labels,
val_paths, val_labels,
test_paths, test_labels,
use_clahe: bool,
batch_size: int = 32,
num_workers: int = 4,
):
"""Create DataLoaders for all three splits."""
train_ds = DRDataset(train_paths, train_labels, transform=build_transforms(use_clahe, training=True))
val_ds = DRDataset(val_paths, val_labels, transform=build_transforms(use_clahe, training=False))
test_ds = DRDataset(test_paths, test_labels, transform=build_transforms(use_clahe, training=False))
train_loader = DataLoader(
train_ds, batch_size=batch_size, shuffle=True,
num_workers=num_workers, pin_memory=True
)
val_loader = DataLoader(
val_ds, batch_size=batch_size * 2, shuffle=False,
num_workers=num_workers, pin_memory=True
)
test_loader = DataLoader(
test_ds, batch_size=batch_size * 2, shuffle=False,
num_workers=num_workers, pin_memory=True
)
return train_loader, val_loader, test_loader
|
