-
Notifications
You must be signed in to change notification settings - Fork 0
/
deband_dataset.py
37 lines (31 loc) · 1.14 KB
/
deband_dataset.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
import pandas as pd
from torch.utils.data import Dataset
from utils import read_image, cv_to_tensor
import albumentations as A
transform = A.Compose(
[
A.Flip(p=0.8),
A.ShiftScaleRotate(p=1.0,rotate_limit=180, shift_limit = 0, scale_limit = 0),
A.RandomBrightnessContrast(p=0.5),
],
additional_targets={'wet': 'image'}
)
class imageDebandDataset(Dataset):
def __init__(self, csv, root_dir= "",test=False ,transform=None):
self.frame = pd.read_csv(csv)
self.root_dir = root_dir
self.test = test
self.transform = transform
def __len__(self):
return len(self.frame)
def __getitem__(self, idx):
dry = self.frame.loc[idx, "dry"]
wet = self.frame.loc[idx, "banded"]
d = read_image(dry, to_tensor=False)
w = read_image(wet, to_tensor=False)
if not self.test and transform:
transformed = transform(image = d, wet = w)
sample = {"dry": cv_to_tensor(transformed["image"]), "wet": cv_to_tensor(transformed["wet"])}
else:
sample = {"dry": cv_to_tensor(d), "wet": cv_to_tensor(w)}
return sample