Update app.py

app.py

@@ -1,114 +1,132 @@
 import streamlit as st
-import pickle
-import numpy as np
 import cv2
-from PIL import Image
-
-# Load the trained model and other necessary files
-model = pickle.load(open('artifacts/license_plate_model.pkl', 'rb'))
-
-# Function to load custom CSS
-def local_css(file_name):
-    with open(file_name) as f:
-        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
-
-# Apply custom CSS for styling
-local_css("style.css")
-
-# Page configuration
-st.set_page_config(page_title='License Plate Detection', layout='centered')
-
-# Add a fancy header with emojis
-st.markdown("""
-    <div class="glass">
-        <h1>🚗 License Plate Detection 🚗</h1>
-        <p>Upload an image to detect license plates</p>
-    </div>
-""", unsafe_allow_html=True)
-
-# File uploader for image input
-st.markdown("""
-    <div class="glass">
-        <p>📸 Upload a car image for license plate detection</p>
-    </div>
-""", unsafe_allow_html=True)
-
-uploaded_file = st.file_uploader("Choose a file", type=["jpg", "jpeg", "png"])
+import numpy as np
+import easyocr
+from ultralytics import YOLO
+
+# Title of the app
+st.title("License Plate Recognition System🚗")
+
+# Add external CSS for custom styling
+with open('style.css') as f:
+    st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
+
+# Load the YOLO model for license plate detection
+@st.cache_resource
+def load_yolo_model():
+    model_path = "best.pt"  # Replace with your model file
+    model = YOLO(model_path)
+    return model
+
+# Load EasyOCR reader
+@st.cache_resource
+def load_easyocr_reader():
+    return easyocr.Reader(['en'], gpu=False)
+
+# Initialize models
+yolo_model = load_yolo_model()
+ocr_reader = load_easyocr_reader()
+
+# Function to process image and detect license plates
+def process_image(image, confidence_threshold=0.5):
+    # Perform license plate detection
+    results = yolo_model(image, conf=confidence_threshold)
+    annotated_image = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
+    st.image(annotated_image, caption="Detected License Plate(s)", use_container_width=True)
+    
+    # Loop through detections and perform OCR
+    for result in results:
+        boxes = result.boxes.xyxy.cpu().numpy().astype(int)
+        if len(boxes) == 0:
+            st.warning("No license plate detected!")
+            return
+        for i, box in enumerate(boxes):
+            x1, y1, x2, y2 = box
+            cropped_image = image[y1:y2, x1:x2]
+            cropped_image_rgb = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)
+            st.image(cropped_image_rgb, caption=f"Cropped License Plate {i+1}", use_container_width=True)
+            
+            # Perform OCR on the cropped image
+            text_results = ocr_reader.readtext(cropped_image_rgb, detail=0)
+            detected_text = " ".join(text_results)
+            st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
+            st.write(f"**Confidence Score:** {result.boxes.conf.cpu().numpy()[i]:.2f}")
+
+# Function to process video and detect license plates
+def process_video(video_path, confidence_threshold=0.5, output_path="output_video.mp4"):
+    # Open the video file
+    cap = cv2.VideoCapture(video_path)
+
+    # Get video frame dimensions
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Create VideoWriter object to save the output video
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for mp4
+    out = cv2.VideoWriter(output_path, fourcc, 20.0, (frame_width, frame_height))  # 20 FPS
+
+    if not cap.isOpened():
+        st.error("Error opening video stream or file")
+        return
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        results = yolo_model(frame, conf=confidence_threshold)
+        annotated_frame = cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)
+
+        # Loop through detections and perform OCR
+        for result in results:
+            boxes = result.boxes.xyxy.cpu().numpy().astype(int)
+            for i, box in enumerate(boxes):
+                x1, y1, x2, y2 = box
+                cropped_plate = frame[y1:y2, x1:x2]
+                cropped_rgb = cv2.cvtColor(cropped_plate, cv2.COLOR_BGR2RGB)
+
+                # Perform OCR on the cropped image
+                text_results = ocr_reader.readtext(cropped_rgb, detail=0)
+                detected_text = " ".join(text_results)
+
+                # Optionally add detected text on the annotated frame
+                cv2.putText(annotated_frame, detected_text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+
+        # Write the annotated frame to the output video
+        out.write(cv2.cvtColor(annotated_frame, cv2.COLOR_RGB2BGR))
+
+    cap.release()
+    out.release()
+
+    st.success(f"Video processing complete. Output video saved to {output_path}.")
+
+    # Provide a download link for the processed video
+    with open(output_path, "rb") as f:
+        st.download_button(label="Download Processed Video", data=f, file_name=output_path)
+
+# Sidebar input for file upload
+uploaded_file = st.file_uploader("Upload an Image or Video", type=["mp4", "avi", "mov", "jpg", "jpeg", "png"])
 
-# If image is uploaded, process the image and make predictions
 if uploaded_file is not None:
-    # Load the image using PIL
-    image = Image.open(uploaded_file)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    # Convert image to OpenCV format (for model processing)
-    img_array = np.array(image)
-    img_gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
-
-    # Use the model to detect license plate (model predictions are made here)
-    # Assuming your model is a classifier or detector for the license plate
-    plate_detected = model.predict(img_gray)  # Modify this depending on your model
-
-    if plate_detected:
-        st.markdown("""
-            <div class="glass">
-                <p>✔️ License Plate Detected!</p>
-            </div>
-        """, unsafe_allow_html=True)
-    else:
-        st.markdown("""
-            <div class="glass">
-                <p>❌ No License Plate Detected!</p>
-            </div>
-        """, unsafe_allow_html=True)
-
-# Apply custom button styling to Streamlit's default button using CSS
-st.markdown("""
-    <style>
-        .stButton > button {
-            background: linear-gradient(135deg, #8B5E3C, #B8860B); /* Warm vintage gold-brown gradient */
-            border: none;
-            color: white; /* White text */
-            padding: 12px 24px;
-            text-align: center;
-            text-decoration: none;
-            display: inline-block;
-            font-size: 16px;
-            margin: 10px 5px;
-            cursor: pointer;
-            border-radius: 12px; /* Rounded corners for a soft vintage touch */
-            box-shadow: 0 6px 12px rgba(0, 0, 0, 0.25); /* Soft, deep shadow for antique feel */
-            font-family: 'Georgia', serif; /* Classic serif font for the button */
-            text-transform: uppercase; /* Uppercase text for a formal touch */
-        }
-
-        .stButton > button:hover {
-            background: linear-gradient(135deg, #704214, #B8860B); /* Slightly darker gold-brown */
-            box-shadow: 0 8px 16px rgba(0, 0, 0, 0.3); /* Stronger shadow for depth on hover */
-            transform: translateY(-2px); /* Lift effect */
-        }
-    </style>
-""", unsafe_allow_html=True)
-
-# Optionally, add more interactivity or styling to the button
-if st.button('Detect License Plate'):
-    # Trigger the license plate detection when the button is pressed
-    pass  # Include any additional logic here, if necessary
-
-# Glassmorphism styling for text elements
-st.markdown("""
-    <style>
-        h1, h2, h3, h4 {
-            font-family: 'Georgia', serif;
-            color: #ffffff;
-            text-align: center;
-            text-shadow: 1px 1px 2px rgba(0, 0, 0, 0.5);
-        }
-
-        p {
-            font-family: 'Georgia', serif;
-            color: #ffffff;
-            text-align: center;
-        }
-    </style>
-""", unsafe_allow_html=True)
+    # Check if it's an image or video
+    file_type = uploaded_file.type
+
+    if file_type.startswith("image"):
+        # Read and process the image
+        image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
+        confidence_threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, 0.01)
+        process_image(image, confidence_threshold)
+
+    elif file_type.startswith("video"):
+        # Save the uploaded video to a temporary file
+        video_bytes = uploaded_file.read()
+        video_path = "/tmp/uploaded_video.mp4"
+        with open(video_path, "wb") as f:
+            f.write(video_bytes)
+
+        # Process the video and save the output
+        output_path = "/tmp/output_video.mp4"
+        process_video(video_path, confidence_threshold=0.5, output_path=output_path)
+
+st.markdown("---")
+st.info("This application uses Fine Tuned YOLOv8 for detection and EasyOCR for text recognition.")