Chapter 5.3 - Link the model to Label Studio¶

Introduction¶

In this chapter, you will begin by creating an API for the model. Next, you will link the model you trained in the previous chapters to Label Studio using this API. This integration will enable the model to make predictions on the unlabeled data within Label Studio, rendering the labeling process "AI-assisted" and thus making it significantly more efficient.

The following diagram illustrates the control flow of the experiment at the end of this chapter:

flowchart TB
    extra -->|upload| labelStudioTasks
    labelStudioTasks -->|label| labelStudioAnnotations
    bento_model -->|load| fastapi
    labelStudioTasks -->|POST /predict| fastapi
    fastapi --> labelStudioPredictions
    labelStudioPredictions -->|submit| labelStudioAnnotations

    subgraph workspaceGraph[WORKSPACE]
        extra[extra-data/extra]
        bento_model[model/classifier.bentomodel]
        fastapi[src/serve_labelstudio.py]
    end

    subgraph labelStudioGraph[LABEL STUDIO]
        labelStudioTasks[Tasks]
        labelStudioAnnotations[Annotations]
        labelStudioPredictions[Predictions]
    end

    style extra opacity:0.4,color:#7f7f7f80
    style labelStudioTasks opacity:0.4,color:#7f7f7f80
    linkStyle 0 opacity:0.4,color:#7f7f7f80
    linkStyle 1 opacity:0.4,color:#7f7f7f80

Steps¶

Install FastAPI¶

Before adding the model to Label Studio, we need to create an API for the model. Python has many packages for building web frameworks. In this guide we will use FastAPI for its simplicity. Add the main fastapi[standard] dependency to the requirements.txt file:

tensorflow==2.21.0
matplotlib==3.11.0
scikit-learn==1.9.0
pyyaml==6.0.3
dvc[gs]==3.67.1
bentoml==1.4.39
pillow==12.3.0
evidently==0.7.23
label-studio==1.23.0
fastapi[standard]==0.139.2

Check the differences with Git to validate the changes:

# Show the differences with Git
git diff requirements.txt

The output should be similar to this:

diff --git a/requirements.txt b/requirements.txt
index e5f490c..5735e62 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,3 +7,4 @@ bentoml==1.4.39
 pillow==12.3.0
 evidently==0.7.23
 label-studio==1.23.0
+fastapi[standard]==0.139.2

Install the package and update the freeze file.

# Install the dependencies
pip install -r requirements.txt

# Freeze the dependencies
pip freeze --local --all > requirements-freeze.txt

# Install the dependencies
uv pip install -r requirements.txt

# Freeze the dependencies
uv pip freeze > requirements-freeze.txt

Create the model API¶

Let's start building the API which will allow Label Studio to talk with our model. Label Studio expects the following API endpoints:

• GET /health - To check if the model is running.
• POST /setup - To configure the model API. It is useful if we have multiple models and we want to pick the right one for the task.
• POST /webhook - To start jobs such as model training from the Label Studio UI to our API.
• POST /predict - To make predictions using the model.

Let's implement these endpoints.

Create a new file src/serve_labelstudio.py and add the following code:

from __future__ import annotations

import uuid
from pathlib import Path
from typing import Dict

import bentoml
from fastapi import FastAPI
from PIL import Image

# 1. Define constants
MODEL_VERSION = "v0.0.1"
DATA_FOLDER_PATH = Path("extra-data/extra")

# 2. Initialize FastAPI app
app = FastAPI()
# Load the BentoML model
bento_model = bentoml.keras.get("celestial_bodies_classifier_model")
preprocess = bento_model.custom_objects["preprocess"]
postprocess = bento_model.custom_objects["postprocess"]
model = bento_model.load_model()


# 3. Define API endpoints
@app.get("/health")
async def health():
    return {"status": "ok"}


@app.post("/setup")
async def setup(_: Dict):
    # Currently, we don't need to do anything with the setup
    # This is why we ignore the input data
    return {"model_version": MODEL_VERSION}


@app.post("/webhook")
async def webhook():
    # Currently, we don't need to do anything with the webhook
    # It is used to start training the model from the Label Studio UI
    # Refer to:
    # https://github.com/HumanSignal/label-studio-ml-backend/blob/master/label_studio_ml/api.py#L118
    return {"status": "Unknown event"}


