Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # 1. Load Models - Fixed paths to match Docker structure | |
| plate_model = YOLO("car_plate_best.pt") | |
| ocr_model = YOLO("model_ocr.pt") | |
| # Explicit mapping to ARABIC-INDIC digits (٠١٢٣٤٥٦٧٨٩) | |
| arabic_mapping = { | |
| "aain": "ع", "alef": "أ", "a": "أ", "alf": "أ", "baa": "ب", | |
| "daal": "د", "dal": "د", "e": "ع", "faa": "ف", "geem": "ج", | |
| "haa": "هـ", "kaaf": "ك", "laam": "ل", "meem": "م", "noon": "ن", | |
| "qaf": "ق", "raa": "ر", "sad": "ص", "seen": "س", "taa": "ط", | |
| "waaw": "و", "waw": "و", "yaa": "ى", "zay": "ز", "dad": "ض", | |
| "0": "٠", "1": "١", "2": "٢", "3": "٣", "4": "٤", | |
| "5": "٥", "6": "٦", "7": "٧", "8": "٨", "9": "٩", | |
| } | |
| def predict_plate(img): | |
| if img is None: | |
| return None, "", "" | |
| try: | |
| # Gradio sends RGB, YOLO likes BGR | |
| img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) | |
| display_img = img_cv.copy() | |
| # 1. Plate Detection | |
| plate_results = plate_model.predict(source=img_cv, conf=0.30, verbose=False) | |
| best_plate_box = None | |
| max_conf = -1 | |
| for box in plate_results[0].boxes: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| conf = float(box.conf) | |
| cv2.rectangle(display_img, (x1, y1), (x2, y2), (0, 255, 0), 3) | |
| if conf > max_conf: | |
| max_conf = conf | |
| best_plate_box = (x1, y1, x2, y2) | |
| letters_str, numbers_str = "", "" | |
| if best_plate_box: | |
| x1, y1, x2, y2 = best_plate_box | |
| h, w = img_cv.shape[:2] | |
| margin = 20 # Increased margin for better OCR on HF | |
| cx1, cy1, cx2, cy2 = max(0, x1-margin), max(0, y1-margin), min(w, x2+margin), min(h, y2+margin) | |
| plate_crop = img_cv[cy1:cy2, cx1:cx2] | |
| if plate_crop.size > 0: | |
| # Optimized OCR params for Hugging Face CPU/GPU | |
| ocr_results = ocr_model.predict(source=plate_crop, conf=0.25, iou=0.45, agnostic_nms=True, verbose=False) | |
| detections = [] | |
| for obox in ocr_results[0].boxes: | |
| cid = int(obox.cls) | |
| ename = ocr_results[0].names[cid] | |
| # Ignore 'License Plate' label | |
| if ename.lower() == "license plate": | |
| continue | |
| achr = arabic_mapping.get(ename, ename) | |
| ox1, _, ox2, _ = obox.xyxy[0] | |
| detections.append(((ox1 + ox2) / 2, achr)) | |
| # Spatial sort Left-to-Right | |
| detections.sort(key=lambda d: d[0]) | |
| arabic_digits_set = set("٠١٢٣٤٥٦٧٨٩") | |
| det_letters = [d[1] for d in detections if d[1] not in arabic_digits_set] | |
| det_numbers = [d[1] for d in detections if d[1] in arabic_digits_set] | |
| # Reverse for RTL display | |
| det_numbers.reverse() | |
| det_letters.reverse() | |
| letters_str = " ".join(det_letters) | |
| numbers_str = " ".join(det_numbers) | |
| return cv2.cvtColor(display_img, cv2.COLOR_BGR2RGB), numbers_str, letters_str | |
| except Exception: | |
| return None, "Error", "Error" | |
| # UI Setup | |
| with gr.Blocks(theme=gr.themes.Default()) as demo: | |
| gr.Markdown("# 🚗 Arabic License Plate OCR") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_img = gr.Image(label="Upload or Camera", type="pil") | |
| btn = gr.Button("Extract Plate Info", variant="primary") | |
| with gr.Column(): | |
| output_img = gr.Image(label="Detection Result") | |
| with gr.Row(): | |
| # Numbers (Left) | Letters (Right) | |
| out_numbers = gr.Textbox(label="Numbers (Arabic Script)", interactive=False) | |
| out_letters = gr.Textbox(label="Arabic Letters", interactive=False) | |
| btn.click(fn=predict_plate, inputs=input_img, outputs=[output_img, out_numbers, out_letters], api_name="predict_plate") | |
| if __name__ == "__main__": | |
| # CRITICAL: Added server_name and server_port for Hugging Face | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |