data.py

import argparse
import os
import ipdb
import numpy as np
import scipy.io as sio
import pandas as pd
import torch
import pickle
from utils import get_datadir, labels_mapping, get_data_stats
from sklearn.metrics import accuracy_score
from sklearn import preprocessing
from torch.utils.data import Dataset


BASE = ['X', 'Y', 'emb', 'mapping', 'classnames']

TRAIN = ['train_X', 'train_Y',
         'train_emb', 'train_mapping', 'train_classnames']
VAL_SEEN = ['val_seen_X', 'val_seen_Y',
             'val_seen_emb', 'val_seen_mapping', 'val_seen_classnames']
VAL_UNSEEN = ['val_unseen_X', 'val_unseen_Y',
               'val_unseen_emb', 'val_unseen_mapping',
               'val_unseen_classnames']
TRAINVAL = ['trainval_X', 'trainval_Y',
            'trainval_emb', 'trainval_mapping', 'trainval_classnames']
TEST_SEEN = ['test_seen_X', 'test_seen_Y',
             'test_seen_emb', 'test_seen_mapping', 'test_seen_classnames']
TEST_UNSEEN = ['test_unseen_X', 'test_unseen_Y',
               'test_unseen_emb', 'test_unseen_mapping',
               'test_unseen_classnames']


def sample_imgs(y, idx, ratio=0.2):
    y_unique, y_counts = np.unique(y, return_counts=True)
    idx_val = []
    for val, n_per_class in zip(y_unique, y_counts):
        cond = np.squeeze(np.argwhere(y == val))
        subset = np.random.choice(idx[cond],
                                  replace=False,
                                  size=int(n_per_class*ratio))
        idx_val.extend(subset)
    idx_train = idx[np.isin(idx, idx_val, invert=True)]
    return idx_train, np.asarray(idx_val)


def create_train_val_splits(y_trainval, idx_trainval, n_unseen,
                            seed=123):
    np.random.seed(seed)

    # pick unseen categories
    y_unique = np.unique(y_trainval)
    y_val_unseen = np.random.choice(y_unique, replace=False, size=n_unseen)

    # get val unseen
    cond = np.isin(y_trainval, y_val_unseen)
    idx_val_unseen = idx_trainval[cond]

    y_seen = y_unique[~np.isin(y_unique, y_val_unseen)]
    cond = np.isin(y_trainval, y_seen)
    idx_seen = idx_trainval[cond]

    idx_train, idx_val_seen = sample_imgs(y_trainval[cond], idx_seen)

    return idx_train, idx_val_seen, idx_val_unseen


