Documentation
5. Image Segmentation

Image Segmentation Guide

This guide walks you through creating and training an image segmentation model using AnyLearning. AnyLearning calls this task Image Segmentation; the literature calls it semantic segmentation. It is a computer vision task that involves classifying each pixel in an image to a specific class (foreground classes). For example, you might want to segment road markings, vehicles, pedestrians, and buildings in street scene images for autonomous driving applications. There is an additional implicit class called background. Those pixels that do not belong to any foreground class will be considered background.

Step 1: Create a Project

First, create a new project specifically for image segmentation:

  1. Click Create project
  2. Select "Image Segmentation” as the project type
  3. Provide a meaningful project name and description

Project creation

Step 2: Data Preparation

2.1. Create the Label Set

The label set defines all possible pixel classes that your model will learn to segment. For example, in a microscopy particle segmentation project, your label might be particle.

To create your label set:

  1. Navigate to the “Overview” tab
  2. Enter each class name individually in the input field
  3. Click “+” after each class name
  4. Ensure class names are descriptive and consistent

Create label

Convenient tip: If you already have an annotated dataset, you can let AnyLearning infer the label set for you without manually creating the labels. PLease see 2.2 for more detail.

2.2. Upload the Datasets

Go to the “Dataset” tab to manage your datasets.

For effective model training, split your data into three sets:

  • Training set: The largest portion (typically 70-80%) used to train the model
  • Validation set: A smaller portion (typically 10-15%) used to tune hyperparameters and prevent overfitting
  • Test set: The remaining portion (typically 10-15%) used to evaluate the final model performance

Upload Process:

  1. Go to the Dataset tab and choose the Training, Validation or Test set.
  2. Click Choose images or a .zip, or drop files straight onto the panel. Several images at once is fine, and so is one .zip per split.
  3. If your images are already annotated, upload the archive as it is: COCO, YOLO, LabelMe and AnyLabeling are all read, and the annotations come in with the images. Ticking Create classes from folder names takes the class list from the archive rather than needing step 2.1 first.
  4. Wait for the upload to finish. The image and label counts under the panel update as it goes.

Trial Dataset: A sample dataset is available on Hugging Face (opens in a new tab).

Important: Ensure masks are correctly labeled, with each pixel corresponding to a class in your label set.

Dataset upload

(1) Labelme json annotation format example:

{
  "version": "4.2.10",
  "flags": {},
  "shapes": [
    {
      "label": "cat",
      "points": [
        [100, 100],
        [150, 100],
        [150, 150],
        [100, 150]
      ],
      "group_id": null,
      "shape_type": "polygon",
      "flags": {}
    },
    {
      "label": "dog",
      "points": [
        [200, 200],
        [300, 200],
        [300, 300],
        [200, 300]
      ],
      "group_id": null,
      "shape_type": "polygon",
      "flags": {}
    }
  ],
  "imagePath": "example.jpg",
  "imageData": null,
  "imageHeight": 400,
  "imageWidth": 400
}

2.3. Label the Data

If your dataset lacks annotations, you can use the built-in labeling tool to manually label your data

  1. Click the “Label Now” button on the dataset tab
  2. Use the polygon tool to draw masks around objects
  3. Select the appropriate class label for each painted region
  4. Repeat for all images in your dataset

Labeling Tips:

  • Zoom in for fine-grained labeling
  • Be consistent in your labeling approach (rules for foreground, background)
  • Use keyboard shortcuts to speed up labeling

Pro Tip: For large datasets, consider using pre-trained models for auto-labeling and refine the masks manually to save time.

Labeling

Step 3: Model Training

Training Configuration:

  1. Go to the “Training” tab

  2. Click Start training

  3. Configure the following hyperparameters:

    • Model variant: Deeplabv3-Lightweight, Deeplabv3-Medium or Deeplabv3-Large
    • Starting weights: Default starts from the pretrained weights that ship with AnyLearning; pick an earlier model of yours to continue training from it
    • Learning rate: how far the model moves at each step (0.001 is a sensible start)
    • Batch size: how many images are processed together (4 or 8). It must not exceed the number of images in a split, or the run trains on nothing
    • Epochs: how many times the model sees the whole training set
    • Image size: leave at Model default unless your objects are small, when a larger size helps and costs speed
    • Hardware: Automatic uses the GPU when there is one. Choose CPU to leave the GPU free
    • Augmentation: flips, rotation and colour jitter, generating variations of your images so the model sees more than it was given

    Every setting has an ⓘ beside it that explains it in place.

  4. Click Start training in the dialog to begin

Create a new training

Monitor Training Progress:

  • View all training sessions in the “Training” tab
  • Click on any session to see detailed information

Create a new training

Training Metrics and Logs:

  • Monitor loss values (segmentation loss)
  • ✅ Check IoU (Intersection over Union) metrics
  • View training logs for detailed progress information
  • Watch for signs of overfitting (validation loss getting worse)

View training lgos

Step 4: Test the Trained Model

After training completes, validate your model’s performance:

  1. Go to the Models tab
  2. Click the “Try” button
  3. Upload test images that weren’t used in training
  4. Analyze the model’s predictions

Go to model

The model will display predictions as segmentation masks overlaid on the input images, highlighting each class with a distinct color.

Make predictioj

Step 5: Export the model and use with your code

Use the download button on the model's row and pick Raw Model for the PyTorch checkpoint, or ONNX Model for the exported graph. The ONNX download arrives as a zip holding the .onnx file and the config.yml the run was trained with -- unzip it before using the code below.

The inference code is shown below.

Export models

5.1. Raw (Pytorch) model usage

  • Install the necessary libraries:
pip install numpy torch torchvision Pillow opencv-python pyyaml
  • Run the code:
