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)