Spaces:
Running
Running
| import torch | |
| from ultralytics import YOLO | |
| import easyocr | |
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| import os | |
| import logging | |
| import re | |
| # Set up logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| # Ensure directories exist | |
| os.makedirs(os.getenv('EASYOCR_MODULE_PATH', '/app/.EasyOCR'), exist_ok=True) | |
| os.makedirs(os.getenv('YOLO_CONFIG_DIR', '/app/.config/Ultralytics'), exist_ok=True) | |
| # Download pretrained model | |
| ANPR_WEIGHTS = "anpr_yolov8.pt" | |
| if not os.path.exists(ANPR_WEIGHTS): | |
| logger.info(f"Downloading model weights to {ANPR_WEIGHTS}") | |
| os.system(f"wget -O {ANPR_WEIGHTS} https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt") | |
| # Load YOLO model | |
| try: | |
| model = YOLO(ANPR_WEIGHTS) | |
| logger.info(f"Successfully loaded YOLO model from {ANPR_WEIGHTS}") | |
| except Exception as e: | |
| logger.error(f"Error loading YOLO model from {ANPR_WEIGHTS}: {str(e)}") | |
| raise | |
| # Load OCR reader | |
| try: | |
| reader = easyocr.Reader(['en'], model_storage_directory=os.getenv('EASYOCR_MODULE_PATH', '/app/.EasyOCR')) | |
| logger.info("Successfully initialized EasyOCR reader") | |
| except Exception as e: | |
| logger.error(f"Error initializing EasyOCR reader: {str(e)}") | |
| raise | |
| def detect_and_read_plate(image): | |
| logger.info("Starting image processing for license plate detection") | |
| try: | |
| detected_texts = [] | |
| results = model(image, conf=0.1) # Low confidence for broader detection | |
| logger.info(f"YOLO model returned {len(results)} results") | |
| for result in results: | |
| boxes = result.boxes.xyxy.cpu().numpy() | |
| confidences = result.boxes.conf.cpu().numpy() | |
| logger.info(f"Detected {len(boxes)} bounding boxes") | |
| for box, conf in zip(boxes, confidences): | |
| x1, y1, x2, y2 = map(int, box) | |
| # Minimal size filter | |
| if (x2 - x1) < 20 or (y2 - y1) < 10: | |
| logger.warning(f"Skipping small box: ({x1}, {y1}, {x2}, {y2})") | |
| continue | |
| # Crop the detected license plate | |
| plate_img = image[y1:y2, x1:x2] | |
| if plate_img.size == 0: | |
| logger.warning("Empty cropped image, skipping") | |
| continue | |
| # Run OCR | |
| logger.info("Running EasyOCR on cropped plate") | |
| ocr_result = reader.readtext(plate_img) | |
| if ocr_result: | |
| for res in ocr_result: | |
| text = res[1] | |
| confidence = res[2] | |
| # Light filtering: prefer alphanumeric, avoid overly long text | |
| if len(text) <= 12 and confidence > 0.2 and re.match(r'^[A-Z0-9\s\-]+$', text): | |
| detected_texts.append(f"{text} (conf: {conf:.2f})") | |
| logger.info(f"Detected Plate: {text} (YOLO conf: {conf:.2f}, OCR conf: {confidence:.2f})") | |
| else: | |
| logger.info(f"Rejected text: {text} (OCR conf: {confidence:.2f})") | |
| # Draw bounding box | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| output_text = "\n".join(detected_texts) if detected_texts else "No license plate detected" | |
| logger.info(f"Output text: {output_text}") | |
| return image, output_text | |
| except Exception as e: | |
| logger.error(f"Error during detection: {str(e)}") | |
| return image, f"Error processing image: {str(e)}" | |
| # Create Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Automatic Number Plate Recognition (ANPR)") | |
| gr.Markdown("Upload an image of a car to detect and read its license plate. Results may take a few seconds.") | |
| with gr.Row(): | |
| image_input = gr.Image(type="numpy", label="Upload an image of a car") | |
| with gr.Row(): | |
| image_output = gr.Image(type="numpy", label="Detected License Plate Image") | |
| text_output = gr.Textbox(label="Detected License Plate Number") | |
| submit_button = gr.Button("Detect License Plate") | |
| submit_button.click( | |
| fn=detect_and_read_plate, | |
| inputs=image_input, | |
| outputs=[image_output, text_output], | |
| show_progress=True | |
| ) | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |