EAST Helper Scripts

msra_to_east.py

# This script converts msra ground truth format to east training format
# Download and extract MSRA, cd <MSRA_DIR>, mkdir out and run this script
# https://github.com/argman/EAST

import os
import math
import shutil
import numpy as np

rotate = lambda xy, theta: [xy[0] * math.cos(theta) - xy[1] * math.sin(theta),xy[0] * math.sin(theta) + xy[1] * math.cos(theta)]
translate = lambda xy, offset: (xy[0] + offset[0], xy[1] + offset[1])
rotate_about = lambda xy, theta, origin: translate(rotate(translate(xy, -origin), theta), origin)
rect2pts = lambda r: [[r[0], r[1]], [r[0]+r[2],r[1]], [r[0]+r[2],r[1]+r[3]], [r[0],r[1]+r[3]]]
center = lambda r: np.array([r[0] + r[2]/2, r[1] + r[3]/2])
rotated_poly = lambda r,t: np.array([rotate_about(pt,t,center(r)) for pt in rect2pts(r)], dtype=int)
poly_to_str = lambda poly: (str(poly[0][0]) + ',' + str(poly[0][1]) + ',' + 
														str(poly[1][0]) + ',' + str(poly[1][1]) + ',' +
														str(poly[2][0]) + ',' + str(poly[2][1]) + ',' +
														str(poly[3][0]) + ',' + str(poly[3][1]))

i = 0
for fn in os.listdir('train/'):
		if fn.endswith('.gt'):
				continue
		print(fn)
		img_name = 'train/' + fn
		gt_name = img_name.replace('.JPG', '.gt')
		shutil.copy(img_name, 'out/' + str(i) + '.jpg')
		lines = None
		with open(gt_name) as f:
				lines = [ln.rstrip() for ln in f.readlines()]
		text = ''
		for ln in lines:
				if ln == '':
						continue
				fields = ln.split(' ')
				rect = [int(fields[2]),int(fields[3]),int(fields[4]),int(fields[5])]
				thet = float(fields[6])
				text += poly_to_str(rotated_poly(rect, thet)) + ',\n'
		with open('out/' + str(i) + '.txt', 'w') as f:
				f.write(text)
		i += 1