data_script/vist_hdf5.py

import argparse
import json
import os

import cv2
import h5py
import numpy as np
from PIL import Image
from tqdm import tqdm


def main(args):
    train_data = json.load(open(os.path.join(args.sis_json_dir, 'train.story-in-sequence.json')))
    val_data = json.load(open(os.path.join(args.sis_json_dir, 'val.story-in-sequence.json')))
    test_data = json.load(open(os.path.join(args.sis_json_dir, 'test.story-in-sequence.json')))

    prefix = ["train", "val", "test"]
    whole_album = {}
    for i, data in enumerate([train_data, val_data, test_data]):
        album_mapping = {}
        for annot_new in data["annotations"]:
            annot = annot_new[0]
            assert len(annot_new) == 1
            if annot['story_id'] not in album_mapping:
                album_mapping[annot['story_id']] = {"flickr_id": [annot['photo_flickr_id']],
                                                    "sis": [annot['original_text']],
                                                    "length": 1}
            else:
                album_mapping[annot['story_id']]["flickr_id"].append(annot['photo_flickr_id'])
                album_mapping[annot['story_id']]["sis"].append(
                    annot['original_text'])
                album_mapping[annot['story_id']]["length"] += 1
        whole_album[prefix[i]] = album_mapping

    for p in prefix:
        deletables = []
        for story_id, story in whole_album[p].items():
            if story['length'] != 5:
                print("deleting {}".format(story_id))
                deletables.append(story_id)
                continue
            d = [os.path.exists(os.path.join(args.img_dir, "{}.jpg".format(_))) for _ in story["flickr_id"]]
            if sum(d) < 5:
                print("deleting {}".format(story_id))
                deletables.append(story_id)
            else:
                pass
        for i in deletables:
            del whole_album[p][i]

    train_data = json.load(open(os.path.join(args.sis_json_dir, 'train.description-in-isolation.json')))
    val_data = json.load(open(os.path.join(args.sis_json_dir, 'val.description-in-isolation.json')))
    test_data = json.load(open(os.path.join(args.sis_json_dir, 'test.description-in-isolation.json')))

    flickr_id2text = {}
    for i, data in enumerate([train_data, val_data, test_data]):
        for l in data['annotations']:
            assert len(l) == 1
            if l[0]['photo_flickr_id'] in flickr_id2text:
                flickr_id2text[l[0]['photo_flickr_id']] = \
                    max([flickr_id2text[l[0]['photo_flickr_id']], l[0]['original_text']], key=len)
            else:
                flickr_id2text[l[0]['photo_flickr_id']] = l[0]['original_text']

    for p in prefix:
        deletables = []
        for story_id, story in whole_album[p].items():
            story['dii'] = []
            for i, flickr_id in enumerate(story['flickr_id']):
                if flickr_id not in flickr_id2text:
                    print("{} not found in story {}".format(flickr_id, story_id))
                    deletables.append(story_id)
                    break
                story['dii'].append(flickr_id2text[flickr_id])
        for i in deletables:
            del whole_album[p][i]

    f = h5py.File(args.save_path, "w")
    for p in prefix:
        group = f.create_group(p)
        story_dict = whole_album[p]
        length = len(story_dict)
        images = list()
        for i in range(5):
            images.append(
                group.create_dataset('image{}'.format(i), (length,), dtype=h5py.vlen_dtype(np.dtype('uint8'))))
        sis = group.create_dataset('sis', (length,), dtype=h5py.string_dtype(encoding='utf-8'))
        dii = group.create_dataset('dii', (length,), dtype=h5py.string_dtype(encoding='utf-8'))
        for i, (story_id, story) in enumerate(tqdm(story_dict.items(), leave=True, desc="saveh5")):
            imgs = [Image.open('{}/{}.jpg'.format(args.img_dir, flickr_id)).convert('RGB') for flickr_id in
                    story['flickr_id']]
            for j, img in enumerate(imgs):
                img = np.array(img).astype(np.uint8)
                img = cv2.imencode('.png', img)[1].tobytes()
                img = np.frombuffer(img, np.uint8)
                images[j][i] = img
            sis[i] = '|'.join([t.replace('\n', '').replace('\t', '').strip() for t in story['sis']])
            txt_dii = [t.replace('\n', '').replace('\t', '').strip() for t in story['dii']]
            txt_dii = sorted(set(txt_dii), key=txt_dii.index)
            dii[i] = '|'.join(txt_dii)
    f.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='arguments for vist hdf5 file saving')
    parser.add_argument('--sis_json_dir', type=str, required=True, help='sis json file directory')
    parser.add_argument('--dii_json_dir', type=str, required=True, help='dii json file directory')
    parser.add_argument('--img_dir', type=str, required=True, help='json file directory')
    parser.add_argument('--save_path', type=str, required=True, help='path to save hdf5')
    args = parser.parse_args()
    main(args)