alpr-ocr / app.py
Japhari's picture
Update app.py
1476c21 verified
Raw
History Blame Contribute Delete
5.32 kB
import gradio as gr
import cv2
import numpy as np
from PIL import Image
from typing import get_args
from fast_alpr import ALPR
from fast_alpr.default_detector import PlateDetectorModel
from fast_alpr.default_ocr import OcrModel
# Default models
DETECTOR_MODELS = list(get_args(PlateDetectorModel))
OCR_MODELS = list(get_args(OcrModel))
# Put global OCR first
OCR_MODELS.remove("cct-s-v1-global-model")
OCR_MODELS.insert(0, "cct-s-v1-global-model")
def process_image(image, detector_model, ocr_model):
"""
Process an image with ALPR system
Args:
image: PIL Image or numpy array
detector_model: Selected detector model
ocr_model: Selected OCR model
Returns:
tuple: (annotated_image, results_text)
"""
if image is None:
return None, "Please upload an image to continue."
try:
# Convert image to numpy array if it's a PIL Image
if isinstance(image, Image.Image):
img_array = np.array(image.convert("RGB"))
else:
img_array = image
# Initialize ALPR with selected models
alpr = ALPR(detector_model=detector_model, ocr_model=ocr_model)
# Run ALPR on the image
results = alpr.predict(img_array)
# Draw predictions on the image
annotated_img_array = alpr.draw_predictions(img_array)
# Convert back to PIL Image for Gradio
annotated_img = Image.fromarray(annotated_img_array)
# Format results text
if results:
results_text = "**OCR Results:**\n"
for i, result in enumerate(results):
plate_text = result.ocr.text if result.ocr else "N/A"
plate_confidence = result.ocr.confidence if result.ocr else 0.0
results_text += f"{i+1}. Detected Plate: `{plate_text}` with confidence `{plate_confidence:.2f}`\n"
else:
results_text = "No license plate detected."
return annotated_img, results_text
except Exception as e:
return None, f"Error processing image: {str(e)}"
# Create Gradio interface
with gr.Blocks(
title="Automatic License Plate Recognition (ALPR)",
theme=gr.themes.Soft(
primary_hue="green",
secondary_hue="blue",
neutral_hue="slate"
)
) as demo:
gr.Markdown("""
# Automatic License Plate Recognition (ALPR)
An automatic license plate recognition (ALPR) system with customizable detector and OCR models.
This system uses the FastALPR library to detect and recognize license plates in images.
""")
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### Model Selection")
detector_dropdown = gr.Dropdown(
choices=DETECTOR_MODELS,
value=DETECTOR_MODELS[0],
label="Choose Detector Model",
info="Select the model for license plate detection"
)
ocr_dropdown = gr.Dropdown(
choices=OCR_MODELS,
value=OCR_MODELS[0],
label="Choose OCR Model",
info="Select the model for text recognition"
)
gr.Markdown("### Upload Image")
image_input = gr.Image(
label="Upload an image of a vehicle with a license plate",
type="pil",
height=300
)
process_btn = gr.Button(
"Process Image",
variant="primary",
size="lg"
)
with gr.Column(scale=1):
gr.Markdown("### Results")
image_output = gr.Image(
label="Annotated Image with OCR Results",
height=300
)
text_output = gr.Markdown(
label="OCR Results",
value="Upload an image and click 'Process Image' to see results."
)
# Add some examples
gr.Markdown("### Example Images")
gr.Examples(
examples=[],
inputs=image_input,
label="Try with these examples (if available)"
)
# Connect the interface
process_btn.click(
fn=process_image,
inputs=[image_input, detector_dropdown, ocr_dropdown],
outputs=[image_output, text_output]
)
gr.Markdown("""
---
### How to Use
1. **Select Models**: Choose your preferred detector and OCR models from the dropdowns
2. **Upload Image**: Upload an image containing a vehicle with a license plate
3. **Process**: Click the "Process Image" button to run ALPR
4. **View Results**: See the annotated image and OCR text results
### Supported Models
- **Detector Models**: Various YOLO-based detection models
- **OCR Models**: CCT OCR models for text recognition including global and specialized models
### Features
- Real-time license plate detection
- Customizable model selection
- Visual annotations with bounding boxes
- Confidence scores for OCR results
- Support for multiple image formats
""")
# Launch the app
if __name__ == "__main__":
demo.launch(share=False)