Update Model

master
tunmx 2 years ago
parent a40ca1f2a1
commit 7230840cf7

@ -72,6 +72,9 @@ class Plate(object):
def to_result(self): def to_result(self):
return [self.plate_code, self.rec_confidence, self.plate_type, self.det_bound_box.tolist(),] return [self.plate_code, self.rec_confidence, self.plate_type, self.det_bound_box.tolist(),]
def to_full_result(self):
return [self.plate_code, self.rec_confidence, self.plate_type, self.det_bound_box.tolist(), self.vertex.tolist()]
def __dict__(self): def __dict__(self):
return self.to_dict() return self.to_dict()

@ -14,7 +14,8 @@ class LicensePlateCatcher(object):
inference: int = INFER_ONNX_RUNTIME, inference: int = INFER_ONNX_RUNTIME,
folder: str = _DEFAULT_FOLDER_, folder: str = _DEFAULT_FOLDER_,
detect_level: int = DETECT_LEVEL_LOW, detect_level: int = DETECT_LEVEL_LOW,
logger_level: int = 3): logger_level: int = 3,
full_result: bool = False):
if inference == INFER_ONNX_RUNTIME: if inference == INFER_ONNX_RUNTIME:
from hyperlpr3.inference.multitask_detect import MultiTaskDetectorORT from hyperlpr3.inference.multitask_detect import MultiTaskDetectorORT
from hyperlpr3.inference.recognition import PPRCNNRecognitionORT from hyperlpr3.inference.recognition import PPRCNNRecognitionORT
@ -31,7 +32,7 @@ class LicensePlateCatcher(object):
raise NotImplemented raise NotImplemented
rec = PPRCNNRecognitionORT(join(folder, ort_cfg['rec_model_path']), input_size=(48, 160)) rec = PPRCNNRecognitionORT(join(folder, ort_cfg['rec_model_path']), input_size=(48, 160))
cls = ClassificationORT(join(folder, ort_cfg['cls_model_path']), input_size=(96, 96)) cls = ClassificationORT(join(folder, ort_cfg['cls_model_path']), input_size=(96, 96))
self.pipeline = LPRMultiTaskPipeline(detector=det, recognizer=rec, classifier=cls) self.pipeline = LPRMultiTaskPipeline(detector=det, recognizer=rec, classifier=cls, full_result=full_result)
else: else:
raise NotImplemented raise NotImplemented

@ -6,10 +6,11 @@ from hyperlpr3.common.tools_process import *
class LPRMultiTaskPipeline(object): class LPRMultiTaskPipeline(object):
def __init__(self, detector, recognizer, classifier): def __init__(self, detector, recognizer, classifier, full_result=False):
self.detector = detector self.detector = detector
self.recognizer = recognizer self.recognizer = recognizer
self.classifier = classifier self.classifier = classifier
self.full_result = full_result
def run(self, image: np.ndarray) -> list: def run(self, image: np.ndarray) -> list:
result = list() result = list()
@ -56,7 +57,10 @@ class LPRMultiTaskPipeline(object):
plate_type = GREEN plate_type = GREEN
plate = Plate(vertex=land_marks, plate_code=plate_code, det_bound_box=np.asarray(rect), plate = Plate(vertex=land_marks, plate_code=plate_code, det_bound_box=np.asarray(rect),
rec_confidence=rec_confidence, dex_bound_confidence=score, plate_type=plate_type) rec_confidence=rec_confidence, dex_bound_confidence=score, plate_type=plate_type)
result.append(plate.to_result()) if self.full_result:
result.append(plate.to_full_result())
else:
result.append(plate.to_result())
return result return result

Loading…
Cancel
Save