Dataset

[flyyy]

I download the LSPet dataset from "http://sam.johnson.io/research/lspet.html", but the directory structure is not "COCO_train2014_000000000001.jpg". And the COCO 2014-All directory structure is not "099992483.jpg".

And you say "As the original InstaVariety is saved in tfrecord format, which is not suitable for use in Pytorch. You could run this script which will extract frames of every tfrecord and save them as jpeg."
I run this script but is very slow to save "jpg" images.

You use many dataset to train the model, but I am very hard to download so many dataset. So could you provide preprocessed data in a '.pt' file? Such as in TCMR "https://github.com/hongsukchoi/TCMR_RELEASE", which is easy to reproduce the model.

[ZiNiU WaN]

Sorry for the inconvenience. I have corrected the directory structure examples in data.md for COCO and LSPest. They are just typos and would not affect training.

Regarding the preprocessing of InstaVariety, since the size of InstaVariety is very large, empirically it could take DAYS to finish the preprocessing. Since the work was done during my internship, I am not able to access the preprocessed dataset right now. Sorry I can't help you on that now.

[Joe Lee]

Here is a revised code that is adapted to the latest Tensorflow and can greatly accelerate the preprocessing process. It only takes 2 hours to finish the whole process.

import os
import sys
sys.path.append('.')

import argparse
import numpy as np
import os.path as osp
from multiprocessing import Process, Pool
from glob import glob
from tqdm import tqdm
import tensorflow as tf
from PIL import Image

from lib.core.config import INSTA_DIR, INSTA_IMG_DIR

def process_single_record(args):
    fname, outdir, split = args
    record_name = fname.split('/')[-1]
    dataset = tf.data.TFRecordDataset([fname])
    for vid_idx, serialized_ex in enumerate(dataset):
        os.makedirs(osp.join(outdir, split, record_name, str(vid_idx)), exist_ok=True)
        example = tf.train.Example()
        example.ParseFromString(serialized_ex.numpy())

        N = int(example.features.feature['meta/N'].int64_list.value[0])
        images_data = example.features.feature['image/encoded'].bytes_list.value

        images = tf.stack([tf.image.decode_jpeg(image, channels=3) for image in images_data], axis=0)
        for i, image in enumerate(images):
            image = Image.fromarray(image.numpy())
            image.save(osp.join(outdir, split, record_name, str(vid_idx), str(i)+".jpg"))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--inp_dir', type=str, help='tfrecords file path', default=INSTA_DIR)
    parser.add_argument('--n', type=int, help='total num of workers', default=4)
    parser.add_argument('--i', type=int, help='current index of worker (from 0 to n-1)', default=0)
    parser.add_argument('--split', type=str, help='train or test', default='train')
    parser.add_argument('--out_dir', type=str, help='output images path', default=INSTA_IMG_DIR)
    args = parser.parse_args()

    fpaths = glob(f'./data/insta_variety/{args.split}/*.tfrecord')
    fpaths = sorted(fpaths)

    total = len(fpaths)
    fpaths = fpaths[args.i*total//args.n : (args.i+1)*total//args.n]

    os.makedirs(args.out_dir, exist_ok=True)

    with Pool(processes=args.n) as pool:
        with tqdm(total=len(fpaths), unit='file', desc=f'Worker {args.i}/{args.n}') as pbar:
            for _ in pool.imap_unordered(process_single_record, [(fname, args.out_dir, args.split) for fname in fpaths]):
                pbar.update(1)