from flask import Flask, render_template, jsonify, request, send_file
import psutil
import json
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import io
import base64

app = Flask(__name__)
modelName = "yolov9c-seg.pt"
model = YOLO(modelName)


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


@app.route("/sysInfo")
def sysInfo():
    ram = psutil.virtual_memory()
    ram_usage = ram.percent
    cpu_usage = psutil.cpu_percent(interval=1)
    data = {"ram": ram_usage, "cpu": cpu_usage}
    return jsonify(data)


@app.route("/processor", methods=["POST"])
def processor():
    global modelName, model
    image = request.form.get("image")
    modelNameForm = request.form.get("model", modelName)
    if modelNameForm != modelName:
        modelName = modelNameForm
        model = YOLO(modelName)
    image = image.split(",")[1]
    image = base64.b64decode(image)
    image = Image.open(io.BytesIO(image))
    image = np.array(image)
    results = model(
        image,
        show=False,
        save=False,
        show_boxes=True,
        show_labels=True,
        imgsz=640,
        iou=0.1,
        max_det=20,
    )
    image = results[0].plot()
    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    image = cv2.imencode(".jpg", image)[1].tobytes()
    image = base64.b64encode(image).decode("utf-8")
    return jsonify({"image": image})


app.run(debug=True, host="0.0.0.0", port=7860)
# app.run(debug=True)