Asis / alpr.py
RivianG's picture
Update alpr.py
8fff871
import os
import cv2
import numpy as np
import tensorflow.compat.v1 as tf
from numpy.linalg import norm
from local_utils import detect_lp
from os.path import splitext
from tensorflow.python.keras.backend import set_session
from tensorflow.keras.models import model_from_json
from tensorflow.compat.v1 import ConfigProto
class DetectLicensePlate:
def __init__(self):
tf.compat.v1.disable_eager_execution()
#config = ConfigProto()
#config.gpu_options.allow_growth = False
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.433)
self.sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
#self.sess = tf.Session(config=config)
self.graph = tf.get_default_graph()
set_session(self.sess)
self.wpod_net_path ="wpod-net.json" # model path
self.wpod_net = self.load_model(self.wpod_net_path)
def load_model(self, path):
try:
path = splitext(path)[0]
with open('%s.json' % path, 'r') as json_file:
model_json = json_file.read()
model = model_from_json(model_json, custom_objects={})
model.load_weights('%s.h5' % path)
print("Loading model successfully...")
self.graph = tf.get_default_graph()
return model
except Exception as e:
print(e)
def preprocess_image(self, image_path, resize=False):
img = image_path
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img / 255
if resize:
img = cv2.resize(img, (224, 224))
return img
def get_plate(self, image_path, Dmax=608, Dmin=608):
vehicle = self.preprocess_image(image_path)
ratio = float(max(vehicle.shape[:2])) / min(vehicle.shape[:2])
side = int(ratio * Dmin)
bound_dim = min(side, Dmax)
_, plates, _, cor = detect_lp(self.graph,self.sess,self.wpod_net, vehicle, bound_dim, lp_threshold=0.5)
return vehicle, plates, cor
def alpr(frame,license_plate):
plate_image = frame.copy()
try:
vehicle, plates, cor = license_plate.get_plate(frame)
return plates[0]
except Exception as e:
print(str(e))