From 7867a42e4eaaf0b85cabf4d1e7c4222d77db6d28 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Montazzolli?= Date: Wed, 26 Dec 2018 17:09:26 -0200 Subject: [PATCH] Small NMS correction in OCR. --- README.md | 2 +- darknet/python/darknet.py | 3 ++- license-plate-ocr.py | 13 +++++++++---- src/label.py | 9 +++++++++ vehicle-detection.py | 2 +- 5 files changed, 22 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 3fffd617..d614dce4 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ In order to easily run the code, you must have installed the Keras framework wit $ cd darknet && make ``` -**The current version was tested in an Ubuntu 16.04 machine, with Keras 2.0.6, TensorFlow 1.5.0 and Python 2.7.** +**The current version was tested in an Ubuntu 16.04 machine, with Keras 2.2.4, TensorFlow 1.5.0 and Python 2.7.** ## Download Models diff --git a/darknet/python/darknet.py b/darknet/python/darknet.py index 59e3679d..c0cab3ab 100644 --- a/darknet/python/darknet.py +++ b/darknet/python/darknet.py @@ -138,9 +138,10 @@ def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45): b = dets[j].bbox res.append((meta.names[i], dets[j].prob[i], (b.x, b.y, b.w, b.h))) res = sorted(res, key=lambda x: -x[1]) + wh = (im.w,im.h) free_image(im) free_detections(dets, num) - return res + return res,wh if __name__ == "__main__": #net = load_net("cfg/densenet201.cfg", "/home/pjreddie/trained/densenet201.weights", 0) diff --git a/license-plate-ocr.py b/license-plate-ocr.py index de510948..7210f9e4 100644 --- a/license-plate-ocr.py +++ b/license-plate-ocr.py @@ -7,6 +7,8 @@ from os.path import splitext, basename from glob import glob from darknet.python.darknet import detect +from src.label import dknet_label_conversion +from src.utils import nms if __name__ == '__main__': @@ -23,7 +25,7 @@ ocr_net = dn.load_net(ocr_netcfg, ocr_weights, 0) ocr_meta = dn.load_meta(ocr_dataset) - imgs_paths = glob('%s/*lp.png' % output_dir) + imgs_paths = sorted(glob('%s/*lp.png' % output_dir)) print 'Performing OCR...' @@ -33,12 +35,15 @@ bname = basename(splitext(img_path)[0]) - R = detect(ocr_net, ocr_meta, img_path ,thresh=ocr_threshold) + R,(width,height) = detect(ocr_net, ocr_meta, img_path ,thresh=ocr_threshold, nms=None) if len(R): - R.sort(key=lambda x: x[2][0]) - lp_str = ''.join([r[0] for r in R]) + L = dknet_label_conversion(R,width,height) + L = nms(L,.45) + + L.sort(key=lambda x: x.tl()[0]) + lp_str = ''.join([chr(l.cl()) for l in L]) with open('%s/%s_str.txt' % (output_dir,bname),'w') as f: f.write(lp_str + '\n') diff --git a/src/label.py b/src/label.py index 25462e13..50deb9bc 100644 --- a/src/label.py +++ b/src/label.py @@ -85,6 +85,15 @@ def lwrite(file_path,labels,write_probs=True): fd.write('%d %f %f %f %f\n' % (cl,cc[0],cc[1],wh[0],wh[1])) +def dknet_label_conversion(R,img_width,img_height): + WH = np.array([img_width,img_height],dtype=float) + L = [] + for r in R: + center = np.array(r[2][:2])/WH + wh2 = (np.array(r[2][2:])/WH)*.5 + L.append(Label(ord(r[0]),tl=center-wh2,br=center+wh2,prob=r[1])) + return L + class Shape(): diff --git a/vehicle-detection.py b/vehicle-detection.py index 52cee3a3..dc07d144 100644 --- a/vehicle-detection.py +++ b/vehicle-detection.py @@ -39,7 +39,7 @@ bname = basename(splitext(img_path)[0]) - R = detect(vehicle_net, vehicle_meta, img_path ,thresh=vehicle_threshold) + R,_ = detect(vehicle_net, vehicle_meta, img_path ,thresh=vehicle_threshold) R = [r for r in R if r[0] in ['car','bus']]