Adicionando implementação para Yolo

yolo_utils.py

+import xml.etree.ElementTree as et
+import os
+from shutil import copyfile
+
+class YoloUtils:
+	def __init__(self):
+		pass
+
+	def createObjectNames(self, dictClasses, file_name):
+		text_out = ''
+		for idx, class_name in dictClasses.items():
+			text_out += class_name + '\n'
+		
+		f = open(file_name, 'w')
+		f.write(text_out)
+		f.close()
+		
+
+	def estruturaBlocos(self, classes):
+		item = {}
+		for block in classes:
+			bls = block.split('\n')
+			variable = []
+			for idx in range(1, len(bls)-1):			
+				for pd in bls[idx].split(':'):
+					s = pd.replace('"',"").replace(" ","")
+					variable.append(s)
+			item[int(variable[1])] = variable[3]
+		return item
+
+	def protoToDict(self, file_name):
+		file = open(file_name, 'r')
+		lines = file.readlines()
+		idx = 0
+		initial_idx = -1
+		final_idx = -1
+		classes = []
+		for line in lines:	
+			if '{' in line:
+				initial_idx = idx
+			if '}' in line:
+				final_idx = idx		
+			if initial_idx >= 0 and final_idx >= 0:
+				block = ''.join(lines[initial_idx:final_idx+1])		
+				initial_idx = block.index('{')
+				final_idx = block.index('}')
+				block = block[initial_idx:final_idx+1]
+				classes.append(block)
+				initial_idx = -1
+				final_idx = -1		
+			idx += 1
+		classes = self.estruturaBlocos(classes)
+		return classes
+
+	def createFileSetYolo(self, path_set, path_yolo_txt):
+		text_out = ''
+		for pth in os.listdir(path_set):
+			path_class = path_set + '/' + pth
+			if os.path.isdir(path_class):
+				for file_in_class in os.listdir(path_class):
+					name, extension = os.path.splitext(file_in_class)
+					if extension.lower() == '.xml'.lower():
+						file_name = path_class + '/' + file_in_class
+						tree = et.parse(file_name)
+						path_img = tree.find("path").text
+						text_out += path_img +'\n'
+		f = open(path_yolo_txt, 'w')
+		f.write(text_out)
+		f.close()
+
+	def createFilesYolo(self, repositorio):
+		for fold in os.listdir(repositorio):
+			fold_src = repositorio + '/' + fold
+			for set_ in os.listdir(fold_src):
+				if set_ != 'test':
+					path_pbtxt = fold_src + '/' + set_+'/' + fold +'_'+set_+'.pbtxt'
+					dictClasses = self.protoToDict(path_pbtxt)
+					obj_name_file = fold_src + '/' + set_+'/' + fold +'_'+set_+'.names'
+					self.createObjectNames(dictClasses, obj_name_file)
+					path_set = fold_src + '/' + set_
+					path_yolo_txt = fold_src + '/' + set_+'/' + fold +'_'+set_+'.txt'
+					self.createFileSetYolo(path_set, path_yolo_txt)
+
+	def searchSourceFile(self,path_search, file_img):
+		for cl in os.listdir(path_search):
+			path_cl = path_search + '/' + cl
+			if os.path.isdir(path_cl):
+				for f in os.listdir(path_cl):
+					if f == file_img:
+						path_dst = path_cl
+						return path_dst
+
+	def sendYoloAnotationRepository(self,output_dir, path_search):
+		for f in os.listdir(output_dir):
+			name, extension = os.path.splitext(f)
+			if extension.lower() == '.txt'.lower():
+				path_src = output_dir +'/'+f
+				file_img = name +'.jpg'
+				dst = searchSourceFile(path_search, file_img)
+				path_dst = dst + '/' + name +'.txt'
+				copyfile(path_src, path_dst)
+
+	def executeScript(self, path_driver, input_csv, input_names, output_dir):
+		parametros = input_csv + ' ' + input_names + ' ' + output_dir
+		os.system('python3 ' + path_driver+ '/generate_yolo_txt.py ' + parametros)
+
+	def createAnnotationYolo(self, path_driver_yolo, repositorio, path_tfrecords):
+		for fold in os.listdir(repositorio):
+			fold_src = repositorio +'/'+fold
+			for set_ in os.listdir(fold_src):
+				if set_ != 'test':
+					input_names = fold_src + '/' + set_+'/' + fold +'_'+set_+'.names'
+					input_csv = path_tfrecords +'/'+fold + '/' + set_+'/' + fold +'_'+set_+'.csv'
+					output_dir = path_tfrecords +'/'+fold + '/' + set_+'/Annotations'
+					self.executeScript(path_driver_yolo, input_csv, input_names, output_dir)
+					path_search = fold_src + '/' + set_ #buscar image para associar ao txt
+					self.sendYoloAnotationRepository(output_dir, path_search)
\ No newline at end of file