Sourudra's picture
Update app.py
9b1dc68 verified
import streamlit as st
import cv2
import numpy as np
import easyocr
from ultralytics import YOLO
# Title of the app with enhanced styling
st.markdown("""
<style>
.stApp {
background-image: url("https://i.postimg.cc/cLpB502d/black-car.png");
background-size: cover;
background-position: center center;
}
h1 {
color: white;
font-size: 3em;
text-shadow: 3px 3px 5px rgba(0, 0, 0, 0.7);
text-align: center;
}
.sidebar .sidebar-content {
background-color: rgba(0, 0, 0, 0.6);
color: white;
}
.stButton>button {
background-color: #4CAF50;
color: white;
font-size: 18px;
padding: 10px 24px;
border-radius: 10px;
border: none;
cursor: pointer;
}
.stButton>button:hover {
background-color: #45a049;
}
.stTextInput>div>div>input {
background-color: #ffffff;
color: black;
}
.stTextInput>div>label {
color: white;
}
.stSlider>div>label {
color: white;
}
.stMarkdown {
color: white;
}
.stInfo {
background-color: rgba(0, 0, 0, 0.6);
color: white;
}
</style>
""", unsafe_allow_html=True)
# Title of the app
st.title("License Plate Recognition System")
# 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)
# Create columns to display the uploaded image, cropped image, and extracted text side by side
c1, c2 = st.columns(2)
with c1:
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)
with c2:
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)
with c2:
st.write(f"**Extracted Text (Plate {i+1}):** {detected_text}")
st.write(f"**Confidence Score:** {result.boxes.conf.cpu().numpy()[i]:.2f}")
# Sidebar input for file upload
uploaded_file = st.file_uploader("Upload an Image or Video", type=["mp4", "avi", "mov", "jpg", "jpeg", "png"])
if uploaded_file is not None:
# 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)
st.markdown("---")
st.info("Please use images as input, because processing videos with the CPU may result in longer times.")
st.info("This application uses Fine Tuned YOLOv8 for detection and EasyOCR for text recognition.")