bool値が格納されたnumpy arrayから、COCOデータセットのRLEを算出する

環境

• Python 3.13.1
• pycocotools 2.0.10
• numpy 2.2.4

背景

bool値が格納されたnumpy arrayから、COCOデータセットのRLEに変換したいです。
ChatGPTに質問しながらコードを書いていたのですが、結構ハマりました。
なので、自分用のメモとして最終的にできあがったコードを残しておきます。

Uncompressed RLEを算出

import numpy
import pycocotools
import pycocotools.mask
from typing import Any


def get_uncompressed_rle_from_boolean_segmentation_array(boolean_segmentation_array: numpy.ndarray) -> dict[str, Any]:
    """
    セグメンテーションを表すboolのnumpy arrayから、RLE形式(Uncompressed)のsegmentationを取得します。
    """
    height, width = boolean_segmentation_array.shape
    uint8_segmentation_array = boolean_segmentation_array.astype(numpy.uint8)

    # run-length counting
    prev = 0
    cnt = 0
    counts = []
    for v in uint8_segmentation_array.flatten(order="F"):
        if v == prev:
            cnt += 1
        else:
            counts.append(cnt)
            cnt = 1
            prev = v
    counts.append(cnt)
    return {"size": [height, width], "counts": counts}

In [3]: bool_segmentation_array = np.array(
   ...:     [
   ...:         [False, True, False],
   ...:         [True, True, False],
   ...:     ],
   ...:     dtype=bool,
   ...: )

In [4]: uncompressed_rle = get_uncompressed_rle_from_boolean_segmentation_array(bool_segmentation_array)

In [70]: uncompressed_rle
Out[70]: {'size': [2, 3], 'counts': [1, 3, 2]}
# countsの意味：Falseが1個、Trueが3個、Falseが2個（列方向に見ていく）

COCOアノテーション（Unompressed RLE）を取得

def get_coco_annotation_with_uncompressed_rle(boolean_segmentation_array: numpy.ndarray) -> dict[str, Any]:
    """
    セグメンテーションを表すboolのnumpy arrayから、COCOアノテーションを取得します。
    `segmentation`はUncompressed RLE形式です。
    """
    height, width = boolean_segmentation_array.shape
    uncompressed_rle = get_uncompressed_rle_from_boolean_segmentation_array(boolean_segmentation_array)
    compressed_rle = pycocotools.mask.frPyObjects(uncompressed_rle, height, width)

    return {
        "segmentation": uncompressed_rle,
        # `float()`を実行している理由：`pycocotools.mask.area`の戻り値はnumpy.int32であり、そのままJSONにシリアライズできないため。
        # ※ `numpy.int32`は`int`のサブクラスでない
        "area": float(pycocotools.mask.area(compressed_rle)),
        # `.tolist()`を実行している理由：`pycocotools.mask.toBbox`の戻り値はnumpy.ndarrayであり、そのままJSONにシリアライズできないため。
        "bbox": pycocotools.mask.toBbox(compressed_rle).tolist(),
    }

In [141]: coco_annotation_with_uncompressed_rle = get_coco_annotation_with_uncompressed_rle(bool_segmentation_array)

In [142]: coco_annotation_with_uncompressed_rle
Out[142]:
{'segmentation': {'size': [2, 3], 'counts': [1, 3, 2]},
 'area': 3.0,
 'bbox': [0.0, 0.0, 2.0, 2.0]}

COCOアノテーション（Compressed RLE）を取得

instances_val20217.jsonは、Uncompressed RLEで記載されていたので、本当にこれで正しくCompressed RLEになるかは自信がありません…

def get_coco_annotation_with_compressed_rle(boolean_segmentation_array: numpy.ndarray) -> dict[str, Any]:
    """
    セグメンテーションを表すboolのnumpy arrayから、COCOアノテーションを取得します。
    `segmentation`はCompressed RLE形式です。
    """
    compressed_rle = pycocotools.mask.encode(numpy.asfortranarray(boolean_segmentation_array.astype(numpy.uint8)))

    segmentation = compressed_rle.copy()
    # `decode("ascii")`が本当に正しいか自信がないです…
    segmentation["counts"] = segmentation["counts"].decode("ascii")
    return {
        "segmentation": segmentation,
        "area": float(pycocotools.mask.area(compressed_rle)),
        "bbox": pycocotools.mask.toBbox(compressed_rle).tolist(),
    }

segmentation["counts"]はbyte型です。JSONへダンプするために、文字列に変換しています。encodingは以下のissueを参考にして、asciiにしています。これが正しいかは分かりません…

ただ、算出したCompressed RLEを再度decodeしたら、期待するnumpy arrayになったので、たぶん問題ないかと思います…

In [136]: coco_annotation_with_compressed_rle = get_coco_annotation_with_compressed_rle(bool_segmentation_array)

In [138]: coco_annotation_with_compressed_rle
Out[138]:
{'segmentation': {'size': [2, 3], 'counts': '132'},
 'area': 3.0,
 'bbox': [0.0, 0.0, 2.0, 2.0]}


In [140]: pycocotools.mask.decode(coco_annotation_with_compressed_rle["segmentation"])
Out[140]:
array([[0, 1, 0],
       [1, 1, 0]], dtype=uint8)