import json
import random
from pathlib import Path
import numpy as np
import scipy.misc
from moviepy.editor import VideoFileClip
from tqdm import tqdm


class WindowNpyGenerator(object):
    def __init__(self,
                 n_train=None,
                 n_val=None,
                 fps=1,
                 timesteps=5,
                 overlap=4,
                 target_dir=None):
        self.video_dirs = []
        self.n_train = n_train or 100
        self.n_val = n_val or 20
        self.fps = fps
        self.timesteps = timesteps
        self.overlap = overlap
        self.target_dir = target_dir or Path('./npz/')

    def extract_windows(self, video_dir):
        video = VideoFileClip(str(video_dir / 'video.mp4'))
        info = json.load((video_dir / 'info.json').open())
        n_frames = int(video.duration) * self.fps
        timesteps = self.timesteps
        overlap = self.overlap

        label = np.zeros(n_frames, dtype=np.uint8)
        for s, e in zip(info['starts'], info['ends']):
            fs = round(s * self.fps)
            fe = round(e * self.fps)
            label[fs:fe + 1] = 1

        windows = [(video, f - timesteps, f, label[f - 1])
                   for f in range(timesteps, n_frames, timesteps - overlap)]
        return windows

    def gen_npz(self, windows, chunk_size=10000):
        for idx in tqdm(range(0, len(windows), chunk_size), desc='chunk'):
            chunk_s, chunk_e = idx, min(idx + chunk_size, len(windows))
            chunk = windows[chunk_s:chunk_e]

            n = len(chunk)
            xs = np.zeros((n, self.timesteps, 224, 224, 3), dtype=np.float32)
            ys = np.zeros(n, dtype=np.uint8)
            for i, (video, s, e, y) in enumerate(tqdm(chunk, desc='data')):
                for j in range(e - s):
                    img = video.get_frame((s + j) / self.fps)
                    xs[i][j] = scipy.misc.imresize(img, (224, 224))
                ys[i] = y

            npz_path = self.target_dir / '{:04d}.npz'.format(idx // chunk_size)
            np.savez(npz_path, xs=xs, ys=ys)

    def fit(self, video_dirs):
        train, val = [], []
        for video_dir in video_dirs:
            windows = self.extract_windows(video_dir)
            random.shuffle(windows)
            pivot = round(
                (self.n_train) / (self.n_train + self.n_val) * len(windows))
            train.extend(windows[:pivot])
            val.extend(windows[pivot:])

        if len(train) < n_train:
            print('Not enough train data: {} < {}'.format(len(train), n_train))
            return
        if len(val) < n_val:
            print('Not enough val data: {} < {}'.format(len(val), n_val))
            return

        train = random.sample(train, k=self.n_train)
        val = random.sample(val, k=self.n_val)

        self.target_dir.mkdir(exist_ok=True)
        self.gen_npz(train)
        self.gen_npz(val)

    def flow(self, batch_size=80):
        npzs = sorted(self.target_dir.glob('*.npz'))

        idx = 0
        x_batch = np.zeros(
            (batch_size, self.timesteps, 224, 224, 3), dtype=np.float32)
        y_batch = np.zeros((batch_size, 1), dtype=np.uint8)

        while True:
            for npz in npzs:
                for x, y in zip(npz['xs'], npz['ys']):
                    x_batch[idx] = x
                    y_batch[idx] = y
                    if idx + 1 == batch_size:
                        yield x_batch, y_batch
                    idx = (idx + 1) % batch_size


def main():
    dataset = Path('~/tthl-dataset/').expanduser()
    video_dirs = sorted(dataset.glob('video*/'))

    gen = WindowNpyGenerator(n_train=10000, n_val=2000, fps=3, timesteps=2, overlap=1)
    gen.fit(video_dirs)


if __name__ == '__main__':
    main()
@@ -2441,16 +2441,60 @@
 vot:%5D)%0A%0A
+            print(video_dir, len(windows))%0A%0A
         
@@ -2562,46 +2562,9 @@
 data
-: %7B%7D %3C %7B%7D'.format(len(train), n_train)
+'
 )%0A  
@@ -2651,42 +2651,9 @@
 data
-: %7B%7D %3C %7B%7D'.format(len(val), n_val)
+'
 )%0A  

Update at 2017-07-21 16-29-13