app.py

# --- 1. Import all the necessary tools ---
import gradio as gr
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import numpy as np
import cv2
import roboflow
from collections import Counter, defaultdict
import re

# --- 2. Load BOTH of your AI models ---
print("Downloading and loading models...")

# --- Model 1: The Character Detector (from Hugging Face) ---
character_model_path = hf_hub_download(
    repo_id="MKgoud/License-Plate-Character-Detector", 
    filename="Charcter-LP.pt"
)
character_model = YOLO(character_model_path)
print("✅ Character Detector loaded.")

# --- Model 2: The Plate Detector (from Roboflow) ---
ROBOFLOW_API_KEY = "YfKCsreNkoXYFD1CfMBY"
DETECTOR_WORKSPACE_ID = "mylprproject"
DETECTOR_PROJECT_ID = "license-plate-yuw1z-kirke"
DETECTOR_VERSION_NUMBER = 1

rf = roboflow.Roboflow(api_key=ROBOFLOW_API_KEY)
project_detector = rf.workspace(DETECTOR_WORKSPACE_ID).project(DETECTOR_PROJECT_ID)
plate_model = project_detector.version(DETECTOR_VERSION_NUMBER).model
print("✅ Plate Detector loaded.")

# --- 3. Enhanced preprocessing functions ---
def enhance_plate_image(plate_crop):
    """
    Apply multiple enhancement techniques to improve character visibility
    """
    enhanced_crops = []
    
    # Original image
    enhanced_crops.append(plate_crop)
    
    # Convert to grayscale and back to RGB for consistent processing
    gray = cv2.cvtColor(plate_crop, cv2.COLOR_RGB2GRAY)
    
    # Enhancement 1: Adaptive histogram equalization
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    enhanced_gray = clahe.apply(gray)
    enhanced_crops.append(cv2.cvtColor(enhanced_gray, cv2.COLOR_GRAY2RGB))
    
    # Enhancement 2: Gaussian blur + unsharp mask
    blurred = cv2.GaussianBlur(gray, (3, 3), 0)
    unsharp = cv2.addWeighted(gray, 1.8, blurred, -0.8, 0)
    unsharp = np.clip(unsharp, 0, 255).astype(np.uint8)
    enhanced_crops.append(cv2.cvtColor(unsharp, cv2.COLOR_GRAY2RGB))
    
    # Enhancement 3: Contrast stretching
    min_val, max_val = np.percentile(gray, [2, 98])
    stretched = np.clip((gray - min_val) * 255 / (max_val - min_val), 0, 255).astype(np.uint8)
    enhanced_crops.append(cv2.cvtColor(stretched, cv2.COLOR_GRAY2RGB))
    
    # Enhancement 4: Gamma correction
    gamma_corrected = np.power(gray / 255.0, 0.7) * 255
    gamma_corrected = gamma_corrected.astype(np.uint8)
    enhanced_crops.append(cv2.cvtColor(gamma_corrected, cv2.COLOR_GRAY2RGB))
    
    return enhanced_crops

def smart_character_correction(text, pattern_analysis=True):
    """
    Intelligent character correction based on license plate patterns
    """
    if not text or len(text) < 3:
        return text
    
    # Remove any spaces first
    text = text.replace(" ", "").upper()
    
    # Philippine license plate patterns analysis
    def analyze_likely_pattern(s):
        """Determine if sequence should be letters or numbers based on position and context"""
        if len(s) < 6:
            return s
        
        # Common Philippine patterns: ABC123, ABC1234, 123ABC
        # Most common is 3 letters + 3-4 numbers
        
        corrected = list(s)
        
        # Pattern 1: First 3 characters are typically letters
        for i in range(min(3, len(corrected))):
            char = corrected[i]
            if char.isdigit():
                # Convert numbers that look like letters
                digit_to_letter = {'0': 'O', '1': 'I', '2': 'Z', '5': 'S', '6': 'G', '8': 'B'}
                if char in digit_to_letter:
                    corrected[i] = digit_to_letter[char]
        
        # Pattern 2: Characters after position 3 are typically numbers
        for i in range(3, len(corrected)):
            char = corrected[i]
            if char.isalpha():
                # Convert letters that look like numbers
                letter_to_digit = {'O': '0', 'I': '1', 'L': '1', 'S': '5', 'G': '6', 'B': '8', 'Z': '2', 'T': '7'}
                if char in letter_to_digit:
                    corrected[i] = letter_to_digit[char]
        
        return ''.join(corrected)
    
    # Apply pattern-based corrections if enabled
    if pattern_analysis and len(text) >= 6:
        text = analyze_likely_pattern(text)
    
    # Additional common OCR error corrections
    ocr_corrections = {
        # Numbers that might be misread as letters
        'Q': '0',  # Q often confused with O/0
        'D': '0',  # D sometimes looks like 0
        # Letters that might be misread as numbers  
        'A': 'A',  # Keep A as is (could be confused with 4 but A is common in plates)
    }
    
    # Apply only high-confidence corrections
    for old_char, new_char in ocr_corrections.items():
        text = text.replace(old_char, new_char)
    
    return text

