Upload 3 files
Browse files- appf.py +80 -0
- facemask_detection_model_f1.pth +3 -0
- requirements.txt +4 -0
appf.py
ADDED
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from flask import Flask, render_template, request, jsonify
|
2 |
+
from PIL import Image
|
3 |
+
import torch
|
4 |
+
import io
|
5 |
+
import base64
|
6 |
+
from torchvision import transforms
|
7 |
+
from face_mask_detection import FaceMaskDetectionModel
|
8 |
+
import numpy as np
|
9 |
+
|
10 |
+
app = Flask(__name__)
|
11 |
+
|
12 |
+
# Load the model
|
13 |
+
model = FaceMaskDetectionModel()
|
14 |
+
# Load the state dictionary
|
15 |
+
model_state_dict = torch.load("models\\facemask_model_statedict1_f.pth", map_location=torch.device('cpu'))
|
16 |
+
# Load the state dictionary into the model
|
17 |
+
model.load_state_dict(model_state_dict)
|
18 |
+
# Set the model to evaluation mode
|
19 |
+
model.eval()
|
20 |
+
|
21 |
+
|
22 |
+
# Define the pre-processing transform
|
23 |
+
transform = transforms.Compose([
|
24 |
+
transforms.Resize((224, 224)),
|
25 |
+
transforms.ToTensor()
|
26 |
+
])
|
27 |
+
|
28 |
+
# Define class labels
|
29 |
+
class_labels = ['without mask', 'with mask']
|
30 |
+
|
31 |
+
|
32 |
+
@app.route('/')
|
33 |
+
def index():
|
34 |
+
return render_template('index.html')
|
35 |
+
|
36 |
+
|
37 |
+
@app.route('/predict', methods=['POST'])
|
38 |
+
def predict():
|
39 |
+
try:
|
40 |
+
# Get the image from the request
|
41 |
+
image = request.files['image']
|
42 |
+
|
43 |
+
# Pre-process the image
|
44 |
+
image_tensor = transform(Image.open(io.BytesIO(image.read())).convert('RGB')).unsqueeze(0)
|
45 |
+
|
46 |
+
# Set the model to evaluation mode
|
47 |
+
model.eval()
|
48 |
+
|
49 |
+
# Make a prediction
|
50 |
+
with torch.no_grad():
|
51 |
+
output = model(image_tensor)
|
52 |
+
print("Output: ", output)
|
53 |
+
|
54 |
+
# Convert the output to probabilities using softmax
|
55 |
+
probabilities = torch.nn.functional.softmax(output[0], dim=0)
|
56 |
+
print("Probabilities: ", probabilities)
|
57 |
+
|
58 |
+
# Get the predicted class
|
59 |
+
predicted_class = torch.argmax(probabilities).item()
|
60 |
+
print("Predicted: ", predicted_class)
|
61 |
+
|
62 |
+
# Get the probability for the predicted class
|
63 |
+
predicted_probability = probabilities[predicted_class].item()
|
64 |
+
|
65 |
+
# Define class labels
|
66 |
+
class_labels = ['without mask', 'with mask']
|
67 |
+
|
68 |
+
print(f"Predicted Class: {class_labels[predicted_class]}")
|
69 |
+
print(f"Probability: {predicted_probability:.4f}")
|
70 |
+
|
71 |
+
# Return the prediction along with the uploaded image
|
72 |
+
image_base64 = base64.b64encode(image.read()).decode('utf-8')
|
73 |
+
return jsonify({'prediction': predicted_class, 'image': image_base64})
|
74 |
+
except Exception as e:
|
75 |
+
return jsonify({'error': str(e)}), 500
|
76 |
+
|
77 |
+
|
78 |
+
|
79 |
+
if __name__ == '__main__':
|
80 |
+
app.run(debug=True)
|
facemask_detection_model_f1.pth
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:7258119cb55a3a49bf18f5d4034690009205313cf309810c8a51cf109d5af1ce
|
3 |
+
size 51776804
|
requirements.txt
ADDED
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
1 |
+
torch
|
2 |
+
torchvision
|
3 |
+
flask
|
4 |
+
numpy
|