Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| from utils import perspective_transform | |
| from ultralytics import YOLO | |
| max_len = 10 | |
| def ocr(plate_image, ocr_model): | |
| try: | |
| if 'yolo' in ocr_model: | |
| model = YOLO(f'./weights/{ocr_model}.pt', task='detect') | |
| # model = YOLO(f'./weights/best.pt', task='detect') | |
| preds = model.predict(plate_image) | |
| pred_data = [list(pred.boxes.data.cpu().numpy()[0]) for pred in preds[0]] | |
| x_sorted_preds = sorted(pred_data, key=lambda x:x[0]) | |
| pred_cls = np.array(x_sorted_preds,dtype=np.int32)[:,-1] | |
| return ''.join([model.names[cls] for cls in pred_cls]) | |
| else: | |
| return '' | |
| except Exception : | |
| return 'error' | |
| def process_image(image, detection_model, ocr_model, yolo_thresh, perpective_width, perpective_height): | |
| task = 'obb' if 'obb' in detection_model else 'detect' | |
| model = YOLO(f'./weights/{detection_model}.pt', task=task) | |
| predict = model(image, conf=yolo_thresh) | |
| if 'obb' in detection_model: | |
| obb_crops, transformed = perspective_transform(predict, dst_width=perpective_width, dst_height=perpective_height) | |
| crop_results = [(plate, ocr(plate, ocr_model)) for plate in obb_crops] | |
| transform_results = [(plate, ocr(plate, ocr_model)) for plate in transformed] | |
| return crop_results, transform_results | |
| else: | |
| boxes = np.array(predict[0].boxes.xyxy.cpu().numpy(),dtype=np.int32) | |
| crops = [image[y1:y2,x1:x2,:] for x1,y1,x2,y2 in boxes] | |
| results = [(plate,ocr(plate,ocr_model)) for plate in crops] | |
| return results, [] | |
| def create_interface(): | |
| with gr.Blocks(css="footer{display:none !important}") as demo: | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| detection_model = gr.Dropdown(label="Detection Model", choices=["yolov8-m", "yolov8-obb-m","yolov8-s", "yolov8-obb-s"], value="yolov8-m") | |
| ocr_model = gr.Dropdown(label="OCR Model", choices=["yolov32c", "trocr"], value="yolov32c") | |
| yolo_thresh = gr.Slider(minimum=0.0, maximum=1.0, label='yolo_threshold') | |
| with gr.Row(): | |
| perpective_width = gr.Slider(minimum=100, maximum=500, label='perpective_width') | |
| perpective_height = gr.Slider(minimum=50, maximum=200, label='perpective_height') | |
| with gr.Column(scale=3): | |
| input_image = gr.Image(label="Upload Image", type="numpy") | |
| with gr.Row() as r1: | |
| c_img1 = gr.Image(visible=False) | |
| c_text1 = gr.Textbox(visible=False) | |
| t_img1 = gr.Image(visible=False) | |
| t_text1 = gr.Textbox(visible=False) | |
| with gr.Row() as r2: | |
| c_img2 = gr.Image(visible=False) | |
| c_text2 = gr.Textbox(visible=False) | |
| t_img2 = gr.Image(visible=False) | |
| t_text2 = gr.Textbox(visible=False) | |
| with gr.Row() as r3: | |
| c_img3 = gr.Image(visible=False) | |
| c_text3 = gr.Textbox(visible=False) | |
| t_img3 = gr.Image(visible=False) | |
| t_text3 = gr.Textbox(visible=False) | |
| with gr.Row() as r4: | |
| c_img4 = gr.Image(visible=False) | |
| c_text4= gr.Textbox(visible=False) | |
| t_img4 = gr.Image(visible=False) | |
| t_text4 = gr.Textbox(visible=False) | |
| with gr.Row() as r5: | |
| c_img5 = gr.Image(visible=False) | |
| c_text5 = gr.Textbox(visible=False) | |
| t_img5 = gr.Image(visible=False) | |
| t_text5 = gr.Textbox(visible=False) | |
| #o = gr.Textbox() | |
| def main_fn(image, detection_model, ocr_model, yolo_thresh, perpective_width, perpective_height): | |
| crop_results, transform_results = process_image(image, detection_model, ocr_model, yolo_thresh, perpective_width, perpective_height) | |
| crop_output = [] | |
| for i,(crop, c_text) in enumerate(crop_results): | |
| crop_output.append(gr.Image(label=f'clp{i+1}', value=crop, visible=True)) | |
| crop_output.append(gr.Textbox(label=f'text{i+1}',value=c_text, visible=True)) | |
| for j in range((max_len - len(crop_output)) //2): | |
| crop_output.append(gr.Image(visible=False)) | |
| crop_output.append(gr.Textbox(visible=False)) | |
| transform_output = [] | |
| for i,(transform, t_text) in enumerate(transform_results): | |
| transform_output.append(gr.Image(label=f'tlp{i+1}', value=transform, visible=True)) | |
| transform_output.append(gr.Textbox(label=f'text{i+1}',value=t_text, visible=True)) | |
| for j in range((max_len - len(transform_output)) //2): | |
| transform_output.append(gr.Image(visible=False)) | |
| transform_output.append(gr.Textbox(visible=False)) | |
| return crop_output + transform_output | |
| # return len(crop_results) | |
| submit_button = gr.Button("Process Image") | |
| submit_button.click( | |
| fn=main_fn, | |
| inputs=[input_image, detection_model, ocr_model,yolo_thresh, perpective_width, perpective_height], | |
| outputs=[c_img1, c_text1, c_img2, c_text2, c_img3, c_text3, c_img4, c_text4, c_img5, c_text5, t_img1, t_text1, t_img2, t_text2, t_img3, t_text3, t_img4, t_text4, t_img5, t_text5] | |
| # outputs = o | |
| ) | |
| return demo | |
| # Run the Gradio app | |
| demo = create_interface() | |
| demo.launch() | |