Add some features

app.py

@@ -4,8 +4,8 @@ import torchvision.transforms as T
 import cv2
 import numpy as np
 from PIL import Image
+import easyocr
 
-# Custom CSS for amazing UI
 st.markdown("""
     <style>
     .main {
@@ -53,7 +53,6 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
-# Load model with explicit `weights_only=False`
 @st.cache_resource
 def load_model():
     try:
@@ -66,9 +65,15 @@ def load_model():
 
 model = load_model()
 
-def detect_numberplate(image):
+@st.cache_resource
+def load_ocr_reader():
+    return easyocr.Reader(['en'], gpu=False)
+
+reader = load_ocr_reader()
+
+def detect_and_crop_numberplate(image):
     if model is None:
-        return image  # Return original image if model fails to load
+        return image, None
     
     transform = T.ToTensor()
     img_tensor = transform(image).unsqueeze(0)
@@ -79,24 +84,31 @@ def detect_numberplate(image):
     boxes = predictions[0]['boxes']
     scores = predictions[0]['scores']
 
-    image_np = np.array(image)  # Convert to numpy array
+    image_np = np.array(image)
+    cropped_plate = None
     for box, score in zip(boxes, scores):
-        if score > 0.5:  # Confidence threshold
+        if score > 0.5:
             x1, y1, x2, y2 = map(int, box.tolist())
-            # Fancy rectangle with gradient effect
             cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 153), 3)
-            # Add score label
             cv2.putText(image_np, f"{score:.2f}", (x1, y1-10), 
-                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
+            cropped_plate = image.crop((x1, y1, x2, y2))
+            break
 
-    return Image.fromarray(image_np)
+    return Image.fromarray(image_np), cropped_plate
+
+def extract_text(cropped_plate):
+    if cropped_plate is None:
+        return "No number plate detected"
+    cropped_np = np.array(cropped_plate)
+    result = reader.readtext(cropped_np)
+    text = " ".join([res[1] for res in result])
+    return text if text else "Unable to read text"
 
-# Streamlit UI
 st.markdown('<div class="main">', unsafe_allow_html=True)
 st.markdown('<p class="title">🚗 Number Plate Wizard</p>', unsafe_allow_html=True)
 st.markdown('<p class="subtitle">Unveil the magic of AI-powered detection!</p>', unsafe_allow_html=True)
 
-# File uploader with a cool design
 with st.container():
     st.markdown('<div class="upload-box">', unsafe_allow_html=True)
     uploaded_file = st.file_uploader("📸 Drop your image here", type=["jpg", "png", "jpeg"], 
@@ -106,12 +118,11 @@ with st.container():
 if uploaded_file is not None:
     image = Image.open(uploaded_file).convert("RGB")
     
-    # Fancy button with spinner
     if st.button("✨ Detect Number Plate"):
         with st.spinner("🔮 Casting detection spell..."):
-            result_image = detect_numberplate(image)
+            result_image, cropped_plate = detect_and_crop_numberplate(image)
+            detected_text = extract_text(cropped_plate)
         
-        # Display results with a sleek layout
         st.markdown("<h3 style='color: #ffffff; text-align: center;'>🎉 Detection Complete!</h3>", unsafe_allow_html=True)
         col1, col2 = st.columns(2)
         with col1:
@@ -121,10 +132,10 @@ if uploaded_file is not None:
             st.image(result_image, caption="Magic Revealed", use_container_width=True)
             st.markdown('<p class="image-caption">Number plate spotted!</p>', unsafe_allow_html=True)
         
-        # Fun success message
+        st.markdown(f"<h4 style='color: #ffffff; text-align: center;'>Detected Text: {detected_text}</h4>", unsafe_allow_html=True)
+        
         st.success("Boom! The number plate has been summoned! 🚘")
 
-# Footer
 st.markdown('<p style="color: #d1e8ff; text-align: center; margin-top: 20px;">Powered by Abhijeet Singh</p>', 
             unsafe_allow_html=True)
 st.markdown('</div>', unsafe_allow_html=True)