Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| import cv2 | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| import pandas as pd | |
| # Model labels | |
| model1Labels = {0: 'single_number_plate', 1: 'double_number_plate'} | |
| model2Labels = { | |
| 0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: 'A', 11: 'B', 12: 'C', | |
| 13: 'D', 14: 'E', 15: 'F', 16: 'G', 17: 'H', 18: 'I', 19: 'J', 20: 'K', 21: 'L', 22: 'M', 23: 'N', 24: 'O', | |
| 25: 'P', 26: 'Q', 27: 'R', 28: 'S', 29: 'T', 30: 'U', 31: 'V', 32: 'W', 33: 'X', 34: 'Y', 35: 'Z' | |
| } | |
| # Load models | |
| model = YOLO("model/LP-detection.pt") | |
| model2 = YOLO("model/Charcter-LP.pt") | |
| def prediction(image): | |
| result = model.predict(source=image, conf=0.5) | |
| boxes = result[0].boxes | |
| height = boxes.xywh | |
| crd = boxes.data | |
| n = len(crd) | |
| lp_number = [] | |
| img_lp_final = None | |
| for i in range(n): | |
| ht = int(height[i][3]) | |
| c = int(crd[i][5]) | |
| xmin = int(crd[i][0]) | |
| ymin = int(crd[i][1]) | |
| xmax = int(crd[i][2]) | |
| ymax = int(crd[i][3]) | |
| img_lp = image[ymin:ymax, xmin:xmax] | |
| img_lp_final = img_lp.copy() # Store the cropped image for display | |
| cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2) | |
| h = np.median(ht) | |
| # Second Model Prediction | |
| result2 = model2.predict(source=img_lp, conf=0.25) | |
| boxes_ocr = result2[0].boxes | |
| data2 = boxes_ocr.data | |
| n2 = len(data2) | |
| xaxis0, xaxis11, xaxis12 = [], [], [] | |
| label0, label11, label12 = [], [], [] | |
| numberPlate = "" | |
| if c == 0: # Single line license plate | |
| for i in range(n2): | |
| x = int(data2[i][2]) | |
| xaxis0.append(x) | |
| l = int(data2[i][5]) | |
| label0.append(l) | |
| # Sort characters by x-axis for single line | |
| sorted_labels = [label0[i] for i in np.argsort(xaxis0)] | |
| numberPlate = ''.join([model2Labels.get(l) for l in sorted_labels]) | |
| lp_number.append(numberPlate) | |
| elif c == 1: # Double line license plate | |
| for i in range(n2): | |
| x = int(data2[i][0]) | |
| y = int(data2[i][3]) | |
| l = int(data2[i][5]) | |
| if y < (h / 2): | |
| xaxis11.append(x) | |
| label11.append(l) | |
| else: | |
| xaxis12.append(x) | |
| label12.append(l) | |
| # Sort characters by x-axis for double line (upper and lower separately) | |
| sorted_labels11 = [label11[i] for i in np.argsort(xaxis11)] | |
| sorted_labels12 = [label12[i] for i in np.argsort(xaxis12)] | |
| numberPlate = ''.join([model2Labels.get(l) for l in sorted_labels11 + sorted_labels12]) | |
| lp_number.append(numberPlate) | |
| return lp_number, img_lp_final | |
| def process_video(video_file): | |
| # Open the video | |
| cap = cv2.VideoCapture(video_file) | |
| license_plate_texts = [] | |
| processed_frames = [] | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| license_plate_text, cropped_plate_img = prediction(frame) | |
| license_plate_texts.append(" ".join(license_plate_text)) # Join the list of texts into a single string | |
| processed_frames.append(cropped_plate_img) | |
| # Save detected texts to Excel | |
| df = pd.DataFrame(license_plate_texts, columns=["License Plate"]) | |
| df.to_excel("detected_license_plates.xlsx", index=False) | |
| # Save processed video with license plates highlighted | |
| output_video_path = 'processed_video.mp4' | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for MP4 | |
| out = cv2.VideoWriter(output_video_path, fourcc, 20.0, (frame.shape[1], frame.shape[0])) | |
| for processed_frame in processed_frames: | |
| out.write(processed_frame) | |
| cap.release() | |
| out.release() | |
| return output_video_path, "detected_license_plates.xlsx" | |
| # Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🚗 License Plate Recognition (Video Upload)") | |
| gr.Markdown("Upload a video to get the license number of vehicles detected in each frame.") | |
| with gr.Row(): | |
| video_input = gr.File(label="Upload Video", type="filepath") # Corrected the file type | |
| video_output = gr.Video(label="Processed Video") | |
| excel_output = gr.File(label="Excel File with Detected License Plates") | |
| video_input.upload(process_video, inputs=video_input, outputs=[video_output, excel_output]) | |
| demo.launch() | |