FabricScan / app.py
thejagstudio's picture
Update app.py
28bf7eb verified
raw
history blame contribute delete
No virus
1.49 kB
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)