data.py

import json
import os
import os.path
import re

import glob
from PIL import Image
import h5py
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
import numpy as np

import config
import utils


preloaded_vocab = None


def get_loader(train=False, val=False, test=False):
    """ Returns a data loader for the desired split """
    split = VQA(
        utils.path_for(train=train, val=val, test=test, question=True),
        utils.path_for(train=train, val=val, test=test, answer=True),
        config.preprocessed_trainval_path if not test else config.preprocessed_test_path,
        answerable_only=train,
        dummy_answers=test,
    )
    #print(split.__len__(), "splitLEn")
    

    loader = torch.utils.data.DataLoader(
        split,
        batch_size=config.batch_size,
        shuffle=train,  # only shuffle the data in training
        pin_memory=True,
        num_workers=config.data_workers,
        collate_fn=collate_fn,
    )
    return loader



def collate_fn(batch):
    # put question lengths in descending order so that we can use packed sequences later
    batch.sort(key=lambda x: x[-1], reverse=True)
    return data.dataloader.default_collate(batch)


class VQA(data.Dataset):
    """ VQA dataset, open-ended """
    def __init__(self, questions_path, answers_path, image_features_path, answerable_only=False, dummy_answers=False):
        super(VQA, self).__init__()
        with open(questions_path, 'r') as fd:
            questions_json = json.load(fd)
        with open(answers_path, 'r') as fd:
            answers_json = json.load(fd)
        if preloaded_vocab:
            vocab_json = preloaded_vocab
        else:
            with open(config.vocabulary_path, 'r') as fd:
                vocab_json = json.load(fd)

        self.question_ids = [q['question_id'] for q in questions_json['questions']]


        # vocab
        self.vocab = vocab_json
        self.token_to_index = self.vocab['question']
        #print(len(self.vocab['question']), len(self.vocab['answer']), "YOLO")
        self.answer_to_index = self.vocab['answer']

        # if answerable_only == False:
        #     self.questions = list(prepare_questions(questions_json))    
        #     self.questions = [self._encode_question(q) for q in self.questions]
        #     print(len(self.questions))
        #     brak


        # q and a
        self.questions = questions_json#list(prepare_questions(questions_json))
       
        self.answers = answers_json #list(prepare_answers(answers_json))
        #self.questions = [self._encode_question(q) for q in self.questions]
        #self.answers = [self._encode_answers(a) for a in self.answers]
        
        self.max_question_length = 23#self.max_question_length
        
        
        # v
        self.image_features_path = image_features_path
        self.coco_id_to_index = self._create_coco_id_to_index()
        self.coco_ids = [q['image_id'] for q in questions_json['questions']]


        self.dummy_answers= dummy_answers

        # only use questions that have at least one answer?
        self.answerable_only = answerable_only
        if self.answerable_only:
            #self.answerable = self._find_answerable(not self.answerable_only)
            self.answerable = list(np.load("answerableList.npy"))
        #np.save("answerableList.npy", self.answerable)
        #brak

        #self.answers = answers_json
        print("Init done!")

    @property
    def max_question_length_old(self):
        if not hasattr(self, '_max_length'):
            self._max_length = max(map(len, self.questions))
        return self._max_length

    @property
    def num_tokens(self):
        return len(self.token_to_index) + 1  # add 1 for <unknown> token at index 0

    def _create_coco_id_to_index(self):
        """ Create a mapping from a COCO image id into the corresponding index into the h5 file """
        with h5py.File(self.image_features_path, 'r') as features_file:
            coco_ids = features_file['ids'][()]
        coco_id_to_index = {id: i for i, id in enumerate(coco_ids)}

       
        return coco_id_to_index

    def _create_coco_id_to_index_custom(self):
        coco_id_to_index = {}
        image_list = glob.glob("/home/user/data/mscoco/images/train2017/*.jpg") + glob.glob("/home/user/data/mscoco/images/val2017/*.jpg")
        for i,image_path in enumerate(image_list):
            ID = image_path.split("/")[-1].split(".")[0].lstrip("0")
            coco_id_to_index[ID] = i

        return coco_id_to_index



        


    def _check_integrity(self, questions, answers):
        """ Verify that we are using the correct data """
        qa_pairs = list(zip(questions['questions'], answers['annotations']))
        assert all(q['question_id'] == a['question_id'] for q, a in qa_pairs), 'Questions not aligned with answers'
        assert all(q['image_id'] == a['image_id'] for q, a in qa_pairs), 'Image id of question and answer don\'t match'
        assert questions['data_type'] == answers['data_type'], 'Mismatched data types'
        assert questions['data_subtype'] == answers['data_subtype'], 'Mismatched data subtypes'

    def _find_answerable(self, count=False):
        """ Create a list of indices into questions that will have at least one answer that is in the vocab """
        answerable = []
        if count:
            number_indices = torch.LongTensor([self.answer_to_index[str(i)] for i in range(0, 8)])
        for i, answers in enumerate(self.answers):
            # store the indices of anything that is answerable
            if count:
                answers = answers[number_indices]
            answer_has_index = len(answers.nonzero()) > 0
            if answer_has_index:
                answerable.append(i)
        return answerable

    def _encode_question(self, question):
        """ Turn a question into a vector of indices and a question length """
        vec = torch.zeros(self.max_question_length).long()
        for i, token in enumerate(question):
            index = self.token_to_index.get(token, 0)
            vec[i] = index
        return vec, len(question)

    def _encode_answers(self, answers):
        """ Turn an answer into a vector """
        # answer vec will be a vector of answer counts to determine which answers will contribute to the loss.
        # this should be multiplied with 0.1 * negative log-likelihoods that a model produces and then summed up
        # to get the loss that is weighted by how many humans gave that answer
        answer_vec = torch.zeros(len(self.answer_to_index))
        for answer in answers:
            index = self.answer_to_index.get(answer)
            if index is not None:
                answer_vec[index] += 1
        return answer_vec

    def _load_image(self, image_id):
        """ Load an image """
        if not hasattr(self, 'features_file'):
            # Loading the h5 file has to be done here and not in __init__ because when the DataLoader
            # forks for multiple works, every child would use the same file object and fail
            # Having multiple readers using different file objects is fine though, so we just init in here.
            self.features_file = h5py.File(self.image_features_path, 'r')
        index = self.coco_id_to_index[image_id]
        img = self.features_file['features'][index]
        boxes = self.features_file['boxes'][index]
        return torch.from_numpy(img).unsqueeze(1), torch.from_numpy(boxes)

    def _load_image_new(self, image_id):

        #index = self.coco_id_to_index[image_id]
        img_path = "/home/user/data/bottom_up_features_coco/features/" + str(image_id).zfill(12) + ".npy"
        boxes_path = "/home/user/data/bottom_up_features_coco/boxes/" + str(image_id).zfill(12) + ".npy"
        img = np.load(img_path)
        boxes = np.load(boxes_path)
        return torch.from_numpy(img).unsqueeze(1), torch.from_numpy(boxes)        



    def _load_item_demo(self, item, question):
        q, q_length = self._encode_question(prepare_question_demo(question))
        image_id = self.coco_ids[item]
        v, b = self._load_image(image_id) 
        v = v.view(2048,36,1)

        return v, q, b, item, q_length


    def __getitem__(self, item):
        if self.answerable_only:
            item = self.answerable[item]
        q, q_length = self._encode_question(prepare_question(self.questions["questions"][item]))
        #print(self.questions["questions"][item]["question"])
        #brak
        if not self.dummy_answers:
            a = self._encode_answers(prepare_answer(self.answers['annotations'][item]))
        else:
            # just return a dummy answer, it's not going to be used anyway
            a = 0
        image_id = self.coco_ids[item]

        v, b = self._load_image(image_id) #2048x1x36 V
       
        v = v.view(2048,36,1)
        
        
        # since batches are re-ordered for PackedSequence's, the original question order is lost
        # we return `item` so that the order of (v, q, a) triples can be restored if desired
        # without shuffling in the dataloader, these will be in the order that they appear in the q and a json's.
        
        return v, q, a, b, item, q_length

    def __len__(self):
        if self.answerable_only:
            return len(self.answerable)
        else:
            return 214354#len(self.questions)


