Spaces:
Sleeping
Sleeping
File size: 3,168 Bytes
bbea32b |
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 |
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from flask import Flask, request, jsonify, render_template
from PIL import Image
import io
from flask_cors import CORS
import torch.nn.functional as F
class ResBlock(nn.Module):
def __init__(self, input_features, output_features):
super(ResBlock, self).__init__()
self.stride = 1 if input_features == output_features else 2
self.features = nn.Sequential(
nn.Conv2d(input_features, output_features,
kernel_size=3, stride=self.stride, padding=1, bias=False),
nn.BatchNorm2d(output_features),
nn.ReLU(inplace=True),
nn.Conv2d(output_features, output_features,
kernel_size=3, stride=1, padding=1, bias=False),
nn.BatchNorm2d(output_features)
)
self.shortcut = nn.Sequential(nn.Identity())
if input_features != output_features:
self.shortcut = nn.Sequential(
nn.Conv2d(input_features, output_features, kernel_size=1, stride=self.stride, bias=False))
def forward(self, x):
residual = self.shortcut(x)
x = self.features(x)
x += residual
x = F.relu(x, inplace=True)
return x
class Resnet18(nn.Module):
def __init__(self, num_of_classes=10):
super(Resnet18, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
ResBlock(64, 64),
ResBlock(64, 64),
ResBlock(64, 128),
ResBlock(128, 128),
ResBlock(128, 256),
ResBlock(256, 256),
ResBlock(256, 512),
ResBlock(512, 512),
nn.AdaptiveAvgPool2d((1, 1))
)
self.classifier = nn.Sequential(
nn.Linear(512, num_of_classes)
)
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
# Load model
device = "cpu"
model = Resnet18().to(device)
model.load_state_dict(torch.load("resnet_mnist_cpu.pth"))
model.eval()
# Define image preprocessing
transform = transforms.Compose([transforms.Grayscale(), transforms.Resize((224, 224)), transforms.ToTensor()])
# Initialize Flask app
app = Flask(__name__)
CORS(app)
@app.route("/")
def home():
return render_template("index.html")
# Route to handle image predictions
@app.route("/predict", methods=["POST"])
def predict():
file = request.files["image"].read()
image = Image.open(io.BytesIO(file))
image = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(image)
_, predicted = torch.max(outputs, 1)
class_labels = ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9" ] # Modify based on your dataset
prediction = class_labels[predicted.item()]
return jsonify({"prediction": prediction})
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860)
|