Skip to content

Commit

Permalink
Added BlazeFace model loading.
Browse files Browse the repository at this point in the history
Added webcam image crop for better recognition.
Todo: use blazeface to find square with face.
  • Loading branch information
kirmorozov committed Mar 31, 2024
1 parent df06673 commit 356ce01
Show file tree
Hide file tree
Showing 3 changed files with 144 additions and 13 deletions.
21 changes: 18 additions & 3 deletions app.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
from flask import Flask, request
from flask import Flask, request, render_template
from smile_detector import SmileDetector
from io import BytesIO
from urllib.request import urlopen

app = Flask(__name__)

smile_detector = SmileDetector()

@app.route('/')
def hello_world(): # put application's code here
return 'Hello World!'
return render_template('index.html')

@app.route("/has_smile", methods = ['POST'])
def has_smile():
Expand All @@ -18,6 +19,7 @@ def has_smile():
return "ok"
return "ko"


@app.route("/find_faces", methods = ['POST'])
def find_faces():
image_data = request.get_data()
Expand All @@ -27,5 +29,18 @@ def find_faces():
return "ok"
return "ko"

@app.route("/has_smile_json", methods = ['POST'])
def has_smile_json():
json_data = request.json
with urlopen(json_data['image']) as response:
image_data = response.read()

# image_data = requests.get(json_data['image'], stream=True).raw
res = smile_detector.smileCheck(image_data)

if res:
return "ok"
return "ko"

if __name__ == '__main__':
app.run()
48 changes: 38 additions & 10 deletions smile_detector.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import torch
from PIL import Image
from model.lennon import LeNNon
from model.blazeface.blazeface import BlazeFace
from torchvision import transforms
from io import BytesIO

Expand All @@ -9,40 +10,67 @@
__main__.LeNNon=LeNNon



class SmileDetector:
def __init__(self):

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def initFaceDetectionModel(self):
modelDir = "model/blazeface/"
front_net = BlazeFace().to(self.device)
front_net.load_weights(modelDir + "blazeface.pth")
front_net.load_anchors(modelDir + "anchors.npy")
back_net = BlazeFace(back_model=True).to(self.device)
back_net.load_weights(modelDir + "blazefaceback.pth")
back_net.load_anchors(modelDir + "anchorsback.npy")

# Optionally change the thresholds:
front_net.min_score_thresh = 0.75
front_net.min_suppression_threshold = 0.3

self.faceDetectionModel = front_net
def initSmileModel(self):
# Load the model an pass it to the proper device
modelPath = 'model/LeNNon-Smile-Detector.pt'
model = torch.load(modelPath)
model = model.to(device)
model.eval()
smileModel = torch.load(modelPath)
smileModel = smileModel.to(self.device)
smileModel.eval()

# This `transform` object will transform our test images into proper tensors
transform = transforms.Compose([
transforms.Resize((100, 100)), # Resize the image to 100x100
transforms.ToTensor(),
])

self.smileModel = smileModel
self.smileTransform = transform

def __init__(self):

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

self.device = device
self.model = model
self.transform = transform
self.initSmileModel()
self.initFaceDetectionModel()

def smileCheck(self, image_bytes: bytes):
# Open and preprocess he image
image_io = BytesIO(image_bytes)
image = Image.open(image_io)
tensor = self.transform(image)
width, height = image.size # Get dimensions
imageCrop = transforms.CenterCrop((min(width,height),min(width,height)))

tensor = self.smileTransform(imageCrop(image))
tensor = tensor.to(self.device)

# forward pass trough the model
with torch.no_grad():

outputs = self.model(tensor)
outputs = self.smileModel(tensor)

# Get the class prediction
_, predicted = torch.max(outputs.data, 1)

return predicted.item() > 0
return predicted.item() > 0

def findFaces(self, image_bytes: bytes):
# res = self.faceDetectionModel.predict_on_image()
pass
88 changes: 88 additions & 0 deletions templates/index.html
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>Display Webcam Stream</title>

<style>
#container {
margin: 0px auto;
width: 640px;
height: 480px;
border: 10px #333 solid;
}
#webcam {
width: 640px;
height: 480px;
background-color: #666;
}
#controls {
margin: 0px auto;
width: 640px;
text-align: center;
}
.d-none {
display: none;
}
</style>
</head>

<body>
<div id="container">
<video id="webcam" autoplay playsinline width="640" height="480"></video>
<canvas id="canvas" class="d-none"></canvas>
</div>
<div id="controls">
<button id="btnStart">Start</button>
<button id="btnStop">Stop</button>
<button id="btnRun">Capture & Check</button>
</div>
<div id="log">

</div>


<script type="text/javascript" src="https://unpkg.com/webcam-easy/dist/webcam-easy.min.js"></script>
<script>
const webcamElement = document.getElementById('webcam');
const canvasElement = document.getElementById('canvas');
const webcam = new Webcam(webcamElement, 'user', canvasElement, null)

const btnStart = document.getElementById('btnStart');
const btnStop = document.getElementById('btnStop');
const btnRun = document.getElementById('btnRun');
btnStart.addEventListener("click", (event) => {
webcam.start()
.then(result =>{
console.log("webcam started");
})
.catch(err => {
console.log(err);
});
});

btnStop.addEventListener("click", (event) => {
webcam.stop();
});

btnRun.addEventListener("click", (event) => {
var picture = webcam.snap();
// console.log(picture);

var pictureJpeg = canvasElement.toDataURL('image/jpeg');
const rawResponse = fetch('/has_smile_json', {
method: 'POST',
headers: {
'Accept': 'application/json',
'Content-Type': 'application/json'
},
body: JSON.stringify({image: pictureJpeg})
}).then((resp) => {
return resp.json();
}).then((jsonRes) => {
console.log(jsonRes);
});
});
</script>
</body>
</html>

0 comments on commit 356ce01

Please sign in to comment.