def advanced_detection_filtering(boxes, character_results, plate_crop_shape, min_confidence=0.25):
    """
    Advanced filtering with clustering and statistical analysis
    """
    if len(boxes) == 0:
        return []
    
    detections = []
    
    # Extract all detection info
    for box in boxes:
        confidence = float(box.conf[0])
        if confidence < min_confidence:
            continue
            
        class_id = int(box.cls[0])
        character = character_results[0].names[class_id]
        x1, y1, x2, y2 = box.xyxy[0]
        
        detections.append({
            'char': character,
            'conf': confidence,
            'x1': float(x1), 'y1': float(y1), 'x2': float(x2), 'y2': float(y2),
            'width': float(x2 - x1),
            'height': float(y2 - y1),
            'center_x': float((x1 + x2) / 2),
            'center_y': float((y1 + y2) / 2),
            'area': float((x2 - x1) * (y2 - y1))
        })
    
    if len(detections) == 0:
        return []
    
    plate_height, plate_width = plate_crop_shape[:2]
    
    # Step 1: Focus on main character area (upper 75% of plate)
    main_area_threshold = plate_height * 0.75
    main_detections = [d for d in detections if d['center_y'] <= main_area_threshold]
    
    if len(main_detections) < 3:  # If too few in main area, expand slightly
        main_area_threshold = plate_height * 0.85
        main_detections = [d for d in detections if d['center_y'] <= main_area_threshold]
    
    if len(main_detections) == 0:
        main_detections = detections
    
    # Step 2: Statistical filtering based on size and position
    heights = [d['height'] for d in main_detections]
    widths = [d['width'] for d in main_detections]
    y_positions = [d['center_y'] for d in main_detections]
    areas = [d['area'] for d in main_detections]
    
    # Calculate robust statistics (using percentiles to avoid outlier influence)
    median_height = np.median(heights)
    median_width = np.median(widths)
    median_y = np.median(y_positions)
    q75_area = np.percentile(areas, 75)
    
    # Step 3: Multi-criteria filtering
    filtered_detections = []
    
    for detection in main_detections:
        # Size consistency check
        height_ratio = detection['height'] / median_height
        width_ratio = detection['width'] / median_width
        
        # Vertical alignment check  
        y_deviation = abs(detection['center_y'] - median_y)
        max_y_deviation = median_height * 0.5
        
        # Minimum size threshold (avoid tiny noise detections)
        min_size_threshold = plate_height * 0.12
        
        # Area-based filtering (avoid unusually small detections)
        area_threshold = q75_area * 0.3
        
        # Apply all criteria
        if (0.4 <= height_ratio <= 2.5 and
            0.3 <= width_ratio <= 3.0 and
            y_deviation <= max_y_deviation and
            detection['height'] >= min_size_threshold and
            detection['area'] >= area_threshold):
            
            filtered_detections.append(detection)
    
    # Step 4: Remove duplicate detections (same character in nearby positions)
    final_detections = []
    used_positions = []
    
    # Sort by confidence first
    filtered_detections.sort(key=lambda x: x['conf'], reverse=True)
    
    for detection in filtered_detections:
        # Check if this position is too close to already used positions
        too_close = False
        for used_x in used_positions:
            if abs(detection['center_x'] - used_x) < median_width * 0.8:
                too_close = True
                break
        
        if not too_close:
            final_detections.append(detection)
            used_positions.append(detection['center_x'])
    
    return final_detections