@app.post("/predict")
async def predict(data: Dict):
    # For request and response types, refer to:
    # https://labelstud.io/guide/ml_create#3-Implement-prediction-logic
    task = data["tasks"][0]
    # This is a simplification to get the filename from the image id.
    # In a real-world scenario, you would store the image on S3
    # and fetch it here.
    filename = "".join(task["data"]["image"].split("-")[1:])
    image_path = DATA_FOLDER_PATH / filename
    image = Image.open(image_path)

    preprocessed_image = preprocess(image)
    prediction = model.predict(preprocessed_image)
    result = postprocess(prediction)
    prediction = result["prediction"]
    score = result["probabilities"][prediction]
    print(score, type(score))

    return {
        "results": [
            {
                "model_version": MODEL_VERSION,
                "score": score,
                "result": [
                    {
                        "value": {"choices": [prediction]},
                        "id": str(uuid.uuid4()),
                        "from_name": "choice",
                        "to_name": "image",
                        "type": "choices",
                    }
                ],
            }
        ],
    }

This file is structured in 3 parts:

1. Definition of some useful constants.
2. Creation of the FastAPI app and loading of the model similarly to BentoML src/serve.py.
3. Definition of the API endpoints required by Label Studio.

Check the changes¶

Check the changes with Git to ensure that all the necessary files are tracked:

# Add all the files
git add .

# Check the changes
git status

The output should look like this:

On branch main
Changes to be committed:
  (use "git restore --staged <file>..." to unstage)
        modified:   requirements-freeze.txt
        modified:   requirements.txt
        new file:   src/serve_labelstudio.py

Commit the changes to Git¶

Commit the changes to Git:

# Commit the changes
git commit -m "Add FastAPI integration to Label Studio"

Run the model API¶

Run the model API by executing the following command:

fastapi run src/serve_labelstudio.py

The FastAPI server will start at http://localhost:8000 and you can view the documentation at http://localhost:8000/docs.

Make sure to keep the server running in the background for the next steps.

Add the model to Label Studio¶

Make sure Label Studio is running at http://localhost:8080.

1. Click on the Settings button on the top right corner of the screen and select Model from the sidebar.

   

2. Click on the Connect Model button and enter the following details:

   • Name: celestial_bodies_classifier_model
   • Backend URL: http://localhost:8000

   It should look like this:

   

3. Click on Validate and Save to add the model.

You should now see your model with a status of Connected:

Label with the model¶

Now that the model is connected, you can start labeling the data.

1. Click on your project name on the top navigation bar to go back to the project view.

   

2. Click on the Label All Tasks button to start labeling the data. This will resume the labeling process from where you left off.

   You should now see the image has been labeled automatically by the model:

   

3. Review the prediction and make any necessary corrections.

4. Click on the Submit button to save the annotation.

Labeling every image by hand does not scale. The next section shows how to label the remaining data more efficiently with active learning.

Label the rest with active learning¶

Instead of labeling images one by one, you can retrieve predictions for all remaining tasks at once and focus your review on the samples the model is least certain about. This is a simple form of uncertainty sampling, a common active-learning strategy.

1. In the project view, select all tasks by checking the ID box at the top of the task list.
2. Open the batch actions dropdown and select Retrieve Predictions then click OK. The prediction process can take a few minutes.

3. Go to Settings > Annotation Settings, and in the Prelabeling section, select the predictions set created by the batch (by default v0.0.1 (1000 predictions)) as opposed to the connected model.

   Info

   The prediction count in the dropdown can take a few minutes to update while Label Studio processes the batch. Wait until the number reflects all remaining tasks before continuing.

4. Once the predictions have been fetched, click on Columns and add the Prediction Score column to the view.

5. Sort the tasks by Prediction Score in ascending order so the lowest-confidence predictions appear at the top of the list.
6. Open the dropdown menu and select Label Tasks As Displayed.

With the lowest-confidence predictions at the top of the labeling queue, review and correct them first. Those are the labels that help the model improve fastest.

Summary¶

Wow! You have successfully linked the model to Label Studio. You can now use the model to make predictions on the data you have in Label Studio to speed up the labeling process. In the next chapter, we will retrain the model using the new data we have labeled.

You fixed some of the previous issues:

• Labeling of supplemental data is accelerated with AI assistance

State of the MLOps process¶

• Labeling of supplemental data can be done systematically and uniformly
• Labeling of supplemental data is accelerated with AI assistance
• Model needs to be retrained using higher-quality data

Continue to the next chapters to address the remaining items.

Sources¶

• FastAPI
• Label Studio
• Label Studio ML Integration