MNISTDigitModel/webserver.py at main · thomas-mckee/MNISTDigitModel · GitHub

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
from flask import Flask, request, jsonify, render_template
from PIL import Image
import torch
from torch import nn
from torchvision import transforms
import matplotlib.pyplot as plt

app = Flask(__name__)

class MNISTModelV1(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
        super().__init__()
        self.layer_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(num_features=hidden_units),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(num_features=hidden_units),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Dropout(p=0.25)
        )

        self.layer_2 = nn.Sequential(
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(num_features=hidden_units),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(num_features=hidden_units),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Dropout(p=0.25)
        )

        self.layer_3 = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=hidden_units * 7 * 7, out_features=hidden_units),
            nn.ReLU(),
            nn.BatchNorm1d(num_features=hidden_units),
            nn.Dropout(p=0.25),
            nn.Linear(in_features=hidden_units, out_features=hidden_units),
            nn.ReLU(),
            nn.BatchNorm1d(num_features=hidden_units),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=hidden_units, out_features=output_shape)  # final classification layer
        )

    def forward(self, x):
        x = self.layer_1(x)
        x = self.layer_2(x)
        x = self.layer_3(x)
        return x


loaded_model = MNISTModelV1(1, 32, 10)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
state_dict = torch.load("model/model_v1.pth", map_location=device)
loaded_model.load_state_dict(state_dict)
loaded_model.to(device)

loaded_model.eval()

transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
])


@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify({'error': 'No image uploaded'}), 400

    image_file = request.files['image']
    image = Image.open(image_file)

    # Optional: convert to grayscale, resize, etc.
    tensor_image = transform(image).unsqueeze(0)  # Add batch dimension


    # Make prediction using model
    with torch.inference_mode():
        output = loaded_model(tensor_image)
        print(torch.softmax(output, dim=1))
        prediction = torch.argmax(output, dim=1).item()

    return jsonify({'prediction': prediction})


if __name__ == '__main__':
    app.run(debug=True)