HiralShah62's picture
Update detect.py
aa45af5 verified
import streamlit as st
from ultralytics import YOLO
import cv2
import easyocr
import numpy as np
import pandas as pd
from PIL import Image
import tempfile
@st.cache_resource
def load_model():
model = YOLO('yolo11n-custom.pt')
model.fuse()
return model
model = load_model()
reader = easyocr.Reader(['en'])
def detect_license_plate(image):
results = model.predict(image, conf=0.15, iou=0.3, classes=[0])
plate_texts = []
img_array = np.array(image)
# img = cv2.imread(image_path)
img = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_height, img_width, _ = img.shape
for result in results:
for bbox in result.boxes.xyxy:
x1, y1, x2, y2 = map(int, bbox.tolist())
plate = img[int(y1):int(y2), int(x1):int(x2)]
scale=2
height, width = plate.shape[:2]
plate = cv2.resize(plate, (width * scale, height * scale), interpolation=cv2.INTER_CUBIC)
lab = cv2.cvtColor(plate, cv2.COLOR_RGB2LAB)
l, a, b = cv2.split(lab)
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
l = clahe.apply(l)
plate = cv2.merge((l, a, b))
plate = cv2.cvtColor(plate, cv2.COLOR_LAB2RGB)
text = reader.readtext(plate, detail=0, allowlist="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-")
text = " ".join(text).upper()
text_scale = max(1, width / 250)
thickness = max(2, width // 200)
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), thickness)
(text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, text_scale, thickness)
text_x = x1 + (width - text_width) // 2 # Centered horizontally
text_y = y1 - 10 if y1 > 50 else y2 + text_height + 20 # Above unless too high
text_box_y1 = text_y - text_height - 5
text_box_y2 = text_y + 5
cv2.rectangle(img, (text_x - 8, text_box_y1 - 3), (text_x + text_width + 8, text_box_y2 + 3), (0, 0, 0), -1)
cv2.rectangle(img, (text_x - 5, text_box_y1), (text_x + text_width + 5, text_box_y2), (255, 255, 255), -1)
cv2.putText(img, text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, text_scale, (0, 0, 0), thickness)
plate_texts.append(text)
image = img
return image, plate_texts
st.title("🚘 Real-Time License Plate Detection", anchor=False)
st.write("For better license plate detection, ensure you use high-quality images. If detection is unclear, try enhancing the image first. Use the Refine Image for Detection tool.")
st.write("Upload an image, upload a video, or use your webcam for real-time license plate detection.")
option = st.radio("Choose Input Source:", ("Upload Image", "Upload Video", "Webcam"), horizontal=True )
if option == "Upload Image":
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
if uploaded_file:
img = Image.open(uploaded_file)
st.write("Processing...")
processed_img, plate_texts = detect_license_plate(img)
st.image(processed_img, caption="Detected Plates Image", use_container_width=True)
st.write("**Detected License Plates:**")
if plate_texts:
plates = pd.DataFrame({"License Plate": plate_texts})
plates.index = plates.index + 1
st.table(plates)
else:
st.write("No license plates detected.")
elif option == "Upload Video":
uploaded_video = st.file_uploader("Choose a video...", type=["mp4", "avi", "mov"])
if uploaded_video is not None:
st.write("Processing video...")
tfile = tempfile.NamedTemporaryFile(delete=False)
tfile.write(uploaded_video.read())
cap = cv2.VideoCapture(tfile.name)
frame_placeholder = st.empty()
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
processed_frame, plate_texts = detect_license_plate(frame)
frame_placeholder.image(processed_frame, caption="Detected Plates Video", use_container_width=True)
cap.release()
elif option == "Webcam":
if "running" not in st.session_state:
st.session_state.running = True
if st.button("Stop"):
st.session_state.running = False
st.write("Starting Webcam... Press **Stop** to end.")
cap = cv2.VideoCapture(0)
frame_placeholder = st.empty()
while cap.isOpened():
ret, frame = cap.read()
if not ret:
st.warning("Failed to capture webcam feed.")
break
processed_frame, plate_texts = detect_license_plate(frame)
frame_placeholder.image(processed_frame, channels="BGR", caption="Webcam Feed", use_container_width=True)
if not st.session_state.running:
break
cap.release()