wehub-resource-sync
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
import os
|
|
from data import SROIETask2
|
|
from tqdm import tqdm
|
|
import shutil
|
|
import zipfile
|
|
|
|
if __name__ == '__main__':
|
|
test_dir = '../SROIE_Task2_Original/test'
|
|
output_dir = 'temp'
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
generate_txt_path = '../generate-test.txt'
|
|
output_file = None
|
|
output_fp = None
|
|
|
|
with open(generate_txt_path, 'r', encoding='utf8') as fp:
|
|
lines = list(fp.readlines())
|
|
while not lines[0].startswith('T-0'):
|
|
lines = lines[1:]
|
|
|
|
_, data = SROIETask2(test_dir, None, None)
|
|
for t in tqdm(data):
|
|
file_name = t['file_name']
|
|
image_id = int(t['image_id'])
|
|
|
|
this_output_file = os.path.basename(file_name).replace('.jpg', '.txt')
|
|
if this_output_file != output_file:
|
|
if output_fp is not None:
|
|
output_fp.close()
|
|
output_file = this_output_file
|
|
output_fp = open(os.path.join(output_dir, output_file), 'w', encoding='utf8')
|
|
|
|
pred_line_id = image_id * 4 + 2
|
|
pred_line = lines[pred_line_id]
|
|
assert pred_line.startswith('D-{:d}'.format(image_id))
|
|
pred_line = pred_line[pred_line.find('\t') + 1:]
|
|
pred_str = pred_line[pred_line.find('\t') + 1:]
|
|
|
|
for word in pred_str.split():
|
|
output_fp.write(word + '\n')
|
|
|
|
if output_fp:
|
|
output_fp.close()
|
|
|
|
zip_fp = zipfile.ZipFile('predictions.zip', 'w')
|
|
for txt_file in os.listdir(output_dir):
|
|
zip_fp.write(os.path.join(output_dir, txt_file), txt_file)
|
|
zip_fp.close()
|
|
shutil.rmtree(output_dir)
|
|
|