app.py

import os
import time
import cv2  # OpenCV for image handling
import easyocr  # For OCR
import torch  # PyTorch (dependency for YOLO/EasyOCR)
import numpy as np  # For numerical operations
from ultralytics import YOLO  # YOLO model from ultralytics
from flask import Flask, request, render_template, redirect, url_for, flash
from werkzeug.utils import secure_filename
import traceback # To print full error tracebacks
import re # Import regular expressions for post-processing

# --- Configuration ---
UPLOAD_FOLDER = os.path.join('static', 'uploads')
RESULT_FOLDER = os.path.join('static', 'results')
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

weights_path = 'best.pt' # Assumes best.pt is in the same folder as app.py
PLATE_CLASS_NAME = 'license_plate'


app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['RESULT_FOLDER'] = RESULT_FOLDER
app.secret_key = 'super secret key' # Needed for flashing messages

# --- Create Folders if they don't exist ---
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULT_FOLDER, exist_ok=True)

# --- Load Models (Load only once when the app starts) ---
reader = None
model = None
model_class_names = {} # Initialize empty

print("Loading EasyOCR Reader...")
try:
    # Use gpu=True if you have a CUDA-enabled GPU and compatible PyTorch/CUDA
    model_dir = os.path.join('.', 'easyocr_models')
    os.makedirs(model_dir, exist_ok=True)
    reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available(), model_storage_directory=model_dir)
    print("EasyOCR Reader loaded successfully.")
except Exception as e:
    print(f"Error loading EasyOCR Reader: {e}. OCR will not work.")
    print(traceback.format_exc()) # Print full traceback

print("Loading YOLO Model...")
try:
    if not os.path.exists(weights_path):
        raise FileNotFoundError(f"YOLO weights file not found at {weights_path}. Please make sure 'best.pt' is in the same directory as app.py or update the path.")
    model = YOLO(weights_path)
    print("YOLO Model loaded successfully.")
    # Store class names if model loaded
    if model:
        model_class_names = model.names
        print("Class names:", model_class_names) # Print class names {index: 'name'}
except Exception as e:
    print(f"Error loading YOLO model from {weights_path}: {e}. Detection will not work.")
    print(traceback.format_exc()) # Print full traceback

# --- Helper Functions ---
def allowed_file(filename):
    """Checks if the filename has an allowed extension."""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def perform_ocr_on_image(img_crop, allowlist='0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'):
    """Performs OCR on a given image crop (NumPy array)."""
    if reader is None:
        print("EasyOCR Reader not loaded. Cannot perform OCR.")
        return "OCR Error"

    try:
        # Convert to grayscale
        gray_img = cv2.cvtColor(img_crop, cv2.COLOR_BGR2GRAY)

        # Resize for consistent input size
        target_height = 64  # You can adjust this value
        scale_ratio = target_height / gray_img.shape[0]
        target_width = int(gray_img.shape[1] * scale_ratio)
        resized_img = cv2.resize(gray_img, (target_width, target_height))

        # Basic preprocessing
        resized_img = cv2.equalizeHist(resized_img)  # Enhance contrast
        
        # Perform OCR
        results = reader.readtext(resized_img, allowlist=allowlist)
        
        # Process results
        if not results:
            return "No text found"
            
        # Take the result with highest confidence
        text = results[0][1]  # Get the text part
        
        # Convert to uppercase and remove any non-alphanumeric characters
        text = text.upper()
        text = re.sub(r'[^A-Z0-9]', '', text)
        
        return text if text else "No text found"


        print(f"Final combined text after post-processing: '{plate_text}'")

        return plate_text if plate_text else "No text found"

    except Exception as e:
        print(f"Error during OCR processing in perform_ocr_on_image: {e}")
        print(traceback.format_exc())
        return "OCR Processing Error"

# ... (rest of run_detection_ocr, Flask routes, and __main__ remain the same) ...

