-
Notifications
You must be signed in to change notification settings - Fork 4
/
data.py
83 lines (72 loc) · 2.56 KB
/
data.py
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
import glob
import os
import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
class BreastCancerDataset(Dataset):
def __init__(self, root, mode, split=0.95):
if mode not in {"train", "valid"}:
raise ValueError
patient_ids = sorted(os.listdir(root))
split = int(split * len(patient_ids))
# train validation split
patient_ids = patient_ids[:split] if mode == "train" else patient_ids[split:]
self.positives = []
self.negatives = []
for patient_id in patient_ids:
for image_path in glob.glob(os.path.join(root, patient_id, "*/*.png")):
if image_path.endswith("1.png"):
self.positives.append(image_path)
else:
self.negatives.append(image_path)
# transforms
self.transforms = None
if mode == "train":
# data augmentation during training
self.transforms = transforms.Compose(
[transforms.RandomVerticalFlip(), transforms.RandomHorizontalFlip()]
)
@property
def pos_weight(self):
return len(self.negatives) / len(self.positives)
def __len__(self):
return len(self.positives) + len(self.negatives)
def __getitem__(self, i):
label = None
image_path = None
if i < len(self.positives):
label = 1.0
image_path = self.positives[i]
else:
label = 0.0
image_path = self.negatives[i - len(self.positives)]
image = Image.open(image_path)
if self.transforms is not None:
image = self.transforms(image)
return image, label
def collate_fn(batch):
"""custom collate_fn to support PIL batch"""
inputs = [item[0] for item in batch]
labels = torch.tensor([item[1] for item in batch])
return inputs, labels
def make_loaders(config):
train_dataset = BreastCancerDataset(config.data_path, "train", config.split)
valid_dataset = BreastCancerDataset(config.data_path, "valid", config.split)
train_loader = DataLoader(
train_dataset,
batch_size=config.batch_size,
num_workers=config.num_workers,
collate_fn=collate_fn,
pin_memory=True,
shuffle=True,
)
valid_loader = DataLoader(
valid_dataset,
batch_size=config.batch_size,
num_workers=config.num_workers,
collate_fn=collate_fn,
pin_memory=True,
shuffle=False,
)
return train_loader, valid_loader