Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| import cv2 | |
| import numpy as np | |
| from paddleocr import PaddleOCR | |
| import re | |
| # Load models once | |
| model_path = 'yolo11x.pt' | |
| model2_path = r"yolo11_anpr_ghd.pt" | |
| model = YOLO(model_path) | |
| model2 = YOLO(model2_path) | |
| # Initialize PaddleOCR globally to avoid re-initializing on each call | |
| # For Hugging Face Spaces, ensure necessary PaddleOCR models are downloaded (e.g., within Dockerfile or setup script) | |
| # Using 'en' for English license plates and disabling angle classification for simplicity | |
| ocr = PaddleOCR(use_doc_orientation_classify=False, | |
| use_doc_unwarping=False, | |
| use_textline_orientation=False, | |
| lang='en') # Added lang='en' for better English OCR | |
| # Example images paths - these might need adjustment for Hugging Face Spaces | |
| # If these are locally stored, they need to be part of the uploaded files in your Space. | |
| example_images = [ | |
| r"examples/car1.png", | |
| r"examples/car2.png", | |
| r"examples/car3.png", | |
| r"examples/car4.png", | |
| r"examples/car5.png", | |
| ] | |
| def detect_and_annotate(image): | |
| # Convert PIL image to OpenCV format | |
| img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) | |
| # Initialize message | |
| output_message = "" | |
| # Start with a copy of the original image for annotations | |
| annotated_frame = img.copy() | |
| # Vehicle detection | |
| results = model(img, classes=[2, 5, 7, 3]) # vehicle classes | |
| # IMPORTANT FIX: The .plot() method on results[0] already returns a new image with annotations. | |
| # We should use this directly as our base annotated_frame. | |
| if results and results[0].boxes and len(results[0].boxes) > 0: # Check if vehicle results exist and have boxes | |
| # Plot vehicle detections. This returns an annotated NumPy array (BGR format). | |
| annotated_frame = results[0].plot(line_width=3) | |
| # License plate detection only if vehicle is detected | |
| license_plate_results = model2(img) | |
| if license_plate_results and len(license_plate_results[0].boxes) > 0: | |
| # Iterate through detected license plates if there are multiple | |
| for license_box_data in license_plate_results[0].boxes: | |
| license_box = license_box_data.xyxy[0].cpu().numpy() | |
| x1, y1, x2, y2 = map(int, license_box) | |
| # Draw rectangle for license plate on the annotated frame | |
| cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (255, 0, 0), 3) | |
| # OCR on license plate region | |
| license_plate_region = img[y1:y2, x1:x2] | |
| try: | |
| # IMPORTANT FIX: Pass the NumPy array directly to PaddleOCR's ocr method | |
| # This avoids the file read/write error. | |
| ocr_result = ocr.ocr(img=license_plate_region, cls=True) # Changed from .predict to .ocr | |
| license_plate_text = "" | |
| if ocr_result and ocr_result[0] and ocr_result[0][0]: # Check for valid OCR result structure | |
| # The OCR result structure for ocr.ocr is different: it's a list of blocks, each containing [bbox, (text, score)] | |
| # We are interested in the text from the first block's first item | |
| text = ocr_result[0][0][1][0] | |
| text = re.sub(r'[^a-zA-Z0-9]', '', text) | |
| license_plate_text = text | |
| output_message += f"Detected License Plate Text: {license_plate_text}\n" | |
| else: | |
| output_message += "No text found on license plate.\n" | |
| # Put text on the annotated frame (still show on image for visual feedback) | |
| display_text_on_image = license_plate_text if license_plate_text else "No text found" | |
| cv2.putText(annotated_frame, display_text_on_image, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) | |
| except Exception as e: | |
| output_message += f"OCR Error: {e}\n" | |
| print(f"OCR error on region: {e}") | |
| else: | |
| output_message = "No number plate detected in vehicle.\n" | |
| else: | |
| output_message = "No vehicles detected." | |
| # If no vehicles are detected, simply return the original image (or a copy) | |
| # as there's no need to annotate it with a "No vehicles detected" message on the image itself. | |
| annotated_frame = img.copy() # Ensure we always return an image | |
| # Convert back to RGB for display in Gradio | |
| annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB) | |
| return annotated_frame, output_message | |
| iface = gr.Interface( | |
| fn=detect_and_annotate, | |
| inputs=gr.Image(type="pil"), | |
| outputs=[gr.Image(type="numpy", label="Annotated Image"), | |
| gr.Textbox(label="Detection Messages")], # Added Textbox for messages | |
| title="Vehicle & License Plate Detection with OCR", | |
| description="Upload an image. The app will detect vehicles and license plates, perform OCR, and display the result with annotations and messages.", | |
| examples=example_images, | |
| cache_examples=False | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch(share=True) |