Commit e699f94e authored by Gilberto Astolfi's avatar Gilberto Astolfi
Browse files

Adicionando implementação para Yolo

parent c77e1aed
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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment