data.py

import os
import glob
import random
import numpy as np
import pandas as pd
from PIL import Image

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets.folder import DatasetFolder
from torchvision import transforms


def normalize_label(df, params):
    r"""Normalize labels for training"""

    if df['visibility'].max() > 20.0:
        vis_norm = 20000.0
    else:
        vis_norm = 20.0
    norm_dict = {'visibility': vis_norm, 'RH': 100.0, 'temperature': 50.0, 'pressure': 1050.0, 'PM25':500.0, 'PM10': 500.0}
    for param in params:
        df[param] = df[param] / norm_dict[param]
    return df


def recovery_df(df, params, vis_norm=20.0):
    r"""Recovery the dataframe"""
    
    norm_dict = {'visibility': vis_norm, 'RH': 100.0, 'temperature': 50.0, 'pressure': 1050.0, 'PM25':500.0, 'PM10': 500.0}
    for param in params:
        df[param] = df[param] * norm_dict[param]
        df['pred_'+param] = df['pred_'+param] * norm_dict[param]
    return df


class Array2Tensor(object):
    r"""Convert a hyperspectral or an RGB image into a tensor.
    """

    def __init__(self):
        super(Array2Tensor, self).__init__()

    def __call__(self, pic):

        if isinstance(pic, np.ndarray):

            pic = pic.astype(np.float)
            img = torch.from_numpy(pic.transpose((2, 0, 1)).copy())
            if img.max().item() > 255.0:
                return img.float().div(4095.0)
            else:
                return img.float().div(255.0)


def __crop_array(img, crop_size):
    oh, ow, oc = img.shape
    th = tw = crop_size
    x_min = oh - th
    y_min = ow - tw
    x1 = np.random.randint(x_min)
    y1 = np.random.randint(y_min)

    if (ow > tw or oh > th):
        return img[x1:x1+th, y1:y1+tw, :]
    return img


def __flip_array(img):
    if torch.rand(1) < 0.5:
        return np.flip(img, axis=1)
    return img


def get_transform(opt, mode):
    transform_list = []

    if mode == 'train':
        if 'crop' in opt.preprocess:
            transform_list.append(transforms.Lambda(
                lambda img: __crop_array(img, opt.crop_size)))

        if 'flip' in opt.preprocess:
            transform_list.append(transforms.Lambda(
                lambda img: __flip_array(img)))

    # convert to tensor
    transform_list += [Array2Tensor()]
    # normalization
    transform_list += [transforms.Normalize((0.5,)
                                            * opt.input_nc, (0.5,)*opt.input_nc)]
    return transforms.Compose(transform_list)

IMG_EXTENSIONS = (
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
    '.tif', '.TIF', '.tiff', '.TIFF', '.mat', '.npy', '.npz'
)


def default_loader(path):
    return Image.open(path).convert('RGB')


class MyImageFolder(DatasetFolder):

    def __init__(self, root, transform=None, loader=default_loader, extensions=IMG_EXTENSIONS):
        self.transform = transform
        DatasetFolder.__init__(
            self, root, loader, transform=transform, extensions=extensions)

    def __getitem__(self, index):
        file, label = self.samples[index]
        name = file.split(os.sep)[-1].split('.')[0]
        image = self.loader(file)
        #print('transform:', self.transform)
        if self.transform is not None:
            image = self.transform(image)
        return image, torch.tensor(label, dtype=torch.long), name


class MyDataset(Dataset):
    def __init__(self, opt, mode):
        self.labels = pd.read_csv(os.path.join(opt.dataroot, 'label.csv'))
        filepath = os.path.join(opt.dataroot, mode)
        self.files = glob.glob(filepath+'/*')
        self.lenx = len(self.files)
        if opt.model_type == 'reg':
            self.output_params = opt.output_params
            self.labels = normalize_label(self.labels, self.output_params)
            self.dtype = torch.float
        elif opt.model_type == 'clf':
            # we do not test this, please check this if needed
            # mask sure that there is a `.csv` file under the `dataset/clf_folder` path
            # and the names of one column is `name`, one column is `label`
            self.output_params = 'label'
            self.dtype = torch.long  
        else:
            raise ValueError(f'Unknown model type {opt.model_type}.')
        self.transform = get_transform(opt, mode)
        self.HSI = opt.HSI

    def __getitem__(self, index):
        file = self.files[index]
        name = file.split(os.sep)[-1].split('.')[0]

        idx = self.labels['name'] == name
        label_list = [float(self.labels[idx][output_param].values.item())
                      for output_param in self.output_params]
        label = torch.tensor(label_list, dtype=self.dtype)

        if self.HSI:
            image = np.load(file, allow_pickle=True)['image'].astype(float)
        else:
            image = np.array(Image.open(file).convert('RGB'))

        return self.transform(image), label, name
    
    def __len__(self):
        return self.lenx


def build_dataset(opt):

    if not opt.isTest:

        if (opt.model_type == 'reg') or (opt.model_type == 'clf'):
            tr_set = MyDataset(opt, 'train')
            val_set = MyDataset(opt, 'val')
        elif opt.model_type == 'clf_multi':
            tr_path = os.path.join(opt.dataroot, 'train')
            val_path = os.path.join(opt.dataroot, 'val')

            tr_set = MyImageFolder(tr_path, 
                        transform=transforms.Compose([
                            transforms.RandomCrop(opt.crop_size),
                            transforms.RandomHorizontalFlip(),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                            ]))

            val_set = MyImageFolder(val_path,
                        transform=transforms.Compose([
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                            ]))
        else:
            raise ValueError(f'Unknown model type {opt.model_type}.')

        print(f'Files in train and validation sets: {len(tr_set)} and {len(val_set)}.')
        tr_loader = DataLoader(tr_set, batch_size=opt.batch_size, shuffle=True,
                            pin_memory=True, num_workers=opt.num_workers, drop_last=True)
        val_loader = DataLoader(val_set, batch_size=opt.batch_size, shuffle=False,
                                pin_memory=True, num_workers=opt.num_workers, drop_last=True)    
        return tr_loader, val_loader

    else:
        if (opt.model_type == 'reg') or (opt.model_type == 'clf'):
            tr_set = MyDataset(opt, 'test')
            val_set = MyDataset(opt, 'test')
            te_set = MyDataset(opt, 'test')
        elif opt.model_type == 'clf_multi':
            # te_path = os.path.join(opt.dataroot, 'test')
            tr_set = MyImageFolder(os.path.join(opt.dataroot, 'test'),
                            transform=transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                ]))
            val_set = MyImageFolder(os.path.join(opt.dataroot, 'test'),
                            transform=transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                ]))
            te_set = MyImageFolder(os.path.join(opt.dataroot, 'test'),
                            transform=transforms.Compose([
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                ]))
        else:
            raise ValueError(f'Unknown model type {opt.model_type}.')
        print(f'Files in the training, validation, and test sets: {len(tr_set)}, {len(val_set)}, {len(te_set)}.')
        tr_loader = DataLoader(tr_set, batch_size=1, shuffle=False,
                                pin_memory=True, num_workers=opt.num_workers, drop_last=False)    
        val_loader = DataLoader(val_set, batch_size=1, shuffle=False,
                                pin_memory=True, num_workers=opt.num_workers, drop_last=False)    
        te_loader = DataLoader(te_set, batch_size=1, shuffle=False,
                                pin_memory=True, num_workers=opt.num_workers, drop_last=False)    
        return tr_loader, val_loader, te_loader