def compute_boxes(indices, binary_tag_mask):
    """Compute 2d bounding boxes for highlighted pixels"""
    H, W = binary_tag_mask.shape
    num_u = indices.max() + 1
    x_min = np.full(num_u, W - 1, dtype=np.int32)
    y_min = np.full(num_u, H - 1, dtype=np.int32)
    x_max = np.full(num_u, -1, dtype=np.int32)
    y_max = np.full(num_u, -1, dtype=np.int32)
    for y in range(H):
        for x in range(W):
            idx = indices[y, x]
            tag_is_present = binary_tag_mask[y, x]
            if tag_is_present:
                x_min[idx] = min(x_min[idx], x)
                x_max[idx] = max(x_max[idx], x)
                y_min[idx] = min(y_min[idx], y)
                y_max[idx] = max(y_max[idx], y)
    return np.stack((x_min, y_min, x_max, y_max), axis=-1)
Compute 2d bounding boxes for highlighted pixels
infinigen/tools/ground_truth/segmentation_lookup.py