def load_data_gan(path_gan, path_embedding,
                  val_classes=10,
                  path_sentences=None,
                  normalize=False):
    '''Load features and embedding

    Returns dict
    split_X: (n, 2048) array
    split_Y: (n) array
        from 0 to C_split classes
    split_emb: (A, C_split) array
        attribute representation for each class in the split
    split_mapping: dict
        `direct` class index mapping from full to split
        `indirect` class index mapping from split to full
    split_classnames: dict
        class name according to the full class index

    Note: the test seen split is leaking into train and val splits
          MUST use trainval split only
          or create our own train and val splits
          rather using than those from Xian et al, TPAMI 2018
    '''
    feat = np.load(path_gan)
    if 'train_Y' in feat.files:
        y = feat['train_Y']
    else:
        y = np.concatenate((feat['train_seen_Y'], feat['train_unseen_Y']))

    get_sentences = False
    if path_sentences is not None:
        get_sentences = True
        with open(path_sentences, 'rb') as f:
            sentences = pickle.load(f)['sentences'].T

    embedding = sio.loadmat(path_embedding)
    if 'allclasses_names' in embedding.keys():
        names = embedding['allclasses_names']
    else:
        names =[[[str(i)]] for i in range(embedding['att'].shape[1])]
    full_classnames = classnames_parser(y, names)

    # train and validation merged
    if 'train_X' in feat.files:
        trainval_X, trainval_Y_orig = feat['train_X'], feat['train_Y']
    else:
        trainval_X = np.concatenate((feat['train_seen_X'],
                                     feat['train_unseen_X']))
        trainval_Y_orig = np.concatenate((feat['train_seen_Y'],
                                          feat['train_unseen_Y']))

    trainval_Y, trainval_mapping = labels_mapping(trainval_Y_orig)
    trainval_emb = embedding['att'][:, np.unique(trainval_Y_orig)]
    trainval_classnames = classnames_parser(trainval_Y_orig, names)

    # testing splits
    test_seen_X, test_seen_Y_orig = feat['test_seen_X'], feat['test_seen_Y']
    test_seen_Y, test_seen_mapping = labels_mapping(test_seen_Y_orig)
    test_seen_emb = embedding['att'][:, np.unique(test_seen_Y_orig)]
    test_seen_classnames = classnames_parser(test_seen_Y_orig, names)

    test_unseen_X, test_unseen_Y_orig = feat['test_unseen_X'], feat['test_unseen_Y']
    test_unseen_Y, test_unseen_mapping = labels_mapping(test_unseen_Y_orig)
    test_unseen_emb = embedding['att'][:, np.unique(test_unseen_Y_orig)]
    test_unseen_classnames = classnames_parser(test_unseen_Y_orig, names)

    # class-attribute matrix for all classes (used in GZSL)
    full_emb = embedding['att']

    if normalize:
        scaler = preprocessing.MinMaxScaler()
        trainval_X = scaler.fit_transform(trainval_X)
        test_seen_X = scaler.transform(test_seen_X)
        test_unseen_X = scaler.transform(test_unseen_X)

    if get_sentences:
        trainval_sentences = sentences[:, np.unique(trainval_Y_orig)]
        test_seen_sentences = sentences[:, np.unique(test_seen_Y_orig)]
        test_unseen_sentences = sentences[:, np.unique(test_unseen_Y_orig)]

        return {# trainval split
                'trainval_X': trainval_X.astype(np.float32),
                'trainval_Y': trainval_Y,
                'trainval_emb': trainval_emb.astype(np.float32),
                'trainval_mapping': trainval_mapping,
                'trainval_classnames': trainval_classnames,
                'trainval_sentences': trainval_sentences.astype(np.float32),
                # 2 test splits
                'test_seen_X': test_seen_X.astype(np.float32),
                'test_seen_Y': test_seen_Y,
                'test_seen_emb': test_seen_emb.astype(np.float32),
                'test_seen_mapping': test_seen_mapping,
                'test_seen_classnames': test_seen_classnames,
                'test_seen_sentences': test_seen_sentences.astype(np.float32),
                'test_unseen_X': test_unseen_X.astype(np.float32),
                'test_unseen_Y': test_unseen_Y,
                'test_unseen_emb': test_unseen_emb.astype(np.float32),
                'test_unseen_mapping': test_unseen_mapping,
                'test_unseen_classnames': test_unseen_classnames,
                'test_unseen_sentences': test_unseen_sentences.astype(np.float32),
                # full cam
                'test_seen_Y_orig': test_seen_Y_orig.astype(np.int),
                'test_unseen_Y_orig': test_unseen_Y_orig.astype(np.int),
                'full_emb': full_emb.astype(np.float32),
                'full_classnames': full_classnames,
                'full_sentences': sentences.astype(np.float32),
                'seen_Y': np.unique(test_seen_Y_orig.astype(np.int)),
                'unseen_Y': np.unique(test_unseen_Y_orig.astype(np.int))}

    return {# trainval split
            'trainval_X': trainval_X.astype(np.float32),
            'trainval_Y': trainval_Y,
            'trainval_emb': trainval_emb.astype(np.float32),
            'trainval_mapping': trainval_mapping,
            'trainval_classnames': trainval_classnames,
            # 2 test splits
            'test_seen_X': test_seen_X.astype(np.float32),
            'test_seen_Y': test_seen_Y,
            'test_seen_emb': test_seen_emb.astype(np.float32),
            'test_seen_mapping': test_seen_mapping,
            'test_seen_classnames': test_seen_classnames,
            'test_unseen_X': test_unseen_X.astype(np.float32),
            'test_unseen_Y': test_unseen_Y,
            'test_unseen_emb': test_unseen_emb.astype(np.float32),
            'test_unseen_mapping': test_unseen_mapping,
            'test_unseen_classnames': test_unseen_classnames,
            # full cam
            'test_seen_Y_orig': test_seen_Y_orig.astype(np.int),
            'test_unseen_Y_orig': test_unseen_Y_orig.astype(np.int),
            'full_emb': full_emb.astype(np.float32),
            'full_classnames': full_classnames,
            'seen_Y': np.unique(test_seen_Y_orig.astype(np.int)),
            'unseen_Y': np.unique(test_unseen_Y_orig.astype(np.int))}


