| import os |
| from typing import Tuple |
|
|
| import numpy as np |
| import scipy.io |
| import torch |
| import torch.utils.data as data |
|
|
| from auxiliary.utils import normalize, bgr_to_rgb, linear_to_nonlinear, hwc_to_chw |
| from classes.data.DataAugmenter import DataAugmenter |
|
|
|
|
| class ColorCheckerDataset(data.Dataset): |
|
|
| def __init__(self, train: bool = True, folds_num: int = 1): |
|
|
| self.__train = train |
| self.__da = DataAugmenter() |
|
|
| path_to_folds = os.path.join("dataset", "folds.mat") |
| path_to_metadata = os.path.join("dataset", "metadata.txt") |
| self.__path_to_data = os.path.join("dataset", "preprocessed", "numpy_data") |
| self.__path_to_label = os.path.join("dataset", "preprocessed", "numpy_labels") |
|
|
| folds = scipy.io.loadmat(path_to_folds) |
| img_idx = folds["tr_split" if self.__train else "te_split"][0][folds_num][0] |
|
|
| metadata = open(path_to_metadata, 'r').readlines() |
| self.__fold_data = [metadata[i - 1] for i in img_idx] |
|
|
| def __getitem__(self, index: int) -> Tuple: |
| file_name = self.__fold_data[index].strip().split(' ')[1] |
| img = np.array(np.load(os.path.join(self.__path_to_data, file_name + '.npy')), dtype='float32') |
| illuminant = np.array(np.load(os.path.join(self.__path_to_label, file_name + '.npy')), dtype='float32') |
|
|
| if self.__train: |
| img, illuminant = self.__da.augment(img, illuminant) |
| else: |
| img = self.__da.crop(img) |
|
|
| img = hwc_to_chw(linear_to_nonlinear(bgr_to_rgb(normalize(img)))) |
|
|
| img = torch.from_numpy(img.copy()) |
| illuminant = torch.from_numpy(illuminant.copy()) |
|
|
| if not self.__train: |
| img = img.type(torch.FloatTensor) |
|
|
| return img, illuminant, file_name |
|
|
| def __len__(self) -> int: |
| return len(self.__fold_data) |
|
|
