Spaces:
Sleeping
Sleeping
File size: 1,577 Bytes
6fc683c |
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 |
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)
|