def load_data_separate(path_features, path_embedding,
                       val_classes=10,
                       path_sentences=None,
                       normalize=False):
    features = sio.loadmat(path_features)
    x = features['features'].T
    y = features['labels'] - 1  # labels start at 1, need to subtract 1

    get_sentences = False
    if path_sentences is not None:
        get_sentences = True
        with open(path_sentences, 'rb') as f:
            sentences = pickle.load(f)['sentences'].T

    embedding = sio.loadmat(path_embedding)
    if 'allclasses_names' in embedding.keys():
        names = embedding['allclasses_names']
    else:
        names =[[[str(i)]] for i in range(embedding['att'].shape[1])]
    full_classnames = classnames_parser(y, names)

    # train and validation merged
    idx = np.squeeze(embedding['trainval_loc']) - 1
    trainval_X, trainval_Y_orig = x[idx, :], np.squeeze(y[idx])
    trainval_Y, trainval_mapping = labels_mapping(trainval_Y_orig)
    trainval_emb = embedding['att'][:, np.unique(trainval_Y_orig)]
    trainval_classnames = classnames_parser(trainval_Y_orig, names)

    # new training and validation splits
    # with 0.8 train_seen and val_seen
    idx_train, idx_val_seen, idx_val_unseen = create_train_val_splits(
        trainval_Y_orig, idx, val_classes)

    # training split
    train_X, train_Y_orig = x[idx_train, :], np.squeeze(y[idx_train])
    train_Y, train_mapping = labels_mapping(train_Y_orig)
    train_emb = embedding['att'][:, np.unique(train_Y_orig)]
    train_classnames = classnames_parser(train_Y_orig, names)

    # validation splits
    val_seen_X, val_seen_Y_orig = x[idx_val_seen, :], np.squeeze(y[idx_val_seen])
    val_seen_Y, val_seen_mapping = labels_mapping(val_seen_Y_orig)
    val_seen_emb = embedding['att'][:, np.unique(val_seen_Y_orig)]
    val_seen_classnames = classnames_parser(val_seen_Y_orig, names)

    val_unseen_X, val_unseen_Y_orig = x[idx_val_unseen, :], np.squeeze(y[idx_val_unseen])
    val_unseen_Y, val_unseen_mapping = labels_mapping(val_unseen_Y_orig)
    val_unseen_emb = embedding['att'][:, np.unique(val_unseen_Y_orig)]
    val_unseen_classnames = classnames_parser(val_unseen_Y_orig, names)

    # testing splits
    idx = np.squeeze(embedding['test_seen_loc']) - 1
    test_seen_X, test_seen_Y_orig = x[idx, :], np.squeeze(y[idx])
    test_seen_Y, test_seen_mapping = labels_mapping(test_seen_Y_orig)
    test_seen_emb = embedding['att'][:, np.unique(test_seen_Y_orig)]
    test_seen_classnames = classnames_parser(test_seen_Y_orig, names)

    idx = np.squeeze(embedding['test_unseen_loc']) - 1
    test_unseen_X, test_unseen_Y_orig = x[idx, :], np.squeeze(y[idx])
    test_unseen_Y, test_unseen_mapping = labels_mapping(test_unseen_Y_orig)
    test_unseen_emb = embedding['att'][:, np.unique(test_unseen_Y_orig)]
    test_unseen_classnames = classnames_parser(test_unseen_Y_orig, names)

    # class-attribute matrix for all classes (used in GZSL)
    full_emb = embedding['att']

    if normalize:
        scaler = preprocessing.MinMaxScaler()
        trainval_X = scaler.fit_transform(trainval_X)
        test_seen_X = scaler.transform(test_seen_X)
        test_unseen_X = scaler.transform(test_unseen_X)

    if get_sentences:
        trainval_sentences = sentences[:, np.unique(trainval_Y_orig)]
        train_sentences = sentences[:, np.unique(train_Y_orig)]
        val_seen_sentences = sentences[:, np.unique(val_seen_Y_orig)]
        val_unseen_sentences = sentences[:, np.unique(val_unseen_Y_orig)]
        test_seen_sentences = sentences[:, np.unique(test_seen_Y_orig)]
        test_unseen_sentences = sentences[:, np.unique(test_unseen_Y_orig)]

        return {# training split
                'train_X': train_X.astype(np.float32),
                'train_Y': train_Y,
                'train_emb': train_emb.astype(np.float32),
                'train_mapping': train_mapping,
                'train_classnames': train_classnames,
                'train_sentences': train_sentences.astype(np.float32),
                # trainval split
                'trainval_X': trainval_X.astype(np.float32),
                'trainval_Y': trainval_Y,
                'trainval_emb': trainval_emb.astype(np.float32),
                'trainval_mapping': trainval_mapping,
                'trainval_classnames': trainval_classnames,
                'trainval_sentences': trainval_sentences.astype(np.float32),
                # 2 validation splits
                'val_seen_X': val_seen_X.astype(np.float32),
                'val_seen_Y': val_seen_Y,
                'val_seen_emb': val_seen_emb.astype(np.float32),
                'val_seen_mapping': val_seen_mapping,
                'val_seen_classnames': val_seen_classnames,
                'val_seen_sentences': val_seen_sentences.astype(np.float32),
                'val_unseen_X': val_unseen_X.astype(np.float32),
                'val_unseen_Y': val_unseen_Y,
                'val_unseen_emb': val_unseen_emb.astype(np.float32),
                'val_unseen_mapping': val_unseen_mapping,
                'val_unseen_classnames': val_unseen_classnames,
                'val_unseen_sentences': val_unseen_sentences.astype(np.float32),
                # 2 test splits
                'test_seen_X': test_seen_X.astype(np.float32),
                'test_seen_Y': test_seen_Y,
                'test_seen_emb': test_seen_emb.astype(np.float32),
                'test_seen_mapping': test_seen_mapping,
                'test_seen_classnames': test_seen_classnames,
                'test_seen_sentences': test_seen_sentences.astype(np.float32),
                'test_unseen_X': test_unseen_X.astype(np.float32),
                'test_unseen_Y': test_unseen_Y,
                'test_unseen_emb': test_unseen_emb.astype(np.float32),
                'test_unseen_mapping': test_unseen_mapping,
                'test_unseen_classnames': test_unseen_classnames,
                'test_unseen_sentences': test_unseen_sentences.astype(np.float32),
                # full cam
                'val_seen_Y_orig': val_seen_Y_orig.astype(np.int),
                'val_unseen_Y_orig': val_unseen_Y_orig.astype(np.int),
                'test_seen_Y_orig': test_seen_Y_orig.astype(np.int),
                'test_unseen_Y_orig': test_unseen_Y_orig.astype(np.int),
                'full_emb': full_emb.astype(np.float32),
                'full_classnames': full_classnames,
                'full_sentences': sentences.astype(np.float32)}

    return {# training split
            'train_X': train_X.astype(np.float32),
            'train_Y': train_Y,
            'train_emb': train_emb.astype(np.float32),
            'train_mapping': train_mapping,
            'train_classnames': train_classnames,
            # trainval split
            'trainval_X': trainval_X.astype(np.float32),
            'trainval_Y': trainval_Y,
            'trainval_emb': trainval_emb.astype(np.float32),
            'trainval_mapping': trainval_mapping,
            'trainval_classnames': trainval_classnames,
            # 2 validation splits
            'val_seen_X': val_seen_X.astype(np.float32),
            'val_seen_Y': val_seen_Y,
            'val_seen_emb': val_seen_emb.astype(np.float32),
            'val_seen_mapping': val_seen_mapping,
            'val_seen_classnames': val_seen_classnames,
            'val_unseen_X': val_unseen_X.astype(np.float32),
            'val_unseen_Y': val_unseen_Y,
            'val_unseen_emb': val_unseen_emb.astype(np.float32),
            'val_unseen_mapping': val_unseen_mapping,
            'val_unseen_classnames': val_unseen_classnames,
            # 2 test splits
            'test_seen_X': test_seen_X.astype(np.float32),
            'test_seen_Y': test_seen_Y,
            'test_seen_emb': test_seen_emb.astype(np.float32),
            'test_seen_mapping': test_seen_mapping,
            'test_seen_classnames': test_seen_classnames,
            'test_unseen_X': test_unseen_X.astype(np.float32),
            'test_unseen_Y': test_unseen_Y,
            'test_unseen_emb': test_unseen_emb.astype(np.float32),
            'test_unseen_mapping': test_unseen_mapping,
            'test_unseen_classnames': test_unseen_classnames,
            # full cam
            'val_seen_Y_orig': val_seen_Y_orig.astype(np.int),
            'val_unseen_Y_orig': val_unseen_Y_orig.astype(np.int),
            'test_seen_Y_orig': test_seen_Y_orig.astype(np.int),
            'test_unseen_Y_orig': test_unseen_Y_orig.astype(np.int),
            'full_emb': full_emb.astype(np.float32),
            'full_classnames': full_classnames}


