File size: 5,036 Bytes
f93376b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 |
import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image
import os
# Set the title of the Streamlit app
st.title("License Plate Detection using YOLO")
# Ensure the temp directory exists
if not os.path.exists("temp"):
os.makedirs("temp")
# Allow users to upload images or videos
uploaded_file = st.file_uploader("Upload an image or video",
type=["jpg", "jpeg", "png", "bmp", "mp4", "avi", "mov", "mkv"])
# Load YOLO model
try:
model = YOLO('best.pt') # Use the relative path to your trained YOLO model
except Exception as e:
st.error(f"Error loading YOLO model: {e}")
def predict_and_save_image(path_test_car, output_image_path):
"""
Predicts and saves the bounding boxes on the given test image using the trained YOLO model.
Parameters:
path_test_car (str): Path to the test image file.
output_image_path (str): Path to save the output image file.
Returns:
str: The path to the saved output image file.
"""
try:
results = model.predict(path_test_car, device='cpu')
image = cv2.imread(path_test_car)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
confidence = box.conf[0]
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image, f'{confidence * 100:.2f}%', (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imwrite(output_image_path, image)
return output_image_path
except Exception as e:
st.error(f"Error processing image: {e}")
return None
def predict_and_plot_video(video_path, output_path):
"""
Predicts and saves the bounding boxes on the given test video using the trained YOLO model.
Parameters:
video_path (str): Path to the test video file.
output_path (str): Path to save the output video file.
Returns:
str: The path to the saved output video file.
"""
try:
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
st.error(f"Error opening video file: {video_path}")
return None
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
results = model.predict(rgb_frame, device='cpu')
for result in results:
for box in result.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
confidence = box.conf[0]
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f'{confidence * 100:.2f}%', (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
out.write(frame)
cap.release()
out.release()
return output_path
except Exception as e:
st.error(f"Error processing video: {e}")
return None
def process_media(input_path, output_path):
"""
Processes the uploaded media file (image or video) and returns the path to the saved output file.
Parameters:
input_path (str): Path to the input media file.
output_path (str): Path to save the output media file.
Returns:
str: The path to the saved output media file.
"""
file_extension = os.path.splitext(input_path)[1].lower()
if file_extension in ['.mp4', '.avi', '.mov', '.mkv']:
return predict_and_plot_video(input_path, output_path)
elif file_extension in ['.jpg', '.jpeg', '.png', '.bmp']:
return predict_and_save_image(input_path, output_path)
else:
st.error(f"Unsupported file type: {file_extension}")
return None
if uploaded_file is not None:
input_path = os.path.join("temp", uploaded_file.name)
output_path = os.path.join("temp", f"output_{uploaded_file.name}")
try:
with open(input_path, "wb") as f:
f.write(uploaded_file.getbuffer())
st.write("Processing...")
result_path = process_media(input_path, output_path)
if result_path:
if input_path.endswith(('.mp4', '.avi', '.mov', '.mkv')):
video_file = open(result_path, 'rb')
video_bytes = video_file.read()
st.video(video_bytes)
else:
image = Image.open(result_path)
st.image(image)
except Exception as e:
st.error(f"Error uploading or processing file: {e}") |