Untitled

import cv2
import pytesseract
import numpy as np
import os
import time

# Define officer plates
officer_plates = ["CTN1111", "CTN202"]
output_directory = "detected_plates"
os.makedirs(output_directory, exist_ok=True)

# Dictionary to store detected plates and prevent duplicates
detected_plates = {}

def pre_process_image(image):
    """Convert the image to grayscale, apply filters, and detect edges."""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    edged = cv2.Canny(gray, 50, 150)

    # Dilate the edges to close gaps in contours
    kernel = np.ones((3, 3), np.uint8)
    edged = cv2.dilate(edged, kernel, iterations=1)
    
    return edged

def color_filter(image):
    """Filter image based on license plate common colors (black and white)."""
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    
    # Define range for white color
    lower_white = np.array([0, 0, 180], dtype=np.uint8)
    upper_white = np.array([255, 55, 255], dtype=np.uint8)
    
    # Define range for black color
    lower_black = np.array([0, 0, 0], dtype=np.uint8)
    upper_black = np.array([180, 255, 70], dtype=np.uint8)

    white_mask = cv2.inRange(hsv, lower_white, upper_white)
    black_mask = cv2.inRange(hsv, lower_black, upper_black)

    return cv2.bitwise_or(white_mask, black_mask)

def detect_plate_contour(edged, original_image):
    """Detect license plate contour based on aspect ratio and size filtering."""
    cnts, _ = cv2.findContours(edged.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:20]  # Increase number of contours to consider

    for c in cnts:
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.018 * peri, True)

        if len(approx) == 4:  # Rectangle approximation
            x, y, w, h = cv2.boundingRect(c)
            aspect_ratio = float(w) / h

            # Filter by size and aspect ratio
            if 100 < w < 400 and 25 < h < 150 and 2 < aspect_ratio < 6:
                return approx
    return None

def save_image_with_text(cropped_image, plate_text, status):
    """Save the cropped image with status text."""
    global detected_plates

    # Avoid duplicate saves
    if plate_text in detected_plates:
        return

    detected_plates[plate_text] = True

    # Ensure the image has 3 channels
    if cropped_image.shape[2] == 4:
        cropped_image = cv2.cvtColor(cropped_image, cv2.COLOR_BGRA2BGR)

    # Add extra space for text
    height, width, _ = cropped_image.shape
    new_image = np.zeros((height + 30, width, 3), dtype=np.uint8)
    new_image[0:height, 0:width] = cropped_image

    # Add the status text below the image
    cv2.putText(new_image, f"Status: {status}", (10, height + 20), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
    
    # Save the image
    timestamp = int(time.time())
    output_path = os.path.join(output_directory, f"plate_{plate_text}_{status}_{timestamp}.jpg")
    cv2.imwrite(output_path, new_image)

def perform_ocr(cropped_image):
    """Perform OCR and determine plate status based on confidence."""
    config = '--psm 7 --oem 3'
    data = pytesseract.image_to_data(cropped_image, config=config, output_type=pytesseract.Output.DICT)

    # Combine text and check OCR confidence
    text = ''.join(data['text']).strip()

    if len(text) > 0 and all([float(conf) > 60 for conf in data['conf'] if conf != '-1']):
        # Check if it's a valid plate length
        if 6 <= len(text) <= 7:
            status = "Staff" if text in officer_plates else "Outsider"
            print(f"Detected Plate: {text} | Status: {status}")
            save_image_with_text(cropped_image, text, status)

def process_frame(frame):
    """Process each frame for license plate detection and OCR."""
    color_mask = color_filter(frame)
    edged = pre_process_image(frame)

    # Apply color filter mask to focus on relevant areas
    edged = cv2.bitwise_and(edged, edged, mask=color_mask)

    plate_contour = detect_plate_contour(edged, frame)

    if plate_contour is not None:
        cv2.drawContours(frame, [plate_contour], -1, (0, 255, 0), 3)

        mask = np.zeros(edged.shape, np.uint8)
        new_image = cv2.drawContours(mask, [plate_contour], 0, 255, -1)
        new_image = cv2.bitwise_and(frame, frame, mask=mask)

        # Crop the plate area
        (x, y) = np.where(mask == 255)
        (topx, topy) = (np.min(x), np.min(y))
        (bottomx, bottomy) = (np.max(x), np.max(y))
        cropped_image = frame[topx:bottomx+1, topy:bottomy+1]

        # Perform OCR on the cropped license plate area
        perform_ocr(cropped_image)

# Open the video file
video_path = 'your_video.mp4'  # Replace with your MP4 file path
cap = cv2.VideoCapture(video_path)

frame_count = 0

# Main loop to capture frames and process
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break  # Exit the loop when the video ends

    # Process every frame or adjust the interval as needed
    if frame_count % 1 == 0:
        process_frame(frame)

    cv2.imshow("Video Preview", frame)
    frame_count += 1

    # Break the loop with 'q'
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()