Spaces:
Build error
Build error
File size: 1,544 Bytes
3d5e231 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import os
import csv
import numpy as np
img_folder = '/playpen-ssd/adyasha/projects/StoryGAN/pororo_png/'
def get_captions_by_split():
video_len = 4
descriptions_original = np.load(os.path.join(img_folder, 'descriptions.npy'), allow_pickle=True,
encoding='latin1').item()
followings = np.load(os.path.join(img_folder, 'following_cache4.npy'))
train_ids, val_ids, test_ids = np.load(os.path.join(img_folder, 'train_seen_unseen_ids.npy'), allow_pickle=True)
filenames = ['descriptions_train.csv', 'descriptions_val.csv', 'descriptions_test.csv']
for ids, filename in zip([train_ids, val_ids, test_ids], filenames):
im_ids = []
for src_img_id in ids:
tgt_img_paths = [str(followings[src_img_id][i])[2:-1] for i in range(video_len)]
tgt_img_ids = [str(tgt_img_path).replace(img_folder, '').replace('.png', '') for tgt_img_path in
tgt_img_paths]
im_ids.extend(tgt_img_ids)
# captions = [descriptions_original[tgt_img_id] for tgt_img_id in tgt_img_ids]
im_ids = list(set(im_ids))
im_ids.sort()
# captions = [descriptions_original[i] for i in im_ids]
with open(os.path.join(img_folder, filename), 'w') as csvfile:
# creating a csv writer object
csvwriter = csv.writer(csvfile)
for i in im_ids:
csvwriter.writerow([i, descriptions_original[i][0]])
get_captions_by_split() |