from flask import Flask, request, jsonify
from ultralytics import YOLO
from paddleocr import PaddleOCR
import cv2
import numpy as np
import base64
import os

app = Flask(__name__)

# Load Models Once at startup
print("Loading YOLO and PaddleOCR models...")
bus_model = YOLO('best.pt') 
ocr = PaddleOCR(use_textline_orientation=True, lang='en')
print("Models loaded successfully.")

@app.route('/detect', methods=['POST'])
def detect():
    try:
        data = request.json
        if not data or 'image' not in data:
            return jsonify({"error": "No image provided"}), 400
            
        # Decode base64 image
        img_data = base64.b64decode(data['image'])
        nparr = np.frombuffer(img_data, np.uint8)
        frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        if frame is None:
            return jsonify({"error": "Invalid image format"}), 400

        # 1. Bus / Plate Detection via YOLOv8
        results = bus_model(frame, imgsz=640)[0]
        
        bus_detected = False
        license_plate_text = ""
        highest_conf = 0.0

        for box in results.boxes:
            cls_id = int(box.cls[0]) # 0 for bus, 1 for plate (check your labels!)
            conf = float(box.conf[0])
            
            if cls_id == 0: 
                bus_detected = True
            elif cls_id == 1: 
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                
                # Add 5% padding to avoid clipping characters
                h, w = frame.shape[:2]
                pad_x = int((x2 - x1) * 0.05)
                pad_y = int((y2 - y1) * 0.05)
                px1, py1 = max(0, x1 - pad_x), max(0, y1 - pad_y)
                px2, py2 = min(w, x2 + pad_x), min(h, y2 + pad_y)
                
                plate_crop = frame[py1:py2, px1:px2]
                
                # 2. Run PaddleOCR on cropped plate
                ocr_result = ocr.ocr(plate_crop, cls=True)
                if ocr_result and ocr_result[0]:
                    for line in ocr_result[0]:
                        text, confidence = line[1]
                        if confidence > highest_conf:
                            highest_conf = confidence
                            license_plate_text = text

        return jsonify({
            "bus_detected": bus_detected,
            "license_plate": license_plate_text,
            "confidence": highest_conf
        })

    except Exception as e:
        print(f"Error during inference: {e}")
        return jsonify({"error": str(e)}), 500

@app.route('/', methods=['GET'])
def health():
    return jsonify({"status": "ML Inference Service is healthy"})

if __name__ == '__main__':
    # HuggingFace spaces commonly use port 7860
    app.run(host='0.0.0.0', port=7860)