Spaces:
Sleeping
Sleeping
File size: 1,171 Bytes
d05562c c031299 d05562c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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)
|