# this is used for normalizing questions
_special_chars = re.compile('[^a-z0-9 ]*')

# these try to emulate the original normalization scheme for answers
_period_strip = re.compile(r'(?!<=\d)(\.)(?!\d)')
_comma_strip = re.compile(r'(\d)(,)(\d)')
_punctuation_chars = re.escape(r';/[]"{}()=+\_-><@`,?!')
_punctuation = re.compile(r'([{}])'.format(re.escape(_punctuation_chars)))
_punctuation_with_a_space = re.compile(r'(?<= )([{0}])|([{0}])(?= )'.format(_punctuation_chars))


def prepare_questions(questions_json):
    """ Tokenize and normalize questions from a given question json in the usual VQA format. """
    questions = [q['question'] for q in questions_json['questions']]
    for question in questions:
        question = question.lower()[:-1]
        question = _special_chars.sub('', question)
        yield question.split(' ')

def prepare_question(question):

    question = question['question'].lower()[:-1]
    question = _special_chars.sub('', question)
    return question.split(' ')

def prepare_question_demo(question):

    question = question.lower()[:-1]
    question = _special_chars.sub('', question)
    return question.split(' ')


def prepare_answers(answers_json):
    """ Normalize answers from a given answer json in the usual VQA format. """
    answers = [[a['answer'] for a in ans_dict['answers']] for ans_dict in answers_json['annotations']]
    # The only normalization that is applied to both machine generated answers as well as
    # ground truth answers is replacing most punctuation with space (see [0] and [1]).
    # Since potential machine generated answers are just taken from most common answers, applying the other
    # normalizations is not needed, assuming that the human answers are already normalized.
    # [0]: http://visualqa.org/evaluation.html
    # [1]: https://github.com/VT-vision-lab/VQA/blob/3849b1eae04a0ffd83f56ad6f70ebd0767e09e0f/PythonEvaluationTools/vqaEvaluation/vqaEval.py#L96

    def process_punctuation(s):
        # the original is somewhat broken, so things that look odd here might just be to mimic that behaviour
        # this version should be faster since we use re instead of repeated operations on str's
        if _punctuation.search(s) is None:
            return s
        s = _punctuation_with_a_space.sub('', s)
        if re.search(_comma_strip, s) is not None:
            s = s.replace(',', '')
        s = _punctuation.sub(' ', s)
        s = _period_strip.sub('', s)
        return s.strip()

    for answer_list in answers:
        yield list(map(process_punctuation, answer_list))

