botch 数据集 标注信息转换为 csv 的 python实现

长期以来都是使用 AWK, 但 botch 这个好像用 awk略麻烦，使用python 试试，就当练手了
主要是 使用 python的文件、字符串 操作，实现了 一个 有限状态机

# -*- coding: utf-8 -*-
"""
Created on Wed Mar 20 14:33:33 2019

@author: haithink
"""

import os

file = open("train.yaml", "r")

iter_f = iter(file)

bIn = False;
labelIn = False;

i = 0;

def getString(line, prefix):
    x_MaxIdx = line.index(prefix)
    commaIdx = line.find(",", x_MaxIdx+1)
    if(commaIdx == -1):
        commaIdx = line.index("}", x_MaxIdx+1)    
    x_max = line[x_MaxIdx+len(prefix):commaIdx]
    
    return x_max

def getFloat(line, prefix):
    x_max = getString(line, prefix)
    return float(x_max)    

def toInt(sss):
    return str(int(float(sss)))

def outputOnePic(path, coordinates):
    for coor in coordinates:
        print(path + ",1280,720" + ","+coor[4] +","
              + toInt(coor[0])+"," + toInt(coor[1])
              + "," + toInt(coor[2]) + "," + toInt(coor[3]) )

path = ""
x_max = 0
x_min = 0
y_max = 0
y_min = 0
coordinates=[]

print("filename,width,height,class,xmin,ymin,xmax,ymax")
for line in iter_f:
    if(bIn == False):
        if("boxes: []" in line):
            noBox = True
        else:
            noBox = False
        bIn = True;
    else:
        if("path" in line):
            pathIdx = line.index("path")
            path = line[pathIdx+6:-1]
            # TODO: 输出一张图所有信息
#           print(path) 
            outputOnePic(path, coordinates)
            coordinates = []
           
        elif(("{label" in line) or labelIn == True):
            if("{label" in line):
                label = getString(line,"label: ")
                #print(label)
            
            if("x_max: " in line):
                x_max = getString(line, "x_max: ")            
                #print(x_max)
            
            if("x_min" in line):
                x_min = getString(line, "x_min: ")
                #print("xmin is " + str(x_min))
            if("y_max" in line):
                y_max = getString(line, "y_max: ")
            
            if("y_min" in line):
                y_min = getString(line, "y_min: ")
            
            if("}" not in line):
                labelIn = True
            else:
                labelIn = False
                temp = [x_min, y_min, x_max, y_max, label]
                coordinates.append(temp)
#    i = i+1
#    if(i == 20):
#        break
file.close()