def classnames_parser(labels, names):
    classnames = [n[0][0] for n in names]
    classnames = dict(zip(np.arange(len(classnames)), classnames))

    # note that labels have to start at 0 (not 1)
    unique_labels = np.unique(labels)

    classmapping = {}
    for l in unique_labels:
        classmapping[l] = classnames[l]
    return classmapping


class AttributeDataset(Dataset):
    def __init__(self, root, dataset,
                 features_path=None,
                 mode='trainval',
                 sentences=False,
                 generalized=False,
                 normalize=False,
                 both=False,
                 on_gpu=False):
        data_dir = get_datadir(root, dataset)
        if features_path is None:
            features_path = os.path.join(data_dir, 'res101.mat')
            load_data = load_data_separate
        else:
            load_data = load_data_gan
        meta_path = os.path.join(data_dir, 'att_splits.mat')

        _, n_val_classes, _, _ = get_data_stats(dataset, validation=True)

        if sentences:
            # works for CUB dataset
            sentences_path = os.path.join(data_dir, 'CUB_supporting_data.p')
            data = load_data(
                features_path, meta_path, n_val_classes,
                sentences_path, normalize=normalize)
        else:
            data = load_data(
                features_path, meta_path, n_val_classes,
                normalize=normalize)

        if mode == 'train':
            keys = TRAIN
        elif mode == 'val_seen':
            keys = VAL_SEEN
        elif mode == 'val_unseen':
            keys = VAL_UNSEEN
        elif mode == 'trainval':
            keys = TRAINVAL
        elif mode == 'test_seen':
            keys = TEST_SEEN
        elif mode == 'test_unseen':
            keys = TEST_UNSEEN

        self.data = {}
        if sentences:
            kk = keys + [mode + '_sentences']
            bb = BASE + ['sentences']
            for k, b in zip(kk, bb):
                if b in ['X', 'Y'] and on_gpu:
                    self.data[b] = torch.from_numpy(data[k]).cuda()
                else:
                    self.data[b] = data[k]
        else:
            for k, b in zip(keys, BASE):
                if b in ['X', 'Y'] and on_gpu:
                    self.data[b] = torch.from_numpy(data[k]).cuda()
                else:
                    self.data[b] = data[k]

        if generalized:
            self.data['Y_orig'] = data[mode + '_Y_orig']
            self.data['full_emb'] = data['full_emb']
            if sentences:
                self.data['full_sentences'] = data['full_sentences']

        self.mode = mode
        self.generalized = generalized
        self.both = both
        if self.both:
            self.data['seen_Y'] = data['seen_Y']
            self.data['unseen_Y'] = data['unseen_Y']

    def __getitem__(self, index):
        feat = self.data['X'][index, :]
        label = self.data['Y'][index]

        if self.generalized:
            label_g = self.data['Y_orig'][index]
            return feat, (label, label_g)
        else:
            if self.both:
                if label in self.data['seen_Y']:
                    seen = np.float32(1.)
                    unseen = np.float32(0.)
                elif label in self.data['unseen_Y']:
                    seen = np.float32(0.)
                    unseen = np.float32(1.)
                return feat, (label, seen, unseen)
            else:
                return feat, label

    def __len__(self):
        return self.data['X'].shape[0]