Update app.py

app.py

@@ -1,51 +1,88 @@
-# Install necessary libraries
 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
 
-# Step 1: Download the YOLO model weights from your Hugging Face repository
-weights_path = hf_hub_download(repo_id="krishnamishra8848/Nepal-Vehicle-License-Plate-Detection", filename="last.pt")
+# Title for the Streamlit App
+st.title("Nepal Vehicle License Plate and Character Detection")
 
-# Step 2: Load the YOLO model
-model = YOLO(weights_path)
+# Description
+st.write("Upload an image to detect license plates and their characters using two advanced YOLO models.")
 
-# Step 3: Function to process and display results
-def detect_license_plate(image):
-    # Convert the PIL image to a numpy array
-    img = np.array(image)
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+# 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)
 
-    # Perform inference
-    results = model(img)
+    # 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)
 
-    # Draw bounding boxes and confidence scores
+    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, conf in zip(result.boxes.xyxy, result.boxes.conf):
-                x1, y1, x2, y2 = map(int, box)  # Convert to integers
-                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)  # Green rectangle
-                label = f"{conf:.2f}"
-                cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+            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)
 
-    # Convert back to RGB for Streamlit display
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(img)
+    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)
 
-# Streamlit Interface
-st.title("Nepal Vehicle License Plate Detection")
-st.write("Upload an image to detect license plates.")
+    return image
 
-# File uploader
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+# Upload an image file
+uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    # Run detection directly after uploading
-    st.write("Processing...")
+    # Load image
     image = Image.open(uploaded_file)
-    result_image = detect_license_plate(image)
 
-    # Display the detection result
-    st.image(result_image, caption="Detection Results", use_container_width=True)
+    # Display uploaded image
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # 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}:")
+                annotated_plate = detect_characters(cropped_image.copy())
+                st.image(cv2.cvtColor(annotated_plate, cv2.COLOR_BGR2RGB), caption=f"License Plate {idx} with Characters", use_column_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_column_width=True)
+
+    st.success("Processing complete!")