npd / main.py
hostel's picture
commit1
cf18722 verified
Raw
History Blame Contribute Delete
5.93 kB
import cv2
import pandas as pd
import numpy as np
import pytesseract
import os
from ultralytics import YOLO
from datetime import datetime
# Configure Tesseract executable path
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
# Load YOLO model
model = YOLO('best.pt')
# Define polygon area for detection (adjust as needed)
area = [(27, 250), (16, 310), (1015, 310), (992, 250)]
# Track processed numbers to avoid duplicates
processed_numbers = set()
# Initialize video capture
cap = cv2.VideoCapture('mycarplate.mp4') # Change this to your video file path
# Load class list from file
with open("coco1.txt", "r") as file:
class_list = file.read().splitlines()
# Define CSV file and write headers if not already present
csv_file = "car_plate_data_stored.csv"
# Ensure the file has headers if it doesn't already exist
try:
pd.read_csv(csv_file)
except FileNotFoundError:
with open(csv_file, "w") as file:
file.write("ImageFile,Date,Time,Confidence\n")
frame_count = 0
# Function to correct commonly misrecognized characters
def correct_characters(text):
replacements = {
'0': 'O', # Replace '0' with 'O' if detected, to avoid confusion with 'D'
'O': 'D', # Replace 'O' with 'D' in typical license plate context
'I': '1', # Replace 'I' with '1' if detected
'Q': '0' # Replace 'Q' with '0' if detected
}
corrected_text = ''.join(replacements.get(c, c) for c in text)
return corrected_text
# Function to format license plate text
def format_plate_text(text):
text = ''.join(filter(str.isalnum, text))
if len(text) == 10: # Assuming full-length plates are correctly detected
return f"{text[:2]} {text[2:4]} {text[4:6]} {text[6:]}"
return text # If not 10 characters, return as-is
# Deskew function to correct any skewed license plates
def deskew_image(image):
# Check if the image is already grayscale (single channel)
if len(image.shape) == 3:
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
gray = image # It's already grayscale, no need to convert
# Apply Canny edge detection
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
# Use Hough Transform to detect lines
lines = cv2.HoughLines(edges, 1, np.pi / 180, 100)
if lines is not None:
for line in lines:
rho, theta = line[0]
a = np.cos(theta)
b = np.sin(theta)
x0 = a * rho
y0 = b * rho
x1 = int(x0 + 1000 * (-b))
y1 = int(y0 + 1000 * (a))
x2 = int(x0 - 1000 * (-b))
y2 = int(y0 - 1000 * (a))
cv2.line(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
return image
# Processing video frames
while True:
ret, frame = cap.read()
frame_count += 1
# Skip frames to process every third frame for efficiency
if frame_count % 3 != 0:
continue
if not ret:
break
# Resize frame for consistent processing
frame = cv2.resize(frame, (1020, 500))
# Perform YOLO object detection
results = model.predict(frame)
detected_boxes = results[0].boxes.data
detections = pd.DataFrame(detected_boxes).astype("float")
for _, row in detections.iterrows():
x1, y1, x2, y2 = int(row[0]), int(row[1]), int(row[2]), int(row[3])
class_id = int(row[5])
class_name = class_list[class_id]
# Calculate center point of the detected box
cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
# Check if center is within the defined polygon area
if cv2.pointPolygonTest(np.array(area, np.int32), (cx, cy), False) >= 0:
# Crop detected area
crop = frame[y1:y2, x1:x2]
gray_crop = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
# Apply bilateral filter for smoothing, and threshold for clarity
gray_crop = cv2.bilateralFilter(gray_crop, 10, 20, 20)
_, threshold_crop = cv2.threshold(gray_crop, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
# Deskew the image if needed
threshold_crop = deskew_image(threshold_crop)
# Generate unique identifier for the detection
current_date = datetime.now().strftime("%d-%m-%Y")
current_time = datetime.now().strftime("%H:%M:%S")
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
# Save the cropped license plate image
detection_folder = "detected_plates"
if not os.path.exists(detection_folder):
os.makedirs(detection_folder)
image_filename = f"plate_{timestamp}.jpg"
image_path = os.path.join(detection_folder, image_filename)
cv2.imwrite(image_path, crop) # Save original crop
# Save detection information to CSV
if image_filename not in processed_numbers: # Use filename as unique identifier
processed_numbers.add(image_filename)
with open(csv_file, "a") as file:
confidence = float(row[4]) # Get detection confidence
file.write(f"{image_filename},{current_date},{current_time},{confidence:.2f}\n")
# Draw bounding box and display cropped image
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 1)
cv2.imshow('Detected Plate', threshold_crop)
# Draw the detection area on the main frame
cv2.polylines(frame, [np.array(area, np.int32)], True, (255, 0, 0), 2)
cv2.imshow("RGB", frame)
# Break loop on 'ESC' key press
if cv2.waitKey(1) & 0xFF == 27:
break
# Release resources
cap.release()
cv2.destroyAllWindows()