|
|
|
|
|
import os |
|
|
import sys |
|
|
import torch |
|
|
sys.path.append('.') |
|
|
|
|
|
import argparse |
|
|
parser = argparse.ArgumentParser(description='Data preparation') |
|
|
parser.add_argument('--dataset', type=str, choices=['refcoco', 'refcoco+','refcocog_google', 'refcocog_umd'], default='refcoco') |
|
|
args = parser.parse_args() |
|
|
|
|
|
def main(args): |
|
|
dataset = args.dataset |
|
|
input_txt_list = os.listdir(f'../ln_data/anns/{dataset}') |
|
|
if not os.path.exists(f'../data/{dataset}'): |
|
|
os.makedirs(f'../data/{dataset}') |
|
|
for input_txt in input_txt_list: |
|
|
split = input_txt.split('_')[-1].split('.')[0] |
|
|
input_txt = os.path.join('../ln_data/anns', dataset, input_txt) |
|
|
res = [] |
|
|
with open(input_txt, encoding='utf-8') as f: |
|
|
lines = f.readlines() |
|
|
for line in lines: |
|
|
line = line.split() |
|
|
stop = len(line) |
|
|
img_name = line[0] |
|
|
for i in range(1,len(line)): |
|
|
if (line[i]=='~'): |
|
|
stop=i |
|
|
break |
|
|
box_ = [list(map(int,box.split(','))) for box in line[1:stop]] |
|
|
box = box_[0][:4] |
|
|
seg_id=box_[0][-1] |
|
|
|
|
|
sent_stop=stop+1 |
|
|
for i in range(stop+1,len(line)): |
|
|
if line[i]=='~': |
|
|
des = '' |
|
|
for word in line[sent_stop:i]: |
|
|
des = des + word + ' ' |
|
|
sent_stop=i+1 |
|
|
des = des.rstrip(' ') |
|
|
res.append((img_name, seg_id, box, des)) |
|
|
des = '' |
|
|
for word in line[sent_stop:len(line)]: |
|
|
des = des + word + ' ' |
|
|
des = des.rstrip(' ') |
|
|
res.append((img_name, seg_id, box, des)) |
|
|
|
|
|
|
|
|
imgset_path = '{0}_{1}.pth'.format(dataset, split) |
|
|
images = torch.save(res, os.path.join("../data", dataset, imgset_path)) |
|
|
print(dataset, " done") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
main(args) |
|
|
|
