import cv2
import easyocr
from ultralytics import YOLO

# 1. Initialize the YOLO detector and EasyOCR Reader
# For custom plate detectors, replace with your 'best.pt' file paths
#model = YOLO("yolov8n.pt")  
model = YOLO("anpr_v8.pt")
#model = YOLO("best.pt")

reader = easyocr.Reader(['en'], gpu=False) # Pi runs on CPU inference

def perform_anpr(frame):
    # 2. Run object detection on the frame
    # If using custom license plate weights, filter for your plate class instead
    results = model(frame, verbose=False)
    
    for result in results:
        boxes = result.boxes
        for box in boxes:
            # Check if detected object is a car (class 2), bus (5), or truck (7)
            # Skip this check if you are using a dedicated license plate model
            cls = int(box.cls[0])
            print(cls)
            if cls not in [0, 2, 5, 7]:
                continue
                
            # Get bounding box coordinates
            xyxy = box.xyxy[0].cpu().numpy()
            x1, y1, x2, y2 = map(int, xyxy)
            
            # 3. Crop the detected vehicle region
            cropped_vehicle = frame[y1:y2, x1:x2]
            
            if cropped_vehicle.size == 0:
                continue
                
            # 4. Perform OCR text extraction on the cropped region
            ocr_results = reader.readtext(cropped_vehicle)
            
            for (bbox, text, prob) in ocr_results:
                if prob > 0.5:  # Confidence threshold for the text
                    # Clean the read text string
                    clean_text = text.upper().strip()
                    
                    # 5. Draw visual boundary markers on the original image frame
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(frame, clean_text, (x1, y1 - 10), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                    
                    print(f"Detected Plate Text: {clean_text} (Confidence: {prob:.2f})")
                    
    return frame

def main():
    # 6. Capture live video feed from the Pi Camera or USB Camera
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    print("Starting ANPR Stream. Press 'q' to stop...")

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        # Process the image frame through our pipeline
        processed_frame = perform_anpr(frame)
        
        # Stream output frame visually
        cv2.imshow("Raspberry Pi ANPR", processed_frame)
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()