import yaml
import json
 
import cv2
import numpy as np
from PIL import Image
import torch
from torchvision import transforms
 
CONFIG_PATH = "/path/to/config.yml"
MODEL_PATH = "/path/to/model.pth"
IMAGE_PATH = "/path/to/image.png"
 
 
with open(CONFIG_PATH, "r") as f:
    config = yaml.safe_load(f)
 
def hex_to_rgb(hex_color: str) -> tuple:
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in range(0, len(hex_color), 2))
 
def get_transformations(img_size, mean, std):
    transform = transforms.Compose(
        [
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
        ]
    )
    return transform
 
label_set = config["data"]["label_set"]
label_id_to_name = {v["id"]: v["name"] for v in label_set}
label_id_to_color = {v["id"]: hex_to_rgb(v["color"]) for v in label_set}
 
model = torch.load(MODEL_PATH)
model.eval()
transform = get_transformations(config["data"]["img_size"],
                                mean=config["data"]["normalize"]["mean"],
                                std=config["data"]["normalize"]["std"])
 
pil_image = Image.open(IMAGE_PATH).convert('RGB')
inp = transform(pil_image)
inp = inp.unsqueeze(0)
 
with torch.no_grad():
    pred_mask = model(inp)
    pred_mask = torch.argmax(pred_mask, dim=1)[0]
 
# Visualize
np_image = np.array(pil_image)
pred_mask = pred_mask.cpu().numpy()
pred_mask = pred_mask.astype(np.uint8)
pred_mask = cv2.resize(pred_mask, (np_image.shape[1], np_image.shape[0]),
                        interpolation=cv2.INTER_NEAREST)
 
visualization_image = np.zeros_like(np_image)
for label_id, color in label_id_to_color.items():
    visualization_image[pred_mask == label_id] = color
 
# Blend with original image
alpha = 0.5
visualization_image = cv2.addWeighted(
    np_image, alpha, visualization_image, 1-alpha, 0
)
 
# save the visualization image
cv2.imwrite("visualization.png", visualization_image)
 
# formatted predictions, to polygon format
formatted_predictions = []
for label_id, name in label_id_to_name.items():
    mask = (pred_mask == label_id).astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    all_polygons = []
    for contour in contours:
        if cv2.contourArea(contour) > 5: # filter out small polygons (noise)
            all_polygons.append(contour.flatten().tolist())
    if all_polygons:
        formatted_predictions.append({
            "id": label_id,
            "label": name,
            "color": label_id_to_color[label_id],
            "points": all_polygons,
            "type": "polygon"
        })
 
# save the formatted predictions
with open("formatted_predictions.json", "w") as f:
    json.dump(formatted_predictions, f)

5.2. Exported (ONNX) model usage

  • Install the necessary libraries:
pip install numpy torch torchvision onnxruntime Pillow opencv-python pyyaml
  • Run the code:
import yaml
import json
import onnxruntime
import cv2
import numpy as np
from PIL import Image
from torchvision import transforms
 
CONFIG_PATH = "/path/to/config.yml"
MODEL_PATH = "/path/to/model.onnx"
IMAGE_PATH = "/path/to/image.png"
 
with open(CONFIG_PATH, "r") as f:
    config = yaml.safe_load(f)
 
def hex_to_rgb(hex_color: str) -> tuple:
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in range(0, len(hex_color), 2))
 
def get_transformations(img_size, mean, std):
    transform = transforms.Compose(
        [
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
        ]
    )
    return transform
 
label_set = config["data"]["label_set"]
label_id_to_name = {v["id"]: v["name"] for v in label_set}
label_id_to_color = {v["id"]: hex_to_rgb(v["color"]) for v in label_set}
 
# Initialize ONNX Runtime session
session = onnxruntime.InferenceSession(MODEL_PATH)
input_name = session.get_inputs()[0].name
 
transform = get_transformations(config["data"]["img_size"],
                              mean=config["data"]["normalize"]["mean"],
                              std=config["data"]["normalize"]["std"])
 
pil_image = Image.open(IMAGE_PATH).convert('RGB')
inp = transform(pil_image)
inp = inp.unsqueeze(0).numpy()  # Convert to numpy array for ONNX Runtime
 
# Run inference with ONNX Runtime
pred_mask = session.run(None, {input_name: inp})[0]
pred_mask = np.argmax(pred_mask, axis=1)[0]
 
# Visualize
np_image = np.array(pil_image)
pred_mask = pred_mask.astype(np.uint8)
pred_mask = cv2.resize(pred_mask, (np_image.shape[1], np_image.shape[0]),
                      interpolation=cv2.INTER_NEAREST)
 
visualization_image = np.zeros_like(np_image)
for label_id, color in label_id_to_color.items():
    visualization_image[pred_mask == label_id] = color
 
# Blend with original image
alpha = 0.5
visualization_image = cv2.addWeighted(
    np_image, alpha, visualization_image, 1-alpha, 0
)
 
# save the visualization image
cv2.imwrite("visualization.png", visualization_image)
 
# formatted predictions, to polygon format
formatted_predictions = []
for label_id, name in label_id_to_name.items():
    mask = (pred_mask == label_id).astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    all_polygons = []
    for contour in contours:
        if cv2.contourArea(contour) > 5:  # filter out small polygons (noise)
            all_polygons.append(contour.flatten().tolist())
    if all_polygons:
        formatted_predictions.append({
            "id": label_id,
            "label": name,
            "color": label_id_to_color[label_id],
            "points": all_polygons,
            "type": "polygon"
        })
 
# save the formatted predictions
with open("formatted_predictions.json", "w") as f:
    json.dump(formatted_predictions, f)