Spaces:
Runtime error
Runtime error
| from paddleocr import PaddleOCR | |
| from PIL import Image | |
| import json | |
| import gradio as gr | |
| import numpy as np | |
| import cv2 | |
| import torch | |
| model = torch.hub.load('ultralytics/yolov5', 'custom', path='./best_Plate.pt') # local model | |
| def get_random_color(): | |
| c = tuple(np.random.randint(0, 256, 3).tolist()) | |
| return c | |
| def draw_ocr_bbox(image, boxes, colors): | |
| print(colors) | |
| box_num = len(boxes) | |
| for i in range(box_num): | |
| box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64) | |
| image = cv2.polylines(np.array(image), [box], True, colors[i], 2) | |
| return image | |
| def inference(img: Image.Image, lang, confidence): | |
| ocr = PaddleOCR(use_angle_cls=True, lang=lang, use_gpu=False) | |
| # img_path = img.name | |
| det_img = model(img) | |
| det_croppeds = det_img.crop(save=False) | |
| img_render = det_img.render()[0] | |
| if len(det_croppeds) > 0: | |
| img = det_croppeds[0]['im'] | |
| img = Image.fromarray(img) | |
| img2np = np.array(img) | |
| results = ocr.ocr(img2np, cls=True) | |
| result = results[0] | |
| if result == None: | |
| return img_render, None, None | |
| image = img.convert('RGB') | |
| boxes = [line[0] for line in result] | |
| txts = [line[1][0] for line in result] | |
| scores = [line[1][1] for line in result] | |
| final_result = [dict(boxes=box, txt=txt, score=score, _c=get_random_color()) for box, txt, score in zip(boxes, txts, scores)] | |
| final_result = [item for item in final_result if item['score'] > confidence] | |
| im_show = draw_ocr_bbox(image, [item['boxes'] for item in final_result], [item['_c'] for item in final_result]) | |
| im_show = Image.fromarray(im_show) | |
| data = [[json.dumps(item['boxes']), round(item['score'], 3), item['txt']] for item in final_result] | |
| return img_render, im_show, data | |
| title = 'License Plate' | |
| description = 'Demo License Plate Recognition' | |
| examples = [ | |
| # ['example_imgs/example.jpg','en', 0.5], | |
| ] | |
| css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}" | |
| if __name__ == '__main__': | |
| demo = gr.Interface( | |
| inference, | |
| [ | |
| gr.Image(type='pil', label='Input'), | |
| gr.Dropdown(choices=['en', 'ar'], value='en', label='Language'), | |
| gr.Slider(0.1, 1, 0.5, step=0.1, label='Confidence Threshold') | |
| ], | |
| [ | |
| gr.Image(type='pil', label='License Plate Detection'), | |
| gr.Image(type='pil', label='License Plate'), | |
| gr.Dataframe(headers=[ 'bbox', 'score', 'text'], label='Result'), | |
| ], | |
| title=title, | |
| description=description, | |
| # examples=examples, | |
| css=css, | |
| ) | |
| demo.queue(max_size=10) | |
| demo.launch(debug=True, server_name="0.0.0.0") | |