anno2txt.py

import tqdm
import imagesize
from pathlib import Path
from os import path

class_names_to_id={'Stop':0, 'Children':1, 'Bus':2, 'TL_Go':3, 'TL_Left':4, 'TL_GoLeft':5, 'TL_Stop':6}

def read_anno(filename, savename, size):
    with open(filename, 'r') as f:
        content=f.readlines()
    content=[x.strip() for x in content]
    print_buffer=[]
    classes={'Stop':0, 'Children':0, 'Bus':0, 'TL_Go':0, 'TL_Left':0, 'TL_GoLeft':0, 'TL_Stop':0}
    classes_arr={'Stop':[], 'Children':[], 'Bus':[], 'TL_Go':[], 'TL_Left':[], 'TL_GoLeft':[], 'TL_Stop':[]}
    #output_arr={'Stop':[], 'Children':[], 'Bus':[], 'TL_Go':[], 'TL_Left':[], 'TL_GoLeft':[], 'TL_Stop':[]}
    current_class='Stop'
    flag_paren=False
    points_arr=[]
    #stop, children, bus, tl_go, tl_left, tl_goleft, tl_stop=0,0,0,0,0,0,0
    for lines in content:
        parts=lines.split(' ')
        first=parts[0].strip(':",')
        if classes.keys().__contains__(first):
            current_class=first
        if first=='[':
            flag_paren=True
        elif first==']' and flag_paren:
            flag_paren=False
            classes_arr[current_class].append(points_arr)
            points_arr=[]
            classes[current_class]=classes[current_class]+1
        elif flag_paren:
            points_arr.append(int(first))
    print(classes_arr)
    for key in classes_arr.keys():
        c=classes_arr[key]
        for x in range(len(c)):
            b_center_x =(c[x][0]+c[x][2])/2/size[0]
            b_center_y =(c[x][1]+c[x][3])/2/size[1]
            b_width =abs((c[x][2]-c[x][0])/size[0])
            b_height =abs((c[x][3]-c[x][1])/size[1])
            print_buffer.append(
                "{} {:.3f} {:.3f} {:.3f} {:.3f}".format(class_names_to_id[key], b_center_x, b_center_y, b_width, b_height))
    print(print_buffer)
    print("\n".join(print_buffer), file=open(savename, "w"))


for p in Path('../train/labels').glob('*'):
    #anno_path= 'data/labels/2_Bus_0001.anno'
    img_path=str(p).replace("labels", "imgs").replace("anno", "jpg")
    if path.exists(img_path):
        txt_path = str(p).replace("labels", "texts").replace("anno", "txt")
        width, height = imagesize.get(img_path)
        read_anno(str(p), txt_path, [width, height])


#class_id center_x center_y width height
#