from flask import Flask, request, jsonify, send_file
from pathlib import Path
from ultralytics import YOLO
import cv2
import numpy as np
import os
import uuid
import logging
from io import BytesIO

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)

# Initialize detector
MODEL_PATH = str(Path(__file__).parent.parent / "runs" / "detect" / "train" / "weights" / "best.pt")
model = YOLO(MODEL_PATH)


@app.route('/')
def index():
    return send_file('test.html')


@app.route('/detect', methods=['POST'])
def detect():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400

    file = request.files['image']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400

    try:
        # Read image directly from memory
        img_bytes = file.read()
        img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)

        # Run prediction
        results = model.predict(img, imgsz=416, conf=0.3)
        
        # Generate annotated image
        annotated = results[0].plot(line_width=2, font_size=0.5)
        
        # Save to results folder
        output_dir = Path("static/results")
        output_dir.mkdir(exist_ok=True)
        filename = f"{uuid.uuid4()}.jpg"
        output_path = output_dir / filename
        cv2.imwrite(str(output_path), annotated)
        
        # Extract detections
        detections = []
        for box in results[0].boxes:
            detections.append({
                'box': box.xyxy[0].tolist(),
                'confidence': float(box.conf[0]),
                'class': int(box.cls[0])
            })
        
        return jsonify({
            'detections': detections,
            'result_image': f"/results/{filename}"
        })
        
    except Exception as e:
        logging.error(f"Detection error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/results/<filename>')
def get_result(filename):
    return send_file(Path("static/results") / filename)

if __name__ == '__main__':
    # Create directories
    Path("static/uploads").mkdir(parents=True, exist_ok=True)
    Path("static/results").mkdir(parents=True, exist_ok=True)
    app.run(host='0.0.0.0', port=5000)