dataloader.py

import os
import json
from six import iteritems
from random import shuffle

import h5py
import numpy as np
from tqdm import tqdm

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset


class VisDialDataset(Dataset):

    @staticmethod
    def add_cmdline_args(parser):
        parser.add_argument_group('Dataloader specific arguments')
        parser.add_argument('-input_img', default='data/data_img.h5', help='HDF5 file with image features')
        parser.add_argument('-input_vid', default='data/data_video.h5', help='HDF5 file with image features')
        parser.add_argument('-input_audio', default='data/data_audio.h5', help='HDF5 file with audio features')
        parser.add_argument('-input_ques', default='data/dialogs.h5', help='HDF5 file with preprocessed questions')
        parser.add_argument('-input_json', default='data/params.json', help='JSON file with image paths and vocab')
        parser.add_argument('-img_norm', default=1, choices=[1, 0], help='normalize the image feature. 1=yes, 0=no')
        return parser

    def __init__(self, args, subsets):
        """Initialize the dataset with splits given by 'subsets', where
        subsets is taken from ['train', 'val', 'test']
        """
        super(VisDialDataset, self).__init__()
        self.args = args
        self.subsets = tuple(subsets)

        print("Dataloader loading json file: {}".format(args.input_json))
        with open(args.input_json, 'r') as info_file:
            info = json.load(info_file)
            # possible keys: {'ind2word', 'word2ind', 'unique_img_(split)'}
            for key, value in iteritems(info):
                setattr(self, key, value)

        # add <START> and <END> to vocabulary
        word_count = len(self.word2ind)
        self.word2ind['<START>'] = word_count + 1
        self.word2ind['<END>'] = word_count + 2
        self.start_token = self.word2ind['<START>']
        self.end_token = self.word2ind['<END>']

        # padding + <START> + <END> token
        self.vocab_size = word_count + 3
        print("Vocab size with <START>, <END>: {}".format(self.vocab_size))

        # construct reverse of word2ind after adding tokens
        self.ind2word = {
            int(ind): word
            for word, ind in iteritems(self.word2ind)
        }

        print("Dataloader loading h5 file: {}".format(args.input_ques))
        ques_file = h5py.File(args.input_ques, 'r')


        if 'image' in args.input_type:
            print("Dataloader loading h5 file: {}".format(args.input_img))
            img_file = h5py.File(args.input_img, 'r')

        if 'video' in args.input_type:
            print("Dataloader loading h5 file: {}".format(args.input_vid))
            vid_file = h5py.File(args.input_vid, 'r')

        if 'audio' in args.input_type:
            print("Dataloader loading h5 file: {}".format(args.input_audio))
            audio_file = h5py.File(args.input_audio, 'r')

        # load all data mats from ques_file into this
        self.data = {}

        # map from load to save labels
        io_map = {
            'ques_{}': '{}_ques',
            'ques_length_{}': '{}_ques_len',
            'ans_{}': '{}_ans',
            'ans_length_{}': '{}_ans_len',
            'img_pos_{}': '{}_img_pos',
            'cap_{}': '{}_cap',
            'cap_length_{}': '{}_cap_len',
            'opt_{}': '{}_opt',
            'opt_length_{}': '{}_opt_len',
            'opt_list_{}': '{}_opt_list',
            'num_rounds_{}': '{}_num_rounds',
            'ans_index_{}': '{}_ans_ind'
        }

        # processing every split in subsets
        for dtype in subsets:  # dtype is in ['train', 'val', 'test']
            print("\nProcessing split [{}]...".format(dtype))
            # read the question, answer, option related information
            for load_label, save_label in iteritems(io_map):
                if load_label.format(dtype) not in ques_file:
                    continue
                self.data[save_label.format(dtype)] = torch.from_numpy(
                    np.array(ques_file[load_label.format(dtype)], dtype='int64'))

            if 'video' in args.input_type:
                print("Reading video features...")
                vid_feats = torch.from_numpy(np.array(vid_file['images_' + dtype]))

                img_fnames = getattr(self, 'unique_img_' + dtype)
                self.data[dtype + '_img_fnames'] = img_fnames
                self.data[dtype + '_vid_fv'] = vid_feats


            if 'image' in args.input_type:
                print("Reading image features...")
                img_feats = torch.from_numpy(np.array(img_file['images_' + dtype]))

                if args.img_norm:
                    print("Normalizing image features...")
                    img_feats = F.normalize(img_feats, dim=1, p=2)

                img_fnames = getattr(self, 'unique_img_' + dtype)
                self.data[dtype + '_img_fnames'] = img_fnames
                self.data[dtype + '_img_fv'] = img_feats

            if 'audio' in args.input_type:
                print("Reading audio features...")
                audio_feats = torch.from_numpy(np.array(audio_file['images_' + dtype]))
                audio_feats = F.normalize(audio_feats, dim=1, p=2)
                self.data[dtype + '_audio_fv'] = audio_feats

            # record some stats, will be transferred to encoder/decoder later
            # assume similar stats across multiple data subsets
            # maximum number of questions per image, ideally 10
            self.max_ques_count = self.data[dtype + '_ques'].size(1)
            # maximum length of question
            self.max_ques_len = self.data[dtype + '_ques'].size(2)
            # maximum length of answer
            self.max_ans_len = self.data[dtype + '_ans'].size(2)

        # reduce amount of data for preprocessing in fast mode
        #TODO
        if args.overfit:
            print('\n \n \n ---------->> NOT IMPLEMENTED OVERFIT CASE  <-----\n \n \n ')

        self.num_data_points = {}
        for dtype in subsets:
            self.num_data_points[dtype] = len(self.data[dtype + '_ques'])
            print("[{0}] no. of threads: {1}".format(dtype, self.num_data_points[dtype]))
        print("\tMax no. of rounds: {}".format(self.max_ques_count))
        print("\tMax ques len: {}".format(self.max_ques_len))
        print("\tMax ans len: {}".format(self.max_ans_len))

        # prepare history
        if 'dialog' in args.input_type or 'caption' in args.input_type:
            for dtype in subsets:
                self._process_history(dtype)

        for dtype in subsets:
            # 1 indexed to 0 indexed
            self.data[dtype + '_opt'] -= 1
            if dtype + '_ans_ind' in self.data:
                self.data[dtype + '_ans_ind'] -= 1

        # default pytorch loader dtype is set to train
        if 'train' in subsets:
            self._split = 'train'
        else:
            self._split = subsets[0]

    @property
    def split(self):
        return self._split

    @split.setter
    def split(self, split):
        assert split in self.subsets  # ['train', 'val', 'test']
        self._split = split

    # ------------------------------------------------------------------------
    # methods to override - __len__ and __getitem__ methods
    # ------------------------------------------------------------------------

    def __len__(self):
        return self.num_data_points[self._split]

    def __getitem__(self, idx):
        dtype = self._split
        item = {'index': idx}
        item['num_rounds'] = self.data[dtype + '_num_rounds'][idx]

        # get video features
        if 'video' in self.args.input_type:
            item['vid_feat'] = self.data[dtype + '_vid_fv'][idx]
            item['img_fnames'] = self.data[dtype + '_img_fnames'][idx]

        # get image features
        if 'image' in self.args.input_type:
            item['img_feat'] = self.data[dtype + '_img_fv'][idx]
            item['img_fnames'] = self.data[dtype + '_img_fnames'][idx]

        # get audio features
        if 'audio' in self.args.input_type:
            item['audio_feat'] = self.data[dtype + '_audio_fv'][idx]

        # get history tokens
        if 'dialog' in self.args.input_type or 'caption' in self.args.input_type:
            item['hist_len'] = self.data[dtype + '_hist_len'][idx]
            item['hist_len'][item['hist_len'] == 0] += 1
            item['hist'] = self.data[dtype + '_hist'][idx]

        # get question tokens
        item['ques'] = self.data[dtype + '_ques'][idx]
        item['ques_len'] = self.data[dtype + '_ques_len'][idx]

        # get options tokens
        opt_inds = self.data[dtype + '_opt'][idx]
        opt_size = list(opt_inds.size())
        new_size = torch.Size(opt_size + [-1])
        ind_vector = opt_inds.view(-1)

        option_in = self.data[dtype + '_opt_list'].index_select(0, ind_vector)
        opt_len = self.data[dtype + '_opt_len'].index_select(0, ind_vector)
        option_in = option_in.view(new_size)
        opt_len = opt_len.view(opt_size)

        item['opt'] = option_in
        item['opt_len'] = opt_len
        #if dtype != 'test':
        ans_ind = self.data[dtype + '_ans_ind'][idx]
        item['ans_ind'] = ans_ind.view(-1)

        # convert zero length sequences to one length
        # this is for handling empty rounds of v1.0 test, they will be dropped anyway
        #if dtype == 'test':
        item['ques_len'][item['ques_len'] == 0] += 1
        item['opt_len'][item['opt_len'] == 0] += 1
        return item

    #-------------------------------------------------------------------------
    # collate function utilized by dataloader for batching
    #-------------------------------------------------------------------------

    def collate_fn(self, batch):
        dtype = self._split
        merged_batch = {key: [d[key] for d in batch] for key in batch[0]}
        out = {}
        for key in merged_batch:
            if key in {'index', 'num_rounds', 'img_fnames'}:
                out[key] = merged_batch[key]
            elif key in {'cap_len'}:
                out[key] = torch.Tensor(merged_batch[key]).long()
            else:
                out[key] = torch.stack(merged_batch[key], 0)
        # Dynamic shaping of padded batch
        if 'hist' in out:
            out['hist'] = out['hist'][:, :, :torch.max(out['hist_len'])].contiguous()
        out['ques'] = out['ques'][:, :, :torch.max(out['ques_len'])].contiguous()
        out['opt'] = out['opt'][:, :, :, :torch.max(out['opt_len'])].contiguous()

        return out

    #-------------------------------------------------------------------------
    # preprocessing functions
    #-------------------------------------------------------------------------

    def _process_history(self, dtype):
        """Process caption as well as history. Optionally, concatenate history
        for lf-encoder."""
        captions = self.data[dtype + '_cap']
        questions = self.data[dtype + '_ques']
        ques_len = self.data[dtype + '_ques_len']
        cap_len = self.data[dtype + '_cap_len']
        max_ques_len = questions.size(2)

        answers = self.data[dtype + '_ans']
        ans_len = self.data[dtype + '_ans_len']
        num_convs, num_rounds, max_ans_len = answers.size()

        if self.args.concat_history:
            self.max_hist_len = min(num_rounds * (max_ques_len + max_ans_len), 400)
            history = torch.zeros(num_convs, num_rounds, self.max_hist_len).long()
        else:
            history = torch.zeros(num_convs, num_rounds, max_ques_len + max_ans_len).long()
        hist_len = torch.zeros(num_convs, num_rounds).long()

        if 'dialog' in self.args.input_type:
            # go over each question and append it with answer
            for th_id in range(num_convs):
                clen = cap_len[th_id]
                hlen = min(clen, max_ques_len + max_ans_len)
                for round_id in range(num_rounds):
                    if round_id == 0:
                        # first round has caption as history
                        history[th_id][round_id][:max_ques_len + max_ans_len] \
                            = captions[th_id][:max_ques_len + max_ans_len]
                    else:
                        qlen = ques_len[th_id][round_id - 1]
                        alen = ans_len[th_id][round_id - 1]
                        # if concat_history, string together all previous question-answer pairs
                        if self.args.concat_history:
                            history[th_id][round_id][:hlen] = history[th_id][round_id - 1][:hlen]
                            history[th_id][round_id][hlen] = self.word2ind['<END>']
                            if qlen > 0:
                                history[th_id][round_id][hlen + 1:hlen + qlen + 1] \
                                    = questions[th_id][round_id - 1][:qlen]
                            if alen > 0:
                                # print(round_id, history[th_id][round_id][:10], answers[th_id][round_id][:10])
                                history[th_id][round_id][hlen + qlen + 1:hlen + qlen + alen + 1] \
                                    = answers[th_id][round_id - 1][:alen]
                            hlen = hlen + qlen + alen + 1
                        # else, history is just previous round question-answer pair
                        else:
                            if qlen > 0:
                                history[th_id][round_id][:qlen] = questions[th_id][round_id - 1][:qlen]
                            if alen > 0:
                                history[th_id][round_id][qlen:qlen + alen] \
                                    = answers[th_id][round_id - 1][:alen]
                            hlen = alen + qlen
                    # save the history length
                    hist_len[th_id][round_id] = hlen
        else: # -- caption only
            # go over each question and append it with answer
            for th_id in range(num_convs):
                clen = cap_len[th_id]
                hlen = min(clen, max_ques_len + max_ans_len)
                for round_id in range(num_rounds):
                    history[th_id][round_id][:max_ques_len + max_ans_len] = captions[th_id][:max_ques_len + max_ans_len]
                    hist_len[th_id][round_id] = hlen

        self.data[dtype + '_hist'] = history
        self.data[dtype + '_hist_len'] = hist_len