File size: 3,693 Bytes
68da213 0c230ca 68da213 bfb63f6 68da213 bfb63f6 68da213 0c230ca 68da213 bfb63f6 0c230ca 68da213 bfb63f6 68da213 bfb63f6 68da213 bfb63f6 68da213 0c230ca bfb63f6 68da213 bfb63f6 44c2dbe bfb63f6 68da213 a82d3ce 68da213 bfb63f6 44c2dbe 68da213 0c230ca bfb63f6 68da213 0c230ca |
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 |
import cv2
import pytesseract
from PIL import Image
from ultralytics import YOLO
import supervision as sv
import gradio as gr
import os
# Set Tesseract path if needed
pytesseract.pytesseract_cmd = "/usr/bin/tesseract"
# Load the YOLO model (replace with the actual model path)
model_yolo = YOLO('./saved_model.pt')
def process_video(video_path):
"""
Process a video frame-by-frame, detect the first license plate, and save the cropped plate.
Returns the cropped plate path and text.
"""
# Open the video file
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return "Error: Unable to open video file.", None
frame_count = 0
# Ensure the cropped plates directory exists
cropped_dir = "cropped_plates"
os.makedirs(cropped_dir, exist_ok=True)
cropped_plate_path = None
number_plate_text = None
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame_count += 1
print(f"Processing frame {frame_count}...")
# Run YOLO predictions
results = model_yolo.predict(source=frame, save=False)
detections = sv.Detections(
xyxy=results[0].boxes.xyxy.cpu().numpy(),
confidence=results[0].boxes.conf.cpu().numpy(),
class_id=results[0].boxes.cls.cpu().numpy().astype(int)
)
# Process the first detected license plate
if len(detections.xyxy) > 0:
x1, y1, x2, y2 = map(int, detections.xyxy[0]) # Take the first detection
cropped_frame = frame[y1:y2, x1:x2]
pil_image = Image.fromarray(cv2.cvtColor(cropped_frame, cv2.COLOR_BGR2RGB))
try:
# Extract text using Tesseract OCR
number_plate_text = pytesseract.image_to_string(
pil_image, config="--psm 7 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
).strip()
if number_plate_text:
# Save the cropped plate image
cropped_plate_path = os.path.join(cropped_dir, "first_detected_plate.jpg")
cv2.imwrite(cropped_plate_path, cropped_frame)
print(f"Detected Number Plate: {number_plate_text}")
break # Exit loop after detecting the first license plate
except Exception as e:
print(f"Error during OCR: {e}")
cap.release()
if cropped_plate_path is None or number_plate_text is None:
print("No license plate found in the video.")
return cropped_plate_path, number_plate_text
def gradio_video_processor(video):
"""
Gradio wrapper to process video and return the first detected license plate and its text.
"""
video_path = video # Use the file path provided by Gradio
cropped_plate_path, plate_text = process_video(video_path)
if cropped_plate_path and plate_text:
html_result = f"<h3>Detected License Plate</h3>"
html_result += f"<p></p><img src='{cropped_plate_path}' width='300'><br>"
return cropped_plate_path, html_result
else:
return None, "<h3>No license plate detected in the video.</h3>"
# Gradio interface
iface = gr.Interface(
fn=gradio_video_processor,
inputs=gr.Video(label="Upload a Video to detect license plate"),
outputs=[
gr.Image(label="Cropped Plate Image"), # Shows the cropped plate
gr.HTML(label="License Plate Detected") # Shows the extracted text
],
title="License Plate Detection",
description="Upload a video to detect the first license plate and extract its text."
)
if __name__ == "__main__":
iface.launch()
|