softmax

softmax

from __future__ import annotations

import typing as t
from typing import TYPE_CHECKING

import numpy as np
from PIL.Image import Image as PILImage

import bentoml
from <http://bentoml.io|bentoml.io> import Image
from <http://bentoml.io|bentoml.io> import NumpyNdarray

if TYPE_CHECKING:
    from numpy.typing import NDArray

import torch
import torch.nn.functional as F
import torchvision.transforms as T

mnist_runner = bentoml.pytorch.get("covid_new").to_runner()

svc = bentoml.Service(name="covid_new", runners=[mnist_runner])

def to_numpy(tensor):
    return tensor.detach().cpu().numpy()

# Transformation pipeline to preprocess the input image
img_size = 224
test_transform = T.Compose(
    [
        T.Resize(size=(img_size, img_size)),
        T.ToTensor(),
        # for 3-channel images, normalize using mean and std of each channel
        T.Normalize(
            [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
        ),
    ]
)

@svc.api(input=Image(), output=NumpyNdarray(dtype="float32"))
async def predict_image(f: PILImage) -> NDArray[t.Any]:
    assert isinstance(f, PILImage)
    img = test_transform(f)
    
    # Add batch dimension to input tensor
    img = img.unsqueeze(0)

    output_tensor = await mnist_runner.async_run(img)
    
    # Apply softmax function to get a probability distribution over the classes
    probabilities = F.softmax(output_tensor, dim=1)
    
    # Get the index of the class with the highest probability
    _, predicted_class = probabilities.topk(1, dim=1)
    
    # Convert the tensor to numpy array and return the predicted class as integer
    return to_numpy(predicted_class.squeeze()).astype(np.int64)