def ensemble_character_voting(all_detections, plate_width, confidence_threshold=0.35):
    """
    Advanced ensemble voting with spatial clustering and confidence weighting
    """
    if not all_detections:
        return []
    
    # Step 1: Spatial clustering - group detections by x-position
    position_groups = defaultdict(list)
    cluster_tolerance = plate_width * 0.15  # 15% of plate width
    
    for detection in all_detections:
        x_pos = detection['center_x']
        
        # Find existing cluster or create new one
        assigned = False
        for cluster_center in list(position_groups.keys()):
            if abs(x_pos - cluster_center) <= cluster_tolerance:
                position_groups[cluster_center].append(detection)
                assigned = True
                break
        
        if not assigned:
            position_groups[x_pos].append(detection)
    
    # Step 2: For each spatial cluster, determine best character
    final_characters = []
    
    for cluster_center, cluster_detections in position_groups.items():
        # Group by character within cluster
        char_groups = defaultdict(list)
        for det in cluster_detections:
            char_groups[det['char']].append(det)
        
        # Calculate weighted score for each character
        best_char = None
        best_score = 0
        best_detection = None
        
        for char, char_detections in char_groups.items():
            # Calculate score: weighted average of confidence + occurrence bonus
            confidences = [d['conf'] for d in char_detections]
            avg_confidence = np.mean(confidences)
            max_confidence = max(confidences)
            occurrence_bonus = min(len(char_detections) * 0.1, 0.3)  # Up to 30% bonus
            
            # Final score combines average confidence, max confidence, and occurrence
            score = (avg_confidence * 0.5 + max_confidence * 0.4 + occurrence_bonus * 0.1)
            
            if score > best_score and avg_confidence > confidence_threshold:
                best_score = score
                best_char = char
                # Use the detection with highest confidence as representative
                best_detection = max(char_detections, key=lambda x: x['conf'])
        
        if best_char and best_detection:
            best_detection['final_char'] = best_char
            best_detection['final_score'] = best_score
            best_detection['cluster_size'] = len(cluster_detections)
            final_characters.append(best_detection)
    
    # Step 3: Sort by x-position for final ordering
    final_characters.sort(key=lambda x: x['center_x'])
    
    return final_characters

