flask-rtod / app.py
dwkurnie
Add application file
ff462b3
from flask import Flask, render_template, Response
import cv2
import torch
# import pathlib
# temp = pathlib.PosixPath
# pathlib.PosixPath = pathlib.WindowsPath
app = Flask(__name__)
# Load custom YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5/weights/model5.pt')
def gen_frames():
cap = cv2.VideoCapture(0)
while True:
success, frame = cap.read()
if not success:
break
else:
# Perform object detection
results = model(frame)
frame = results.render()[0]
# Encode frame as JPEG
ret, buffer = cv2.imencode('.jpg', frame)
frame = buffer.tobytes()
yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
@app.route('/')
def index():
return render_template('index.html')
@app.route('/video_feed')
def video_feed():
return Response(gen_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')
if __name__ == '__main__':
app.run(debug=True)