-
Notifications
You must be signed in to change notification settings - Fork 0
/
demo.py
36 lines (32 loc) · 890 Bytes
/
demo.py
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
from flask import Flask, send_from_directory
from jinja2 import Template,Environment,FileSystemLoader
import torch as th
from model import CNN
app= Flask(__name__,static_folder="static") # create an app for the website
file_loader = FileSystemLoader('.')
env = Environment(loader=file_loader)
t = env.get_template('index.html')
d = range(20)
m = CNN()
m.load_state_dict(th.load('cnn.pt'))
m.eval()
dataset = th.load("face_dataset.pt")# load face image dataset
X = dataset["X"]
@app.route("/")
def result():
return t.render(face_ids=d)
@app.route("/predict/<int:face_id>")
def predict(face_id):
if face_id%2 == 0:
i = face_id //2
else:
i = face_id //2 + 10
x = X[i].reshape(1,1,64,64)
z=m(x)
if z[0,0]>0:
p = 'Owner'
else:
p = 'Not Owner'
return t.render(face_ids=[face_id], p=p)
if __name__ == "__main__":
app.run()