from flask import Flask, render_template, request, jsonify
import cv2
import numpy as np
from ultralytics import YOLO
import base64

app = Flask(__name__)
model = YOLO("yolov8n.pt")

def readb64(base64_string):
    if ',' not in base64_string:
        raise ValueError("Invalid base64 image string.")
    header, encoded = base64_string.split(',', 1)
    if not encoded:
        raise ValueError("Empty base64 image data.")
    nparr = np.frombuffer(base64.b64decode(encoded), np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("Image decoding failed.")
    return img


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

@app.route('/detect', methods=['POST'])
def detect():
    data = request.get_json()
    frame_data = data['image']

    frame = readb64(frame_data)
    results = model.predict(frame, conf=0.5)
    annotated = results[0].plot()

    _, buffer = cv2.imencode('.jpg', annotated)
    annotated_b64 = base64.b64encode(buffer).decode('utf-8')
    return jsonify({'image': f'data:image/jpeg;base64,{annotated_b64}'})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)