def run_detection_ocr(image_path, result_filename):
    """Runs YOLO detection and EasyOCR on the image."""
    if model is None or reader is None:
        print("Models not loaded. Cannot process image.")
        return None, ["Model Loading Error"]

    detected_texts = []
    result_image_path = None

    try:
        # Read the image
        img = cv2.imread(image_path)
        if img is None:
            print(f"Error: Could not read image at {image_path}")
            return None, ["Image Reading Error"]

        img_copy = img.copy() # Make a copy for drawing

        # Run YOLO detection with original confidence threshold
        print(f"Running YOLO prediction on {image_path}...")
        results = model.predict(img, conf=0.25) # Restore original confidence for detection
        print("YOLO prediction complete.")

        # Process detections
        for result in results:
            boxes = result.boxes
            if boxes is None or len(boxes) == 0:
                print("No boxes detected in this result.")
                continue

            for box in boxes:
                try:
                    # Get class index and name using the loaded names
                    class_id = int(box.cls[0])
                    class_name = model_class_names.get(class_id, 'Unknown')
                    confidence = float(box.conf[0])
                    x1_temp, y1_temp, x2_temp, y2_temp = map(int, box.xyxy[0])
                    print(f"Detected '{class_name}' (Conf: {confidence:.2f}) at [{x1_temp}, {y1_temp}, {x2_temp}, {y2_temp}]")

                    # Check with the CORRECTED class name
                    if class_name.lower() == PLATE_CLASS_NAME.lower(): # Comparison should now work
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        padding = 10
                        crop_y1 = max(0, y1 - padding)
                        crop_y2 = min(img.shape[0], y2 + padding)
                        crop_x1 = max(0, x1 - padding)
                        crop_x2 = min(img.shape[1], x2 + padding)

                        if crop_y2 > crop_y1 and crop_x2 > crop_x1:
                            plate_crop = img[crop_y1:crop_y2, crop_x1:crop_x2]

                            # Added print before calling OCR
                            print(f"--- Cropped {PLATE_CLASS_NAME} at [{crop_x1}, {crop_y1}, {crop_x2}, {crop_y2}], attempting OCR ---")
                            plate_text = perform_ocr_on_image(plate_crop) # Function updated with preprocessing/postprocessing
                            print(f"-> OCR Result returned for {PLATE_CLASS_NAME}: {plate_text}")
                            detected_texts.append(plate_text)

                            # Drawing logic
                            label = f"{class_name}: {plate_text}" # Label uses result from OCR call
                            cv2.rectangle(img_copy, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            text_pos = (x1, y1 - 10) if y1 > 20 else (x1, y2 + 20)
                            cv2.putText(img_copy, label, text_pos,
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                        else:
                            print(f"Skipping invalid crop dimensions for {PLATE_CLASS_NAME}: y1={crop_y1}, y2={crop_y2}, x1={crop_x1}, x2={crop_x2}")
                    else:
                       # Draw blue boxes for other detected objects (optional)
                       cv2.rectangle(img_copy, (x1_temp, y1_temp), (x2_temp, y2_temp), (255, 0, 0), 1) # Blue box, thinner
                       cv2.putText(img_copy, f"{class_name} {confidence:.2f}", (x1_temp, y1_temp - 5),
                                   cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)


                except Exception as e_inner:
                    print(f"Error processing a specific detection: {e_inner}")
                    print(traceback.format_exc())
                    detected_texts.append("Detection Proc. Error")

        # Save the result image
        result_image_path = os.path.join(app.config['RESULT_FOLDER'], result_filename)
        cv2.imwrite(result_image_path, img_copy)
        print(f"Result image saved to: {result_image_path}")

        return result_image_path, detected_texts

    except Exception as e:
        print(f"Error during detection/OCR: {e}")
        print(traceback.format_exc())
        return None, ["Overall Processing Error"]

# --- Flask Routes ---
@app.route('/', methods=['GET', 'POST'])
def upload_file():
    error_message = None
    if request.method == 'POST':
        if 'file' not in request.files:
            flash('No file part')
            return redirect(request.url)
        file = request.files['file']
        if file.filename == '':
            flash('No selected file')
            return redirect(request.url)

        if file and allowed_file(file.filename):
            try:
                original_filename = secure_filename(file.filename)
                timestamp = time.strftime("%Y%m%d-%H%M%S")
                unique_filename = f"{timestamp}_{original_filename}"
                upload_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
                file.save(upload_path)
                print(f"File saved to: {upload_path}")

                # --- Call Processing Logic ---
                if model is None or reader is None:
                    return render_template('result.html',
                                           original_image_url=url_for('static', filename=f'uploads/{unique_filename}'),
                                           result_image_url=None,
                                           detected_texts=None,
                                           error_message="AI models failed to load. Cannot process image.")

                result_filename = f"result_{unique_filename}"
                result_image_path, detected_texts = run_detection_ocr(upload_path, result_filename)

                # Generate URLs for the templates
                original_url = url_for('static', filename=f'uploads/{unique_filename}')
                result_url = url_for('static', filename=f'results/{result_filename}') if result_image_path else None

                # *** Filter out error messages before displaying ***
                display_texts = [text for text in detected_texts if "Error" not in text and text != "No text found" and text != "Crop too small" and text != "Invalid Crop"]
                if not display_texts and detected_texts: # If only errors or 'no text', show generic message
                    display_texts = ["No text recognized"]
                elif not detected_texts and not display_texts: # If detection worked but OCR returned nothing valid
                    display_texts = ["No text recognized"]
                elif not detected_texts: # Fallback if processing completely failed
                     display_texts = ["Processing Error"]


                return render_template('result.html',
                                       original_image_url=original_url,
                                       result_image_url=result_url,
                                       detected_texts=display_texts, # Use filtered list
                                       error_message=None if result_image_path else "An error occurred during image processing.")

            except Exception as e:
                error_message = f"An error occurred during file upload or processing: {e}"
                print(f"Error in upload_file route: {e}")
                print(traceback.format_exc())
                # Render index again, show the error
                return render_template('index.html', error_message=error_message)

        else:
            error_message = "Invalid file type. Please upload a PNG, JPG, or JPEG image."
            return render_template('index.html', error_message=error_message)

    # For GET requests
    return render_template('index.html', error_message=None)


# --- Run the App ---
if __name__ == "__main__":
    app.run(debug=True)