bitmap2svg/bitmap_gen_from_caption.py at master · awesome-davian/bitmap2svg · GitHub

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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import os
from xml.dom import minidom
import argparse
import cairosvg

def parse_caption(cap_arr, doc):

	color_list = ['red', 'orange', 'yellow', 'lime', 'green', 'spring_green', 'cyan',
	              'skyblue','blue', 'purple', 'pink', 'deep_pink']

	if cap_arr[0] == 'circle':
		polygon = doc.createElement('circle')
		polygon.setAttribute('class', 'circle')
		polygon.setAttribute('cx', cap_arr[1])
		polygon.setAttribute('cy', cap_arr[2])
		polygon.setAttribute('r', cap_arr[3])

		#change color to hsl
		hsl = color_list.index(cap_arr[4])
		hsl = str(int(hsl) * 30 )
		style = "fill: hsl(" + hsl + ",100%,50%);"
		polygon.setAttribute('style', style)

		cap_arr = cap_arr[5:]

	elif cap_arr[0] == 'rect':
		polygon = doc.createElement('rect')
		polygon.setAttribute('class', 'rect')
		polygon.setAttribute('x', cap_arr[1])
		polygon.setAttribute('y', cap_arr[2])
		polygon.setAttribute('width', cap_arr[3])
		polygon.setAttribute('height', cap_arr[4])

		#change color to hsl
		hsl = color_list.index(cap_arr[5])
		hsl = str(int(hsl) * 30 )
		style = "fill: hsl(" + hsl + ",100%,50%);"
		polygon.setAttribute('style', style)

		cap_arr = cap_arr[6:]

	return polygon, cap_arr


def main(args):

	root_path = args.root_path
	cap_path = root_path + 'caption/'
	svg_from_trg = root_path + 'gen/svg_from_caption/'
	bitmap_from_trg = root_path + 'gen/bitmap_from_caption/'
	if not os.path.exists(svg_from_trg):
		os.makedirs(svg_from_trg)
	if not os.path.exists(bitmap_from_trg):
		os.makedirs(bitmap_from_trg)

	file_list = os.listdir(cap_path)
	cnt = 0

	for f_list in file_list:
		if cnt >2:
			break;
		#cnt+=1
		#read caption
		with open(os.path.join(cap_path, f_list), 'r') as f:
			caption = f.read()

		#make svg
		doc = minidom.Document()
		svg = doc.createElement('svg')
		svg.setAttribute("xmlns", "http://www.w3.org/2000")
		svg.setAttribute("width", "500")
		svg.setAttribute("height", "500")
		doc.appendChild(svg)

		#parse caption
		cap_arr = caption.split(" ")
		while len(cap_arr) != 0:
			polygon, cap_arr = parse_caption(cap_arr, doc)
			svg.appendChild(polygon)

        #write svg
		with open(os.path.join(svg_from_trg, f_list), 'w+') as f:
			f.write(doc.toxml())

        #convert and save as bitmap
		svg_path = svg_from_trg + f_list
		bitmap_path = bitmap_from_trg + f_list.replace('.svg', '.png')
		cairosvg.svg2png(url=svg_path, write_to=bitmap_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--caption_path', type=str,
                        default='data/nobject/caption',
                        help='path for train annotation file')

    parser.add_argument('--svg_path', type=str,
                        default='data/nobject/svg',
                        help='ath for train annotation file')

    parser.add_argument('--root_path', type=str,
                        default='data/nobject_test/',
                        help='root path for train annotation file')

    args = parser.parse_args()
    main(args)