# --- 4. Enhanced main prediction function ---
def detect_license_plate(input_image):
    """
    Enhanced version with improved filtering and ensemble voting
    """
    print("New image received. Starting enhanced 2-stage pipeline...")
    output_image = input_image.copy()
    
    # --- STAGE 1: Find the license plate ---
    plate_predictions = plate_model.predict(input_image, confidence=40, overlap=30).json()['predictions']

    if not plate_predictions:
        return output_image, "No license plate found."

    # Get the highest confidence plate detection
    plate_box = max(plate_predictions, key=lambda x: x['confidence'])
    x1, y1, x2, y2 = [int(p) for p in [plate_box['x'] - plate_box['width'] / 2, 
                                       plate_box['y'] - plate_box['height'] / 2,
                                       plate_box['x'] + plate_box['width'] / 2, 
                                       plate_box['y'] + plate_box['height'] / 2]]
    
    # Optimized padding - minimal vertical to avoid extra text
    h_padding = 10
    v_padding = 2
    y1 = max(0, y1 - v_padding)
    x1 = max(0, x1 - h_padding)
    y2 = min(input_image.shape[0], y2 + v_padding)
    x2 = min(input_image.shape[1], x2 + h_padding)
    
    plate_crop = input_image[y1:y2, x1:x2]
    plate_height, plate_width = plate_crop.shape[:2]
    
    # Focus on main number area (top 75% of plate)
    main_number_crop = plate_crop[:int(plate_height * 0.75), :]
    
    # --- STAGE 2: Multi-enhancement character detection ---
    enhanced_crops = enhance_plate_image(main_number_crop)
    
    all_detections = []
    
    # Process each enhanced version with different confidence thresholds
    confidence_levels = [0.25, 0.3, 0.35, 0.25]  # Different thresholds for each enhancement
    
    for i, (enhanced_crop, conf_threshold) in enumerate(zip(enhanced_crops, confidence_levels)):
        try:
            character_results = character_model(enhanced_crop, conf=conf_threshold, iou=0.3)
            
            if character_results and hasattr(character_results[0], 'boxes') and len(character_results[0].boxes) > 0:
                boxes = character_results[0].boxes.cpu().numpy()
                filtered_detections = advanced_detection_filtering(
                    boxes, character_results, main_number_crop.shape, min_confidence=conf_threshold
                )
                
                print(f"Enhancement {i}: {len(boxes)} raw -> {len(filtered_detections)} filtered")
                
                for detection in filtered_detections:
                    detection['enhancement_method'] = i
                    detection['enhancement_conf'] = conf_threshold
                    all_detections.append(detection)
                    
        except Exception as e:
            print(f"Error processing enhancement {i}: {e}")
            continue
    
    # --- STAGE 3: Advanced ensemble voting ---
    final_detections = ensemble_character_voting(all_detections, plate_width, confidence_threshold=0.3)
    
    print(f"Ensemble voting: {len(all_detections)} total -> {len(final_detections)} final")
    
    # --- STAGE 4: Generate and post-process text ---
    if final_detections:
        # Sort by x position
        final_detections.sort(key=lambda x: x['center_x'])
        raw_text = "".join([d['final_char'] for d in final_detections])
        
        # Apply smart character correction
        corrected_text = smart_character_correction(raw_text, pattern_analysis=True)
        
        # Additional validation - remove obviously wrong characters
        if len(corrected_text) > 8:  # If too long, might have false positives
            # Keep only the most confident detections
            final_detections = sorted(final_detections, key=lambda x: x['final_score'], reverse=True)[:7]
            final_detections.sort(key=lambda x: x['center_x'])
            raw_text = "".join([d['final_char'] for d in final_detections])
            corrected_text = smart_character_correction(raw_text, pattern_analysis=True)
    else:
        raw_text = ""
        corrected_text = ""
    
    # --- STAGE 5: Draw results ---
    # Draw the main plate box
    cv2.rectangle(output_image, (x1, y1), (x2, y2), (0, 0, 255), 2)
    cv2.putText(output_image, f"Plate: {plate_box['confidence']:.1f}%", 
                (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
    
    # Draw detection area boundary
    main_area_y = y1 + int(plate_height * 0.75)
    cv2.line(output_image, (x1, main_area_y), (x2, main_area_y), (255, 255, 0), 2)
    cv2.putText(output_image, "Detection Area", (x1, main_area_y - 5), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 0), 1)
    
    # Draw character detections
    for i, detection in enumerate(final_detections):
        abs_x1 = x1 + int(detection['x1'])
        abs_y1 = y1 + int(detection['y1'])
        abs_x2 = x1 + int(detection['x2'])
        abs_y2 = y1 + int(detection['y2'])
        
        # Color code by confidence
        color = (0, 255, 0) if detection['final_score'] > 0.7 else (0, 255, 255)
        
        cv2.rectangle(output_image, (abs_x1, abs_y1), (abs_x2, abs_y2), color, 2)
        cv2.putText(output_image, f"{detection['final_char']}", 
                   (abs_x1, abs_y1 - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
        cv2.putText(output_image, f"{detection['final_score']:.2f}", 
                   (abs_x1, abs_y1 - 3), cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
    
    # Prepare result text
    if raw_text != corrected_text and corrected_text:
        result_text = f"Detected: {raw_text}\nCorrected: {corrected_text}\nConfidence: {len(final_detections)} chars"
    elif corrected_text:
        result_text = f"Result: {corrected_text}\nConfidence: {len(final_detections)} characters detected"
    else:
        result_text = "No characters detected with sufficient confidence"
    
    print(f"Final result: {result_text}")
    
    return output_image, result_text

# --- 5. Create the Gradio Web Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# Enhanced High-Accuracy License Plate Detector")
    gr.Markdown("""
    **Improved Features:**
    - Advanced statistical filtering with spatial clustering
    - Smart character correction based on license plate patterns
    - Enhanced ensemble voting with confidence weighting
    - Optimized detection area focusing
    - Multi-level confidence thresholds
    """)
    
    with gr.Row():
        image_input = gr.Image(type="numpy", label="Upload License Plate Image")
        image_output = gr.Image(type="numpy", label="Detection Results")
    
    text_output = gr.Textbox(label="Detected License Plate", lines=3)
    predict_button = gr.Button(value="Detect License Plate", variant="primary")
    
    predict_button.click(
        fn=detect_license_plate, 
        inputs=image_input, 
        outputs=[image_output, text_output]
    )

# --- 6. Launch the application ---
demo.launch()