Spaces:
Sleeping
Sleeping
File size: 4,583 Bytes
01e1043 4001e3d 01e1043 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
from flask import Flask, request, render_template, send_from_directory
from flask import flash, request, redirect, url_for, Response, make_response
from werkzeug.utils import secure_filename
from super_gradients.training import models
from deep_sort_torch.deep_sort.deep_sort import DeepSort
from super_gradients.training.utils.distributed_training_utils import setup_device
from super_gradients.training.processing import ComposeProcessing
import torch
from model_tools import vid_predict, img_predict
from dotenv import load_dotenv
import os
import urllib.request
myhost = os.uname()[1]
import socket
# Get the fully qualified domain name
fqdn = socket.getfqdn()
print("Fully qualified domain name of this computer is:");
print(fqdn)
load_dotenv()
secret_key = os.getenv("secret_key")
dir = os.getcwd()+ f'/build'
dir_static= dir + '/static'
dir_ckpt = os.getcwd()+ f'/checkpoints'
ckpt_path = dir_ckpt + "/best181-8376/ckpt_latest.pth"
best_model = models.get('yolo_nas_s',
num_classes=1,
checkpoint_path=ckpt_path)
best_model = best_model.to("cuda" if torch.cuda.is_available() else "cpu")
best_model.eval()
#### Initiatize tracker
tracker_model = "./checkpoints/ckpt.t7"
tracker = DeepSort(model_path=tracker_model,max_age=30,nn_budget=100, max_iou_distance=0.7, max_dist=0.2)
tracker = tracker.to("cuda" if torch.cuda.is_available() else "cpu")
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp4', 'avi','webm'}
app = Flask(__name__, template_folder=dir,static_folder=dir_static)
app.config['UPLOAD_FOLDER'] = dir_static
app.config['MAX_CONTENT_LENGTH'] = 20*1024*1024
app.secret_key = secret_key
def allowed_file(filename):
return '.' in filename and \
filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@app.route('/')
def index():
predictions = False
return render_template('index.html', predictions=predictions)
@app.route('/upload', methods=["GET", "POST"])
def upload():
if request.method == 'POST':
print('Form',request.form.get('options'))
try:
filename = request.form.get('options')
if filename:
save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
filetype = 'video'
new_filename = filename[:-4] + ".webm"
save_to = vid_predict(save_path,best_model,tracker, out_path=dir_static, filename=new_filename)
save_to = url_for('static', filename=new_filename)
predictions = True
return render_template('index.html', predictions=predictions, saved_outout=save_to, ft=filetype)
except:
pass
# check if the post request has the file part
if ('file') and ('media') not in request.files:
flash('No file part')
return redirect(request.url)
try:
file = request.files['file']
except:
file = request.files['media']
# If the user does not select a file, the browser submits an
# empty file without a filename.
if file.filename == '':
flash('No selected file')
return redirect(request.url)
if file and allowed_file(file.filename):
filename = secure_filename(file.filename)
save_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
file.save(save_path)
new_filename = filename[:-4] + ".webm"
if filename[-3:] in ['mp4','avi','webm']:
print("VIDEO")
filetype = 'video'
save_to = vid_predict(save_path,best_model,tracker, out_path=dir_static, filename=new_filename)
save_to = url_for('static', filename=new_filename)
else:
print("IMAGE")
filetype = 'image'
save_to = img_predict(save_path,best_model, out_path=dir_static, filename=new_filename)
save_to = url_for('static', filename="pred_0.jpg")
predictions = True
return render_template('index.html', predictions=predictions, saved_outout=save_to, ft=filetype)
@app.route('/static/<folder>/<file>')
def css(folder,file):
''' User will call with with thier id to store the symbol as registered'''
path = folder+'/'+file
return send_from_directory(directory=dir_static,path=path)
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|