read_utils.py

import numpy as np
import pickle
import os


class Dataset(object):
    def __init__(self, lst_sub_faces, batch_size, num_frames, seed=66):
        self.lst_sub_faces = lst_sub_faces
        self.batch_size = batch_size
        self.num_frames = num_frames
        self.seed = seed
        np.random.seed(seed=self.seed)

        lst_sub_idx = []
        lst_sub_cursor = []
        for i in range(len(self.lst_sub_faces)):
            idx = np.random.permutation(np.size(lst_sub_faces[i], 0))
            lst_sub_idx.append(idx)
            lst_sub_cursor.append(0)

        self.lst_sub_idx = lst_sub_idx
        self.lst_sub_cursor = lst_sub_cursor

    def get_one_set(self, sub):
        sub_faces = self.lst_sub_faces[sub]
        sub_idx = self.lst_sub_idx[sub]

        if self.num_frames > sub_faces.shape[0]:
            idx = np.random.choice(sub_faces.shape[0], self.num_frames, replace=True)
        else:
            if self.lst_sub_cursor[sub] + self.num_frames > sub_faces.shape[0]:
                sub_idx = np.random.permutation(sub_faces.shape[0])
                self.lst_sub_idx[sub] = sub_idx
                self.lst_sub_cursor[sub] = 0

            idx = sub_idx[self.lst_sub_cursor[sub]: self.lst_sub_cursor[sub] + self.num_frames]
            self.lst_sub_cursor[sub] += self.num_frames

        return sub_faces[idx, :]

    def next_pair_batch(self, set_per_sub=3):
        sub_c = np.random.choice(len(self.lst_sub_faces), self.batch_size, replace=True)
        batch = []
        lst_len = []
        lst_label = []
        # set_per_sub = 3

        for s in sub_c:
            for i in range(0, set_per_sub):
                set = self.get_one_set(s)
                batch.append(set)
                lst_len.append(set.shape[0])
                lst_label.append(s)

        return batch, lst_label, lst_len

    def next_batch(self):
        sub_c = np.random.choice(len(self.lst_sub_faces), self.batch_size, replace=True)
        batch = []
        lst_len = []
        lst_label = []
        for s in sub_c:
            set = self.get_one_set(s)
            batch.append(set)
            lst_len.append(set.shape[0])
            lst_label.append(s)

        return batch, lst_label, lst_len


if __name__ == '__main__':

    num_sub = 100
    meta_dir = './data/Celeb/close'
    lst_sub_faces  = pickle.load(open(os.path.join(meta_dir, 'train_{}.bin'.format(num_sub)), 'rb'))
    dataset = Dataset(lst_sub_faces, 512, 20)
    while True:
        batch, lst_label, lst_len = dataset.next_batch()