67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
|
from ultralytics import YOLO
|
||
|
|
from flask import request, Flask, jsonify
|
||
|
|
from waitress import serve
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
#my changes
|
||
|
|
import os
|
||
|
|
|
||
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||
|
|
# Change the working directory to the script's directory
|
||
|
|
os.chdir(script_dir)
|
||
|
|
|
||
|
|
#app start
|
||
|
|
|
||
|
|
app = Flask(__name__)
|
||
|
|
|
||
|
|
@app.route("/")
|
||
|
|
def root():
|
||
|
|
"""
|
||
|
|
Site main page handler function.
|
||
|
|
:return: Content of index.html file
|
||
|
|
"""
|
||
|
|
with open("index.html") as file:
|
||
|
|
return file.read()
|
||
|
|
|
||
|
|
|
||
|
|
@app.route("/detect", methods=["POST"])
|
||
|
|
def detect():
|
||
|
|
"""
|
||
|
|
Handler of /detect POST endpoint
|
||
|
|
Receives uploaded file with a name "image_file", passes it
|
||
|
|
through YOLOv8 object detection network and returns and array
|
||
|
|
of bounding boxes.
|
||
|
|
:return: a JSON array of objects bounding boxes in format [[x1,y1,x2,y2,object_type,probability],..]
|
||
|
|
"""
|
||
|
|
buf = request.files["image_file"]
|
||
|
|
boxes = detect_objects_on_image(buf.stream)
|
||
|
|
return jsonify(boxes)
|
||
|
|
|
||
|
|
|
||
|
|
def detect_objects_on_image(buf):
|
||
|
|
"""
|
||
|
|
Function receives an image,
|
||
|
|
passes it through YOLOv8 neural network
|
||
|
|
and returns an array of detected objects
|
||
|
|
and their bounding boxes
|
||
|
|
:param buf: Input image file stream
|
||
|
|
:return: Array of bounding boxes in format [[x1,y1,x2,y2,object_type,probability],..]
|
||
|
|
"""
|
||
|
|
model = YOLO("best.pt")
|
||
|
|
results = model.predict(Image.open(buf))
|
||
|
|
result = results[0]
|
||
|
|
output = []
|
||
|
|
for box in result.boxes:
|
||
|
|
x1, y1, x2, y2 = [
|
||
|
|
round(x) for x in box.xyxy[0].tolist()
|
||
|
|
]
|
||
|
|
class_id = box.cls[0].item()
|
||
|
|
prob = round(box.conf[0].item(), 2)
|
||
|
|
output.append([
|
||
|
|
x1, y1, x2, y2, result.names[class_id], prob
|
||
|
|
])
|
||
|
|
return output
|
||
|
|
|
||
|
|
|
||
|
|
serve(app, host='0.0.0.0', port=8080)
|