app.py

import streamlit as st
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import tempfile
import os

# Title for the Streamlit App
st.title("Nepal Vehicle License Plate and Character Detection")

# Description
st.write("Upload an image to detect license plates and their characters using two advanced YOLO models.")

# Download YOLO model weights from Hugging Face
@st.cache_resource
def load_models():
    # Full license plate detection model
    full_plate_model_path = hf_hub_download(
        repo_id="krishnamishra8848/Nepal-Vehicle-License-Plate-Detection", filename="last.pt"
    )
    full_plate_model = YOLO(full_plate_model_path)

    # Character detection model
    character_model_path = hf_hub_download(
        repo_id="krishnamishra8848/Nepal_Vehicle_License_Plates_Detection_Version3", filename="best.pt"
    )
    character_model = YOLO(character_model_path)

    return full_plate_model, character_model

# Load models
full_plate_model, character_model = load_models()

# Function to detect and crop license plates
def detect_and_crop_license_plate(image):
    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    results = full_plate_model(img_bgr)
    
    cropped_images = []
    for result in results:
        if hasattr(result, 'boxes') and result.boxes is not None:
            for box in result.boxes.xyxy:
                x1, y1, x2, y2 = map(int, box)
                cropped_image = img_bgr[y1:y2, x1:x2]
                cropped_images.append(cropped_image)

    return cropped_images, img_bgr

# Function to detect characters
def detect_characters(image):
    results = character_model(image)
    for result in results:
        if hasattr(result, 'boxes') and result.boxes is not None:
            for box in result.boxes.xyxy:
                x1, y1, x2, y2 = map(int, box)
                cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return image

# Upload an image file
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    # Load image
    image = Image.open(uploaded_file)

    # Detect license plates and characters
    with st.spinner("Processing image..."):
        cropped_plates, original_image = detect_and_crop_license_plate(image)

        if cropped_plates:
            st.write(f"Detected {len(cropped_plates)} license plate(s). Showing results:")
            for idx, cropped_image in enumerate(cropped_plates, 1):
                st.write(f"License Plate {idx}:")
                # Pass cropped license plate to character detection
                annotated_plate = detect_characters(cropped_image.copy())
                st.image(cv2.cvtColor(annotated_plate, cv2.COLOR_BGR2RGB), caption=f"License Plate {idx} with Characters", use_container_width=True)
        else:
            st.write("No license plates detected. Running character detection on the full image.")
            annotated_image = detect_characters(original_image.copy())
            st.image(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB), caption="Full Image with Characters", use_container_width=True)

    st.success("Processing complete!")