from __future__ import print_function
import torch.utils.data as data
import os
import sys
import numpy as np
import h5py

class MLMLoader(data.Dataset):
    def __init__(self, data_path, partition, mismatch=0.5):

        if data_path == None:
            raise Exception('No data path specified.')

        if partition is None or partition not in ['train', 'val', 'test']:
            raise Exception('Unknown partition type %s.' % partition)
        else:
            self.partition = partition

        self.h5f = h5py.File(os.path.join(data_path, f'{partition}.h5'), 'r')

        self.ids = self.h5f['ids']
        self.mismatch = mismatch

        self.smr_counter = 0

    def __getitem__(self, index):
        instanceId = self.ids[index]
        # we force a mismatch give the probability
        match = np.random.uniform() > self.mismatch if self.partition == 'train' else True

        target = match and 1 or -1

        if target == 1:  # load positive samples
            all_img = self.h5f[f'{instanceId}_images'][()]
        else:
            # Negative samples are generated by picking random images
            all_idx = range(len(self.ids))
            rndindex = np.random.choice(all_idx)
            # random index to pick image ids at random
            while rndindex == index:
                rndindex = np.random.choice(all_idx)  # pick a random index

            # load negative samples of images
            rndId = self.ids[rndindex]
            all_img = self.h5f[f'{rndId}_images']

        # Other modalities remain unchanged
        all_smr = self.h5f[f'{instanceId}_summaries'][()]
        all_cls = self.h5f[f'{instanceId}_classes'][()]
        cell_id = self.h5f[f'{instanceId}_onehot'][()]

        if self.partition == 'train':
            # for training we sample random image, summary and classes for the given id
            image = all_img[np.random.choice(range(all_img.shape[0]))]
            summary = all_smr[np.random.choice(range(all_smr.shape[0]))]
            classes = all_cls[np.random.choice(range(all_cls.shape[0]))]
        else:
            # for test and validation select the first example for images and triples
            # for summaries we set a counter since we have 3 languages (en, de, fr)
            image = all_img[0]
            summary = all_smr[self.smr_counter] if self.smr_counter < all_smr.shape[0] else all_smr[0]
            classes = all_cls[0]

            # update summary counter
            self.smr_counter = 0 if self.smr_counter == 2 else self.smr_counter + 1

        return {
            'id': instanceId,
            'image': image,
            'summary': summary,
            'classes': classes,
            'target_ir': target,
            'target_le': np.argmax(cell_id)
        }

    def __len__(self):
        return len(self.ids)
from __future__ import print_function
import torch.utils.data as data
import os
import sys
import numpy as np
import h5py

class MLMLoader(data.Dataset):
    def __init__(self, data_path, partition, mismatch=0.5):

        if data_path == None:
            raise Exception('No data path specified.')

        if partition is None or partition not in ['train', 'val', 'test']:
            raise Exception('Unknown partition type %s.' % partition)
        else:
            self.partition = partition

        self.h5f = h5py.File(os.path.join(data_path, f'{partition}.h5'), 'r')

        self.ids = self.h5f['ids']
        self.mismatch = mismatch

        self.smr_counter = 0

    def __getitem__(self, index):
        instanceId = self.ids[index]
        # we force a mismatch give the probability
        match = np.random.uniform() > self.mismatch if self.partition == 'train' else True

        target = match and 1 or -1

        if target == 1:  # load positive samples
            all_img = self.h5f[f'{instanceId}_images'][()]
        else:
            # Negative samples are generated by picking random images
            all_idx = range(len(self.ids))
            rndindex = np.random.choice(all_idx)
            # random index to pick image ids at random
            while rndindex == index:
                rndindex = np.random.choice(all_idx)  # pick a random index

            # load negative samples of images
            rndId = self.ids[rndindex]
            all_img = self.h5f[f'{rndId}_images']

        # Other modalities remain unchanged
        all_smr = self.h5f[f'{instanceId}_summaries'][()]
        all_cls = self.h5f[f'{instanceId}_classes'][()]
        cell_id = self.h5f[f'{instanceId}_onehot'][()]

        if self.partition == 'train':
            # for training we sample random image, summary and classes for the given id
            image = all_img[np.random.choice(range(all_img.shape[0]))]
            summary = all_smr[np.random.choice(range(all_smr.shape[0]))]
            classes = all_cls[np.random.choice(range(all_cls.shape[0]))]
        else:
            # for test and validation select the first example for images and triples
            # for summaries we set a counter since we have 3 languages (en, de, fr)
            image = all_img[0]
            summary = all_smr[self.smr_counter] if self.smr_counter < all_smr.shape[0] else all_smr[0]
            classes = all_cls[0]

            # update summary counter
            self.smr_counter = 0 if self.smr_counter == 2 else self.smr_counter + 1

        return {
            'id': instanceId,
            'image': image,
            'summary': summary,
            'classes': classes,
            'target_ir': target,
            'target_le': np.argmax(cell_id)
        }

    def __len__(self):
        return len(self.ids)
# we force a mismatch give the probability # load positive samples # Negative samples are generated by picking random images # random index to pick image ids at random # pick a random index # load negative samples of images # Other modalities remain unchanged # for training we sample random image, summary and classes for the given id # for test and validation select the first example for images and triples # for summaries we set a counter since we have 3 languages (en, de, fr) # update summary counter