Segmentation Mask Representation

import numpy as np

# Build a 3x3 segmentation mask where each pixel holds a class index
# Classes: 0=background, 1=cat, 2=dog
mask = np.zeros((3, 3), dtype=np.float32)
mask[0, 0] = 1
mask[1, 1] = 2
mask[2, 2] = 1

# Count how many pixels belong to the 'cat' class
cat_pixels = np.sum(mask == 1)
print(cat_pixels)