<reponame>rcorrero/slrgd<filename>slrgd/torch_tiles/charcoal_dataset.py<gh_stars>0
from typing import Optional, Union

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset


class CharcoalDataset(Dataset):
    '''
    Inherits from `torch` class `Dataset`. Handles input of training samples to `torch` models.

    Parameters
    ----------
    data_dict : dict
        A dataset dictionary produced by the `make_data` function.

    mode : str, default=train
        Determines whether to return the tuple `(X, y)` (when `mode=train`) or `X` (when `mode=test`).

    C : int, default=4
        Number of channels for each sample.

    size : int, default=5
        The width and height of each sample.

    img_max : int, default=10000
        The maximum attainable value for the samples. For Planet imagery, reflectances (which range from 0 to 1) are multiplied by 10000, resulting in images with integer values ranging from 0 to 10000. These values may be too large for gradient calculations.
    '''
    def __init__(self, data_dict: dict, mode: Optional[str] = 'train', 
                 C: Optional[int] = 4, size: Optional[int] = 5, 
                 img_max: Optional[int] = 10000) -> None:
        self.data_dict = data_dict
        assert mode in ['train', 'prod'], "Mode must be one of ('train', 'prod')."
        self.mode = mode
        self.idx_map, self.len = self._calc_len(data_dict)
        self.C = C
        self.size = size
        self.img_max = img_max


    def __getitem__(self, idx, unflatten: Optional[bool] = True,
                    transform: Optional[bool] = True):
        '''
        Returns samples from `data_dict`.

        Parameters
        ----------
        idx : int
            The index of the sample to return.

        unflatten : bool, default=True
            Convert the raw features (which are `size*size*C` arrays) into `(C, size, size)` multidimensional arrays.

        transform : bool, default = True
            Divide the raw feature values by `self.img_max`. See the docstring for the `CharcoalDataset` class for more information about `img_max`.

        Returns
        -------
        {(X_idx, y_idx), X_idx} : {tuple, array-like}
        '''
        filepath = self._get_path(idx)
        assert filepath is not None
        min_idx, max_idx = self.idx_map[filepath]
        assert idx >= min_idx and idx < max_idx
        idx -= min_idx
        X_idx = self.data_dict[filepath]['X'][idx]
        if unflatten:
            assert X_idx.shape[0] == self.C * self.size * self.size, "Passed values of `C` \
                and/or `size` do not match feature size."
            X_idx = torch.tensor(X_idx, 
                                 dtype=torch.int).unflatten(-1, (self.C, 
                                                                     self.size, 
                                                                     self.size))
            if transform:
                X_idx = X_idx.float() / self.img_max
        if self.mode == 'train':
            y_idx = torch.tensor(int(self.data_dict[filepath]['y'][idx]), 
                                 dtype=torch.float)
        else:
            return X_idx
        return X_idx, y_idx


    def __len__(self) -> int:
        return self.len


    def get_data(self, elem: str, 
                 filepath: Optional[str] = None) -> Union[np.ndarray, int]:
        '''
        Returns `elem` from `data_dict` corresponding to `filepath`.

        Parameters
        ----------
        elem : str
            The key for the desired element. Must be one of ('W', 'H', 'X', 'y').

        filepath : str, default = list(self.data_dict.keys())[0]
            The path associated with the desired image data. Defaults to the 
            first path in `data_dict.keys()`.
        '''
        valid_keys = list(self.data_dict.values())[0].keys()
        assert elem in valid_keys, "`elem` %s is invalid" % elem
        if filepath is None:
            filepath = list(self.data_dict.keys())[0]
        assert filepath in self.data_dict.keys(), "Filepath not found."
        return self.data_dict[filepath][elem]


    def get_original_dims(self, filepath: Union[str, dict]) -> dict:
        '''
        Get the width and height of the original image from which a set of
        features was generated.

        Parameters
        ----------
        filepath : str, default = list(self.data_dict.keys())[0]
            The path associated with the desired image data. Defaults to the 
            first path in `data_dict.keys()`.

        Returns
        -------
        dims : dict
            A dictionary containing the original image's width (W) and height 
            (H). Note that the image, when loaded as an array, had dimensions
            (C, H, W) where C denotes the number of channels for the image.
        '''
        if type(filepath) == dict:
            filepath = list(filepath.keys())[0]
        orig_img = Image.open(filepath)
        W, H = orig_img.size
        return {'W': W, 'H': H}


    def _calc_len(self, data_dict: dict) -> Union[dict, int]:
        idx_map = {} # `dict` mapping indices to filepaths
        num_samples = 0
        for filepath, data in data_dict.items():
            new_num_samples = data['y'].shape[0]
            idx_map[filepath] = (num_samples, num_samples + new_num_samples)
            num_samples += new_num_samples
        return idx_map, num_samples


    def _get_path(self, idx: int) -> str:
        '''
        Returns the path to file for which `idx` is in range.
        '''
        for filepath, (min_idx, max_idx) in self.idx_map.items():
            if idx >= min_idx and idx < max_idx:
                assert filepath is not None, 'Index out of range.'
                return filepath


    def _get_idx_map(self) -> dict:
        return self.idx_map