YOLO 训练数据准备

namesFile = "TrafficSign.names"
f = open(namesFile)

nameDic = {}

iter_f = iter(f)
classId = 0
for line in iter_f:
    tmp = line.split("\n")[0]
    nameDic[tmp] = classId
    classId = classId + 1

print(nameDic)



csvFile = "train-new.csv"

f = open(csvFile)

iter_f = iter(f)

# 注意，这个写死是为了省事， 需要根据图片 实际情况进行调整
width = 2048
height = 2048


for line in iter_f:
    tmp = line.split(",")
    
    fileName = tmp[0]
    if(fileName == "filename"):
        continue
    
    ftxt = open("yolo\\"+fileName.split(".")[0]+".txt", 'a')
    
    #filename,width,height,class,xmin,ymin,xmax,ymax
    #<object-class> <x_center> <y_center> <width> <height>
    obj = nameDic[tmp[3]]
    
    x1 = float(tmp[4]) / width
    y1 = float(tmp[5]) / height
    x2 = float(tmp[6]) / width
    y2 = float(tmp[7]) / height
    
    xCenter = (x2 + x1)/2
    yCenter = (y2 + y1)/2
            
    objW = (x2 - x1)
    objH = (y2 - y1)
    
    ftxt.write(str(obj) + " " + str(xCenter)[0:5] +" "+ str(yCenter)[0:5] +" "+ str(objW)[0:5] +" "+ str(objH)[0:5] + "\n")
    ftxt.close()