TORCHSCRIPT_CLASSIFIER

Download Flojoy Studio to try this app

Execute a TorchScript classifier against an input image. Params: input_image : Image The image to classify. class_names : DataFrame A dataframe containing the class names. model_path : str The path to the torchscript model. Returns: out : DataFrame A dataframe containing the class name and confidence score.

Python Code

from flojoy import DataFrame, Image, flojoy


@flojoy(
    deps={
        "torch": "2.0.1",
        "torchvision": "0.15.2",
    }
)
def TORCHSCRIPT_CLASSIFIER(
    input_image: Image, class_names: DataFrame, model_path: str
) -> DataFrame:
    """Execute a TorchScript classifier against an input image.

    Parameters
    ----------
    input_image : Image
        The image to classify.
    class_names : DataFrame
        A dataframe containing the class names.
    model_path : str
        The path to the torchscript model.

    Returns
    -------
    DataFrame
        A dataframe containing the class name and confidence score.
    """

    import numpy as np
    import pandas as pd
    import PIL.Image
    import torch
    import torchvision

    # Load model
    model = torch.jit.load(model_path)
    channels = [input_image.r, input_image.g, input_image.b]
    mode = "RGB"

    if input_image.a is not None:
        channels.append(input_image.a)
        mode += "A"

    input_image_pil = PIL.Image.fromarray(
        np.stack(channels).transpose(1, 2, 0), mode=mode
    ).convert("RGB")
    input_tensor = torchvision.transforms.functional.to_tensor(
        input_image_pil
    ).unsqueeze(0)

    # Run model
    with torch.inference_mode():
        output = model(input_tensor)

    # Get class name and confidence score
    _, pred = torch.max(output, 1)
    class_name = class_names.m.iloc[pred.item()].item()
    confidence = torch.nn.functional.softmax(output, dim=1)[0][pred.item()].item()

    return DataFrame(
        df=pd.DataFrame({"class_name": [class_name], "confidence": [confidence]})
    )

Find this Flojoy Block on GitHub

Example App

Having problems with this example app? Join our Discord community and we will help you out!

Press enter or space to select a node.You can then use the arrow keys to move the node around. Press delete to remove it and escape to cancel.

Press enter or space to select an edge. You can then press delete to remove it or escape to cancel.

Download this Flojoy App

In this example, the TORCHSCRIPT_CLASSIFIER loads the user-provided .torchscript model as well as the .csv table that maps class indices to class names. The model is then used to classify the provided input image, and it outputs the class name of the predicted class as well as a confidence score.