import numpy as np import pandas as pd import streamlit as st from ultralytics import YOLO import cv2 import pytesseract from pytesseract import Output import os import re import shutil import tempfile st.set_page_config( page_title = "YOLO Car Licence Plate Image and Video Processing", page_icon = ":car:", initial_sidebar_state = "expanded", ) st.title('YOLO Car Licence Plate :red[Image and Video Processing]') pytesseract.pytesseract.tesseract_cmd = None # search for tesseract binary in path @st.cache_resource def find_tesseract_binary() -> str: return shutil.which("tesseract") # set tesseract binary path pytesseract.pytesseract.tesseract_cmd = find_tesseract_binary() if not pytesseract.pytesseract.tesseract_cmd: st.error("Tesseract binary not found in PATH. Please install Tesseract.") # pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe" # (only for local) # 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"]) def remove_non_alphanum(text): return re.sub(r'[^a-zA-Z0-9]', ' ', text) # Load YOLO model try: model = YOLO('best.pt') except Exception as e: st.error(f"Error loading YOLO model: {e}") def predict_and_save_image(path_test_car:str, output_image_path:str)-> str: """ 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) gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) 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:.1f}%', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 178, 102), 2, cv2.LINE_AA) roi = gray_image[y1:y2, x1:x2] # Perform OCR on the cropped image text = pytesseract.image_to_string(roi,lang='eng', config = r'--oem 3 --psm 6') text = remove_non_alphanum(text) cv2.putText(image, f'{text}', (x1 , y1 + 2 * (y2 - y1)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (51, 255, 255), 2, cv2.LINE_AA) st.code(f"License Number: {text}", language='text') image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # Ensure the directory exists before saving os.makedirs(os.path.dirname(output_image_path), exist_ok= True) # Save the image 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:str, output_path:str)-> str: """ 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(*'h264') output_dir = os.path.dirname(output_path) if not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True) 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) gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) 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:.1f}%', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 178, 102), 2, cv2.LINE_AA) roi = gray_frame[y1:y2, x1:x2] # Perform OCR on the cropped image text = pytesseract.image_to_string(roi, lang='eng', config = r'--oem 3 --psm 6') text = remove_non_alphanum(text) cv2.putText(frame, f'{text}', (x1 , y1 + 2 * (y2 - y1)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (51, 255, 255), 2, cv2.LINE_AA) 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:str, output_path:str) -> str: """ 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 temp_directory = 'temp' if not os.path.exists(temp_directory): os.makedirs(temp_directory) if st.button("Proceed"): 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()) with st.spinner('Processing...'): result_path = process_media(input_path, output_path) if result_path: if input_path.endswith(('.h264','.mp4', '.avi', '.mov', '.mkv')): # video_file = open(result_path, 'rb') with open(result_path, "rb") as video_file: video_bytes = video_file.read() st.video(video_bytes) else: st.image(result_path) except Exception as e: st.error(f"Error uploading or processing file: {e}")