import csv
from os import listdir, getcwd
import pickle
import os
from os.path import join


classes = ["car", "truck"]
def convert_annotation(image_id):
    cl=[]
    coord=[]
    b=()
    out_file = open('labels/%s.txt'%(image_id), 'w')
    with open('udacity.csv', 'rb') as csvfile:
        spamreader = csv.reader(csvfile, delimiter=' ', quotechar='"')
        for row in spamreader:
            if row[0]==image_id:
                for i in range(0,len(row)/5):
                    c= int(row[(i*5)+5])
                    if c not in [1, 3, 4]:
                        cl.append(c)
                        coord.append((row[i*5+1],row[i*5+3],row[i*5+2],row[i*5+4]))
                    else:
                        continue

        for i in range(0,len(cl)):
            b=(coord[i][0], coord[i][1], coord[i][2], coord[i][3])
            w=int(coord[i][1])-int(coord[i][0])
            h=int(coord[i][3])-int(coord[i][2])
            if cl[i] == 0:
                out_file.write("car (" + " ".join([str(a) for a in b]) +") "+str(w)+" "+str(h)+ '\n')
            else:
                out_file.write("truck (" + " ".join([str(a) for a in b]) +") "+str(w)+" "+str(h)+ '\n')

wd = getcwd()
if not os.path.exists('labels/'):
    os.makedirs('labels/')
list_file = open('imageSet.txt', 'w')
with open('udacity.csv', 'rb') as csvfile:
    reader = csv.reader(csvfile, delimiter=' ',quotechar='"')
    for row in reader:
        list_file.write('%s/JPEGImages/%s\n'%(wd, row[0]))
        convert_annotation(row[0])
    list_file.close()