Spaces:
Sleeping
Sleeping
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) | |
def home(): | |
return render_template("index.html") | |
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) | |
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) |