from flask import Flask, request, jsonify

import torch

from torchvision import models, transforms

from PIL import Image

import io

app = Flask(__name__)

# Load the pre-trained PyTorch model (replace with your trained model)

model = models.resnet18(pretrained=True)  # Example model

model.eval()  # Set to evaluation mode for inference

# Define image transformation for input

transform = transforms.Compose([

    transforms.Resize(224),

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

])

@app.route(‘/predict’, methods=[‘POST’])

def predict():

    try:

        # Get image file from the request

        img_file = request.files[‘image’]

        img = Image.open(io.BytesIO(img_file.read()))

        # Preprocess the image

        img_tensor = transform(img).unsqueeze(0)

        # Make prediction

        with torch.no_grad():

            outputs = model(img_tensor)

            _, predicted_class = torch.max(outputs, 1)

        # Return the prediction in JSON format

        return jsonify({‘predicted_class’: int(predicted_class.item())})

    except Exception as e:

        return jsonify({‘error’: str(e)}), 400

if __name__ == “__main__”:

    app.run(debug=True)

In this Flask app:

• /predict endpoint: Accepts a POST request where an image is sent via the request body, and the server responds with the predicted class of the image.
• Image Processing: The image is processed to fit the input format expected by the PyTorch model.
• Prediction: The model performs inference on the processed image, and the class with the highest probability is returned.