def prepare_answer(answer):

    answers = [a['answer'] for a in answer['answers']]
    
    return list(map(process_punctuation_custom, answers))





def process_punctuation_custom(s):
    # the original is somewhat broken, so things that look odd here might just be to mimic that behaviour
    # this version should be faster since we use re instead of repeated operations on str's
    if _punctuation.search(s) is None:
        return s
    s = _punctuation_with_a_space.sub('', s)
    if re.search(_comma_strip, s) is not None:
        s = s.replace(',', '')
    s = _punctuation.sub(' ', s)
    s = _period_strip.sub('', s)
    return s.strip()




class CocoImages(data.Dataset):
    """ Dataset for MSCOCO images located in a folder on the filesystem """
    def __init__(self, path, transform=None):
        super(CocoImages, self).__init__()
        self.path = path
        self.id_to_filename = self._find_images()
        self.sorted_ids = sorted(self.id_to_filename.keys())  # used for deterministic iteration order
        print('found {} images in {}'.format(len(self), self.path))
        self.transform = transform

    def _find_images(self):
        id_to_filename = {}
        for filename in os.listdir(self.path):
            if not filename.endswith('.jpg'):
                continue
            id_and_extension = filename.split('_')[-1]
            id = int(id_and_extension.split('.')[0])
            id_to_filename[id] = filename
        return id_to_filename

    def __getitem__(self, item):
        id = self.sorted_ids[item]
        path = os.path.join(self.path, self.id_to_filename[id])
        img = Image.open(path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)
        return id, img

    def __len__(self):
        return len(self.sorted_ids)