Spaces:
Sleeping
Sleeping
File size: 2,464 Bytes
087df0e |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import os
import argparse
import torch
try:
import ruamel_yaml as yaml
except ModuleNotFoundError:
import ruamel.yaml as yaml
from model.prismer_caption import PrismerCaption
from dataset import create_dataset, create_loader
from tqdm import tqdm
parser = argparse.ArgumentParser()
parser.add_argument('--mode', default='')
parser.add_argument('--port', default='')
parser.add_argument('--exp_name', default='', type=str)
args = parser.parse_args()
# load config
config = yaml.load(open('configs/caption.yaml', 'r'), Loader=yaml.Loader)['demo']
# generate expert labels
if len(config['experts']) > 0:
script_name = f'python experts/generate_depth.py'
os.system(script_name)
print('***** Generated Depth *****')
script_name = f'python experts/generate_edge.py'
os.system(script_name)
print('***** Generated Edge *****')
script_name = f'python experts/generate_normal.py'
os.system(script_name)
print('***** Generated Surface Normals *****')
script_name = f'python experts/generate_objdet.py'
os.system(script_name)
print('***** Generated Object Detection Labels *****')
script_name = f'python experts/generate_ocrdet.py'
os.system(script_name)
print('***** Generated OCR Detection Labels *****')
script_name = f'python experts/generate_segmentation.py'
os.system(script_name)
print('***** Generated Segmentation Labels *****')
# load datasets
_, test_dataset = create_dataset('caption', config)
test_loader = create_loader(test_dataset, batch_size=1, num_workers=4, train=False)
# load pre-trained model
model = PrismerCaption(config)
state_dict = torch.load(f'logging/caption_{args.exp_name}/pytorch_model.bin', map_location='cuda:0')
model.load_state_dict(state_dict)
tokenizer = model.tokenizer
# inference
model.eval()
with torch.no_grad():
for step, (experts, data_ids) in enumerate(tqdm(test_loader)):
captions = model(experts, train=False, prefix=config['prefix'])
captions = tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
caption = captions.to(experts['rgb'].device)[0]
caption = tokenizer.decode(caption, skip_special_tokens=True)
caption = caption.capitalize() + '.'
# save caption
save_path = test_loader.dataset.data_list[data_ids[0]]['image'].replace('jpg', 'txt')
with open(save_path, 'w') as f:
f.write(caption)
print('All Done.')
|