Annotation_Helper/xml_helper.py at master · git-hamza/Annotation_Helper · 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
import xml.etree.ElementTree as ET
import os
import argparse

def list_all_labels(IN_ANN,Count_Flag=False):
	##Shows you all the unique labels in xmls
	Labels = {}
	for xml_file in os.listdir(IN_ANN):
		tree = ET.parse(os.path.join(IN_ANN,xml_file))
		root = tree.getroot()
		bndbox = root.findall('object')
		for bb in bndbox:
			txt = bb[0].text
			if len(Labels) ==0:
				Labels[txt] = 1
			else:
				if txt in Labels.keys():
					Labels[txt] = Labels[txt] + 1
					pass
				else:
					Labels[txt] = 1
	if Count_Flag == False:
		print("Number of Labels : ", (len(Labels)))
		print("####__All Labels__######")
		for label in Labels.keys():
			print(label)
	else:
		return Labels

def count_of_each_label(IN_ANN):
	##Shows you thr number of instances of each unique labels
	Labels = list_all_labels(IN_ANN,Count_Flag=True)
	Labels = sorted(Labels.items(), key=lambda kv: kv[1])
	Labels.reverse()
	for label in Labels:
		print(label)

def change_label_name(IN_ANN,before='',after=''):
	##You can change labels name in all xmls
	for xml_file in os.listdir(IN_ANN):
		tree = ET.parse(os.path.join(IN_ANN,xml_file))
		root = tree.getroot()
		bndbox = root.findall('object')
		for bb in bndbox:
			txt = bb[0].text
			if txt == before:
				bb[0].text = after
		tree.write(os.path.join(IN_ANN,xml_file))
	print("Done")

def bbox_exceed_image_shape(IN_ANN):
	##tackle the error caused by bbox excedding image boundaries
	for xml_file in os.listdir(IN_ANN):
		tree = ET.parse(os.path.join(IN_ANN,xml_file))
		root = tree.getroot()
		filename = root[1].text
		bndbox = root.findall('object')
		siize = root.findall('size')
		wd = int(siize[0][0].text)
		ht = int(siize[0][1].text)
		for bb in bndbox:
			if int(bb[4][0].text) < 0:  ##xmin
				print(filename)
				print("Before : " + str(bb[4][0].text))
				bb[4][0].text = '0'
				print("After : " + str(bb[4][0].text))
			if int(bb[4][1].text) < 0:   ##ymin
				print(filename)
				print("Before : " +str(bb[4][1].text))
				bb[4][1].text = '0'
				print("After : " + str(bb[4][1].text))
			if int(bb[4][2].text) > wd:
				print(filename)
				print("Before : "+str(bb[4][2].text))
				bb[4][2].text = str(wd)
				print("After : " +str(bb[4][2].text))
			if int(bb[4][3].text) > ht:
				print(filename)
				print("Before : "+str(bb[4][3].text))
				bb[4][3].text = str(ht)
				print("After : "+str(bb[4][3].text))

		tree.write(os.path.join(IN_ANN,xml_file))
	print("Done")

def main():
	# change_label_name("Data/annotations/","Gadi","Car")
	# list_all_labels("Data/annotations/")
	# count_of_each_label("Data/annotations/")
	bbox_exceed_image_shape("Data/annotations/")

if __name__ == '__main__':
	main()