Update app.py

app.py

@@ -4,17 +4,17 @@ from ultralytics import YOLO
 import cv2
 import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
+from tensorflow.keras.models import load_model
 import tempfile
 import os
 
 # Title for the Streamlit App
-st.title("Nepal Vehicle License Plate and Character Detection")
+st.title("Nepal Vehicle License Plate and Character Recognition")
 
 # Description
-st.write("Upload an image to detect license plates and their characters using two advanced YOLO models.")
+st.write("Upload an image to detect license plates, segment characters, and recognize each character using advanced YOLO and CNN models.")
 
-# Download YOLO model weights from Hugging Face
+# Download YOLO and CNN model weights from Hugging Face
 @st.cache_resource
 def load_models():
     # Full license plate detection model
@@ -29,10 +29,16 @@ def load_models():
     )
     character_model = YOLO(character_model_path)
 
-    return full_plate_model, character_model
+    # Character recognition model
+    recognition_model_path = hf_hub_download(
+        repo_id="krishnamishra8848/Nepal_Vehicle_License_Plates_Character_Recognisation", filename="model.h5"
+    )
+    recognition_model = load_model(recognition_model_path)
+
+    return full_plate_model, character_model, recognition_model
 
 # Load models
-full_plate_model, character_model = load_models()
+full_plate_model, character_model, recognition_model = load_models()
 
 # Function to detect and crop license plates
 def detect_and_crop_license_plate(image):
@@ -51,16 +57,36 @@ def detect_and_crop_license_plate(image):
 
     return cropped_images, detected_image
 
-# Function to detect characters
-def detect_characters(image):
+# Function to detect and crop characters
+def detect_and_crop_characters(image):
     results = character_model(image)
+    character_crops = []
     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
+                character_crops.append(image[y1:y2, x1:x2])
+    return character_crops
+
+# Function to recognize characters
+def recognize_characters(character_crops):
+    class_labels = [
+        'क', 'को', 'ख', 'ग', 'च', 'ज', 'झ', 'ञ', 'डि', 'त', 'ना', 'प',
+        'प्र', 'ब', 'बा', 'भे', 'म', 'मे', 'य', 'लु', 'सी', 'सु', 'से', 'ह',
+        '०', '१', '२', '३', '४', '५', '६', '७', '८', '९'
+    ]
+    recognized_characters = []
+    for char_crop in character_crops:
+        # Preprocess the cropped character for recognition model
+        resized = cv2.resize(char_crop, (64, 64))
+        normalized = resized / 255.0
+        reshaped = np.expand_dims(normalized, axis=0)  # Add batch dimension
+
+        # Predict the character
+        prediction = recognition_model.predict(reshaped)
+        predicted_class = class_labels[np.argmax(prediction)]
+        recognized_characters.append(predicted_class)
+    return recognized_characters
 
 # Upload an image file
 uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])
@@ -77,15 +103,24 @@ if uploaded_file is not None:
         st.image(cv2.cvtColor(detected_image, cv2.COLOR_BGR2RGB), caption="Detected License Plates", use_container_width=True)
 
         if cropped_plates:
-            st.write(f"Detected {len(cropped_plates)} license plate(s). Showing results:")
+            st.write(f"Detected {len(cropped_plates)} license plate(s).")
             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)
+                st.write(f"Processing License Plate {idx}:")
+
+                # Detect and crop characters
+                character_crops = detect_and_crop_characters(cropped_image)
+
+                if character_crops:
+                    # Recognize characters
+                    recognized_characters = recognize_characters(character_crops)
+
+                    # Show each cropped character and prediction
+                    for i, char_crop in enumerate(character_crops):
+                        st.image(cv2.cvtColor(char_crop, cv2.COLOR_BGR2RGB), caption=f"Character {i+1}")
+                        st.write(f"Predicted Character: {recognized_characters[i]}")
+                else:
+                    st.write("No characters detected in this license plate.")
         else:
-            st.write("No license plates detected. Running character detection on the full image.")
-            annotated_image = detect_characters(detected_image.copy())
-            st.image(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB), caption="Full Image with Characters", use_container_width=True)
+            st.write("No license plates detected.")
 
     st.success("Processing complete!")