sushreen/AlphabetDoodleRecognition

1import torch2import torch.nn as nn3import torchvision.transforms as transforms4from flask import Flask, request, jsonify, render_template5from PIL import Image6import io7from flask_cors import CORS8import torch.nn.functional as F9from PIL import ImageOps10 11 12class ResBlock(nn.Module):13    def __init__(self, input_features, output_features):14        super(ResBlock, self).__init__()15        self.stride = 1 if input_features == output_features else 216        17        #main convolutional path18        self.features = nn.Sequential(19            nn.Conv2d(input_features, output_features, kernel_size=3, stride=self.stride, padding=1, bias=False),20            nn.BatchNorm2d(output_features),21            nn.ReLU(inplace=True),22            nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1, bias=False),23            nn.BatchNorm2d(output_features)24        )25 26        #shortcut connection27        self.shortcut = nn.Identity()28        if input_features != output_features:29            self.shortcut = nn.Sequential(30                nn.Conv2d(input_features, output_features, kernel_size=1, stride=self.stride, bias=False),31                nn.BatchNorm2d(output_features)32            )33 34    def forward(self, x):35        residual = self.shortcut(x)36        x = self.features(x)37        x += residual38        x = F.relu(x, inplace=True)39        return x40 41class Resnet18(nn.Module):42    def __init__(self, num_of_classes=26):  #classes is 26 for EMNIST letters43        super(Resnet18, self).__init__()44        45        self.features = nn.Sequential(46            nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False), 47            nn.BatchNorm2d(64),48            nn.ReLU(inplace=True),49            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),50 51            ResBlock(64, 64),52            ResBlock(64, 64),53 54            ResBlock(64, 128),55            ResBlock(128, 128),56 57            ResBlock(128, 256),58            ResBlock(256, 256),59 60            ResBlock(256, 512),61            ResBlock(512, 512),62 63            nn.AdaptiveAvgPool2d((1, 1))64        )65 66        self.classifier = nn.Sequential(67            nn.Linear(512, num_of_classes)68        )69 70    def forward(self, x):71        x = self.features(x)72        x = torch.flatten(x, 1)73        x = self.classifier(x)74        return x75 76# Load model77device = "cpu"78model = Resnet18().to(device)79model.load_state_dict(torch.load("resnet_emnist_letters_cpu.pth"))80model.eval()81 82# Define image preprocessing83transform = transforms.Compose([transforms.Grayscale(), transforms.Resize((224, 224)), transforms.ToTensor()])84 85# Initialize Flask app86app = Flask(__name__)87CORS(app)88 89@app.route("/")90def home():91    return render_template("index.html")92 93# Route to handle image predictions94@app.route("/predict", methods=["POST"])95def predict():96    file = request.files["image"].read()97    image = Image.open(io.BytesIO(file)).convert("L")98    image = image.rotate(-90, expand=True)  # EMNIST orientation99    image = image.transpose(Image.FLIP_LEFT_RIGHT)100    image = transform(image).unsqueeze(0).to(device)101 102 103      104    with torch.no_grad():105        outputs = model(image)106        _, predicted = torch.max(outputs, 1)107 108    class_labels = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']109    prediction = class_labels[predicted.item()]110 111    return jsonify({"prediction": prediction})112 113if __name__ == "__main__":114    app.run(host="0.0.0.0", port=7860)115