vehicle-detection.py

import sys
import cv2
import numpy as np

import darknet.python.darknet as dn

from src.label 				import Label, lwrite
from os.path 				import splitext, basename, isdir
from os 					import makedirs
from src.utils 				import crop_region, image_files_from_folder
from darknet.python.darknet import detect


if __name__ == '__main__':
	
	input_dir  = sys.argv[1]
	output_dir = sys.argv[2]

	vehicle_threshold = .5

	vehicle_weights = 'data/vehicle-detector/yolo-voc.weights'
	vehicle_netcfg  = 'data/vehicle-detector/yolo-voc.cfg'
	vehicle_dataset = 'data/vehicle-detector/voc.data'

	vehicle_net  = dn.load_net(vehicle_netcfg, vehicle_weights, 0)
	vehicle_meta = dn.load_meta(vehicle_dataset)

	imgs_paths = image_files_from_folder(input_dir)
	imgs_paths.sort()

	if not isdir(output_dir):
		makedirs(output_dir)

	print 'Searching for vehicles using YOLO...'

	for i,img_path in enumerate(imgs_paths):

		print '\tScanning %s' % img_path

		bname = basename(splitext(img_path)[0])

		R = detect(vehicle_net, vehicle_meta, img_path ,thresh=vehicle_threshold)

		print '\t\t%d cars found' % len(R)

		if len(R):

			Iorig = cv2.imread(img_path)
			WH = np.array(Iorig.shape[1::-1],dtype=float)#(start:end:step)
			Lcars = []

			for i,r in enumerate(R):

				cx,cy,w,h = (np.array(r[2])/np.concatenate( (WH,WH) )).tolist()#to list
				tl = np.array([cx - w/2., cy - h/2.])#top_left
				br = np.array([cx + w/2., cy + h/2.])#bottom_right
				label = Label(0,tl,br)#create an object of class Label
				Icar = crop_region(Iorig,label)# Icar is an image  croped by the Iorig

				Lcars.append(label)

				cv2.imwrite('%s/%s_%dcar.png' % (output_dir,bname,i),Icar)

			lwrite('%s/%s_cars.txt' % (output_dir,bname),Lcars)