|
import os |
|
import cv2 |
|
import gradio as gr |
|
import numpy as np |
|
|
|
from scripts.charExtraction import CharExtraction |
|
from scripts.bboxAnnotator import BBOXAnnotator |
|
|
|
wPathPlat = "licence_plat.pt" |
|
wPathChar = "licence_character.pt" |
|
classList = np.array(['A','B','C','D','E','F','G','H','I','J','K','L','M', |
|
'N','O','P','Q','R','S','T','U','V','W','X','Y','Z', |
|
'0','1','2','3','4','5','6','7','8','9']) |
|
sizePlat = (416,200) |
|
|
|
|
|
|
|
extractor = CharExtraction(wPlatePath=wPathPlat, wCharacterPath=wPathChar, classList=classList, |
|
sizePlate=sizePlat, conf=0.5) |
|
annotator = BBOXAnnotator() |
|
|
|
def getAnnotatedImage(frame, conf): |
|
bbox, plateNum, confidence = extractor.predict(image=frame, conf=conf) |
|
annotateImg, plateNum = annotator.draw_bbox(frame, bbox, plateNum) |
|
prob_dict = {key: value for key, value in zip(plateNum, confidence)} |
|
|
|
|
|
|
|
|
|
return annotateImg, prob_dict |
|
|
|
def video_identity(frame, conf=0.45): |
|
image, plateNumber = getAnnotatedImage(frame, conf) |
|
|
|
return image, plateNumber |
|
|
|
|
|
demo = gr.Interface(fn=video_identity, |
|
inputs=[gr.Image(label='Input image'), gr.Slider(minimum=0.1, maximum=1.0, label='Conf value')], |
|
outputs=[gr.Image(label='Result'), gr.Label(label='Plate Number Detected')], |
|
examples = [["sample_image2.jpg"], ["sample_image3.jpg"], ["sample_image5.jpg"], ["sample_image6.jpg"]], |
|
cache_examples=True) |
|
|
|
if __name__ == "__main__": |
|
demo.queue() |
|
demo.launch() |