hy clear Blog

【YOLO v11】Tensorflow.jsで物体検出/OBB/セグメンテーションを実行し表示する【React】

物体検出TensorFlowフロントエンド

2025/01/09

2025/01/09

📰 アフィリエイト広告を利用しています

ブラウザでYOLO v11を実行するメモ。
物体検出、OBB(回転のある物体検出)、セグメンテーションの3つの検出方法について

デモサイト

https://yu2728.github.io/ObjectDetectionInBrowser/

github

https://github.com/yu2728/ObjectDetectionInBrowser

モデルの準備

Tensorflow.jsで使える形式のモデルを出力します。
公式のDockerイメージではCUDAを使用しているのかNvidiaのGPUがないとエラーになりました。

Docker Imageの準備

公式のDockerイメージがあるので使用します。

docker pull ultralytics/ultralytics:latest

YOLOコンテナの作成

jupyter labを使いたいので8888のポートを開けています。

docker run -it --ipc=host -p 8888:8888 --gpus all --name yolo ultralytics/ultralytics:latest

jupyter labを設定する

コンテナを起動後、open in terminalで以下のコマンドを実行。

pip install jupyterlab
jupyter lab  --allow-root --ip 0.0.0.0

モデルを出力

検出したいモデルをTensorflow.js形式で出力する。

検出方法

モデル名

物体検出

yolo11n.pt

OBB

yolo11n-obb.pt

セグメンテーション

yolo11n-seg.pt

from ultralytics import YOLO

# モデル名(yolo11n.pt, yolo11n-obb.pt, yolo11n-seg.pt)
# 精度のいいモデルを使うときはnを変更
model = YOLO("yolo11n.pt")

model.export(format="tfjs")

実行するとmodel.json, metadata.yaml, 分割データが出力されるので圧縮してダウンロードする

!zip -r yolo11n.zip yolo11n_web_model

[共通]ライブラリのインストール

必要なライブラリをインストールします。
js-yamlはYOLOのmetadata.yamlを解析するの処理に使用します。
obbを使用する場合はturf.jsも追加します。

npm install @tensorflow/tfjs js-yaml @types/js-yaml @turf/turf

[共通]モデルとメタデータの読み込み

モデルを読み込みます。読み込みの処理・インターフェイスはすべてのモデルで同一です。まずメタデータのモデルを定義します。

metadata.tsx
import { load } from "js-yaml"
/**
 * モデルの種類
 */
export enum ModelTaskType {
    DETECT = 'detect',
    SEGMENT = 'segment',
    ORIENTED = 'oriented'
}

/**
 * YOLOのメタデータのデータ型
 */
export interface YOLOMetadata {
  description: string;
  author: string;
  date: string;
  version: string;
  license: string;
  docs: string;
  stride: number;
  task: ModelTaskType;
  batch: number;
  imgsz: [number, number];
  names: { [key: number]: string };
}
/**
 * メタデータを読み込む
 * @param {string} path メタデータのフォルダパス.最後にスラッシュはつけない
 * @returns {YOLOMetadata} メタデータをロードするPromiseを返す
 */
export async function loadMetadata(path: string): Promise<YOLOMetadata | null> {
  let metadata: YOLOMetadata | null = null;
  await fetch(`${path}/metadata.yaml`)
    .then((response) => response.text())
    .then((text) => load(text))
    .then((yamlData) => (metadata = yamlData as YOLOMetadata))
    .catch((error) => console.error("YAML読み込みエラー:", error));
  return metadata;
}

　YOLOのモデルを読み込んでウォームアップを実行します。最初の実行に時間がかかるためです。imgszはmetadataから取得した画像のサイズを渡します。metadataで定義されたimgsz以外で推論を実行するとエラーになります。

model.tsx
import "@tensorflow/tfjs-backend-cpu";
import "@tensorflow/tfjs-backend-webgl";

import * as tf from "@tensorflow/tfjs";


/**
 * YOLOのモデルをロードし、初期化を行う
 * @param {string} path モデルが保存されているフォルダパス.最後にスラッシュはつけない
 * @param {[number, number]} imgsz metadataで取得した画像サイズ
 * @returns {Promise<tf.GraphModel<string | tf.io.IOHandler>>} モデルをロードするPromiseを返す
 */
export async function loadYOLOModel(
  path: string,
  imgsz: [number, number]
): Promise<tf.GraphModel<string | tf.io.IOHandler>> {
  const model = await tf.loadGraphModel(`${path}/model.json`);
  // warm up
  tf.tidy(() => {
    const zeroTensor = tf.zeros([1, imgsz[0], imgsz[1], 3], "float32");
    model.execute(zeroTensor);
  });
  return model;
}

[共通]画像の準備

画像を推論できるように加工します。
注意点として、画像のサイズはmetadataのimgszと同じあることと、縦横比を変えてしまうと見ている限り精度が落ちることです。

そのため、縦横比を維持したまま入力サイズに合わせます。
のちに表示するときは復元する必要があります。

convert_inage_element.ts
import * as tf from "@tensorflow/tfjs";

/**
 * ImageElementをimgszのCanvasに変換します。
 * 縦横比を維持したままリサイズし、足りない部分は黒で埋めます。
 * @param {HTMLImageElement | HTMLCanvasElement | HTMLVideoElement} image 変換したい画像の要素
 * @param imgsz YOLOのmetadataで取得した画像サイズ
 * @returns {HTMLCanvasElement} 変換後のCanvas要素
 */
export const convertImageElement = (
  image: HTMLImageElement | HTMLCanvasElement | HTMLVideoElement,
  imgsz: [number, number]
) => {
  const canvas = document.createElement("canvas");
  const context = canvas.getContext("2d");
  if (!context) {
    throw new Error("Canvas context is not available");
  }
  canvas.width = imgsz[0];
  canvas.height = imgsz[1];

  const originalWidth =
    image instanceof HTMLVideoElement ? image.videoWidth : image.width;
  const originalHeight =
    image instanceof HTMLVideoElement ? image.videoHeight : image.height;

  const scale = Math.min(imgsz[0] / originalWidth, imgsz[1] / originalHeight);

  const newWidth = originalWidth * scale;
  const newHeight = originalHeight * scale;

  context.fillStyle = "black";
  context.fillRect(0, 0, imgsz[0], imgsz[1]);

  context.drawImage(image, 0, 0, newWidth, newHeight);

  return canvas;
};

/**
 * ImageElementをTensorに変換します。
 * @param {HTMLImageElement | HTMLCanvasElement | HTMLVideoElement} image 変換したい画像の要素
 * @param imgsz YOLOのmetadataで取得した画像サイズ
 * @returns {tf.Tensor<tf.Rank>} 変換後のTensor
 */
export const tensorFromPixel = (
  image: HTMLImageElement | HTMLCanvasElement | HTMLVideoElement,
  imgsz: [number, number]
): tf.Tensor<tf.Rank> => {
  let imageTensor = tf.browser
    .fromPixels(image)
    .toFloat()
    .div(tf.scalar(255.0));
  imageTensor = imageTensor.resizeBilinear(imgsz);
  return (imageTensor = imageTensor.expandDims(0));
};

[共通]物体検出を実行する

モデルと画像の準備ができたので、推論を実行しBBOXを取得します。
推論を実行するまではすべて共通の処理ですが、結果の形式が違うのでここで分岐します。セグメンテーションのみ結果が配列になります。詳細は変換する処理で解説します。

tf.tensorを処理するとき、tf.tidyを使用しないとメモリリークが発生します。

predict.ts
import "@tensorflow/tfjs-backend-cpu";
import "@tensorflow/tfjs-backend-webgl";

import * as tf from "@tensorflow/tfjs";
import { resultToDetectBbox } from "./result_to_detect";
import { resultToOrientedBbox } from "./result_to_obb";
import { resultToSegBbox } from "./result_to_seg";
import { DetectBbox, ModelTaskType, OrientedBbox, SegBbox } from "./types";

/**
 * 推論を行い、対応したBBOXを返す
 * @param {ModelTaskType} taskType モデルのタスクタイプ
 * @param {tf.GraphModel<string | tf.io.IOHandler>} model YOLOのモデル
 * @param {tf.Tensor} imageTensor 推論する画像のTensor
 * @param {number} labelCount ラベルの数
 * @param {number} minScore 検出する最低スコア
 * @param {number | undefined} targetId 絞りこむときはラベルのIDを指定
 * @returns {DetectBbox[] | SegBbox[] | OrientedBbox[]} BBOXの配列
 */
export const predict = (
  taskType: ModelTaskType,
  model: tf.GraphModel<string | tf.io.IOHandler>,
  imageTensor: tf.Tensor,
  labelCount: number,
  minScore: number,
  targetId: number | undefined
): DetectBbox[] | SegBbox[] | OrientedBbox[] => {
  const maxOutputSize = 200;
  const iouThreshold = 0.5;
  const bbox: DetectBbox[] | SegBbox[] | OrientedBbox[] = tf.tidy(() => {
    switch (taskType) {
      case ModelTaskType.DETECT: {
        const result = model!.predict(imageTensor) as tf.Tensor<tf.Rank>;
        return resultToDetectBbox(result, labelCount, maxOutputSize, iouThreshold, minScore, targetId) as [];
      }
      case ModelTaskType.ORIENTED: {
        const result = model!.predict(imageTensor) as tf.Tensor<tf.Rank>;
        return resultToOrientedBbox(result, labelCount, maxOutputSize, iouThreshold, minScore, targetId) as [];
      }
      case ModelTaskType.SEGMENT: {
        const result = model!.predict(imageTensor) as tf.Tensor<tf.Rank>[];
        return resultToSegBbox(result, labelCount, maxOutputSize, iouThreshold, minScore, targetId) as [];
      }
    }
  });
  return bbox;
};

結果のデータ構造

出力されたデータの構造について、物体検出では出力は [1, 4+ラベル数, 検出した数]の形式で出力されます。

[
　[x, x, ...], 
　[y, y, ...], 
　[w, w, ...], 
　[h, h, ...], 
　[class1_score, class1_score, ...],
　[class2_score, class2_score, ...],
  ...
]

出力されたデータは(x, y)座標が中心となる幅がw、高さがhのバウンティングボックス(bbox)です。そのbboxのclass_scoreの値がラベルの数だけ続きます。デフォルトのモデルだと80個のnamesがあるのでbboxの4つを合わせて[1, 84]の形式になります。推論した結果のbboxが8400検出されるので[1, 84, 8400]となります。

物体検出以外でも上記のデータが基本となり、必要なデータを追加した構造で出力されます。

結果を加工

検出されたbboxの中にはスコアが閾値以下のbboxや重なっているbboxが大量に含まれています。Non-Maximum Suppressionアルゴリズムで不要なデータを削除します。これはTensorflow.jsのtf.image.nonMaxSuppression()で行えるのでこの関数に必要なデータの形式に変換していきます。

tf.image.nonMaxSuppression(boxes, scores, maxOutputSize, iouThreshold, scoreThreshold)

iouThresholdは二つのbboxがどの程度重なっている場合に同じものとして検出するかの閾値です。
tf.image.nonMaxSuppression()では[y1, x1, y2, x2]のデータ形式が必要になるのでそれに合うようにデータを作成します。

物体検出

types.ts
/**
 * 物体検出のバウンディングボックス
 */
export interface DetectBbox {
  x: number;
  y: number;
  w: number;
  h: number;
  label: number;
  score: number;
}
result_to_detect.ts

import "@tensorflow/tfjs-backend-cpu";
import "@tensorflow/tfjs-backend-webgl";

import * as tf from "@tensorflow/tfjs";
import { DetectBbox } from "./types";

/**
 * 物体検出のBBOXに変換する
 * @param result
 * @param labelCount
 * @param maxOutputSize
 * @param iouThreshold
 * @param minScore
 * @param targetId
 * @returns
 */
export const resultToDetectBbox = (result: tf.Tensor<tf.Rank>, labelCount: number, maxOutputSize: number, iouThreshold: number, minScore: number, targetId?: number | undefined): DetectBbox[] => {
  const bbox = tf.tidy(() => {
    const temp = result.squeeze();
    // x, y, w, hを取り出し、[y1, x1, y2, x2]形式に変換
    const x = temp.slice([0, 0], [1, -1]);
    const y = temp.slice([1, 0], [1, -1]);
    const w = temp.slice([2, 0], [1, -1]);
    const h = temp.slice([3, 0], [1, -1]);
    const x1 = tf.sub(x, tf.div(w, 2));
    const y1 = tf.sub(y, tf.div(h, 2));
    const x2 = tf.add(x1, w);
    const y2 = tf.add(y1, h);
    const boxes = tf.stack([y1, x1, y2, x2], 2).squeeze();

    // スコアが一番いいラベルのスコア一覧を取得。 
    const maxScores = targetId === undefined ? temp.slice([4, 0], [labelCount, -1]).max(0) : temp.slice([targetId + 4, 0], [1, -1]).max(0);
    // スコアが一番いいラベルのインデックス一覧を取得。 
    const labelIndexes = targetId === undefined ? temp.slice([4, 0], [labelCount, -1]).argMax(0) : tf.fill(maxScores.shape, targetId);
    // NMSで不要なデータを削除。インデックス一覧を取得
    const boxIndexes = tf.image.nonMaxSuppression(boxes.as2D(boxes.shape[0], boxes.shape[1]!), maxScores.as1D(), maxOutputSize, iouThreshold, minScore);

    // NMSで取得したインデックスを基にデータを絞り込み
    const resultBboxes = boxes.gather(boxIndexes, 0).arraySync() as [];
    const resultScores = maxScores.gather(boxIndexes, 0).arraySync() as [];
    const resultLabels = labelIndexes.gather(boxIndexes, 0).arraySync() as [];
    // DetectBbox[]形式で結果を返す
    return resultBboxes.map((bbox, index) => {
      return {
        x: bbox[1],
        y: bbox[0],
        w: bbox[3] - bbox[1],
        h: bbox[2] - bbox[0],
        score: resultScores[index],
        label: resultLabels[index],
      };
    });
  });
  return bbox as DetectBbox[];
};

OBB

回転を含むOBBの場合はラジアンが追加され、 [1, 4+ラベル数(デフォルト15)+ラジアン, 検出数]となり、[1, 20, 21504]が出力されます。

手順は物体検出と同じですが、tf.image.nonMaxSuppression()が対応していないため作成しました。
作成しましたがtf.image.nonMaxSuppression()をそのまま使ってもそれっぽい出力になります。
テストもほぼしてないので、以下のnonMaxSuppressionWithRotateは正常に動作するかわからないです。表示された結果だけ見れば動いているように見えます。
rotationMatrixで回転した点を演算して、IoUの計算にはturf.jsを使用しています。


import "@tensorflow/tfjs-backend-cpu";
import "@tensorflow/tfjs-backend-webgl";

import * as tf from "@tensorflow/tfjs";
import * as turf from "@turf/turf";
import { Feature, GeoJsonProperties, Polygon } from 'geojson';

/**
 * Bboxの候補
 */
interface Candidate {
  score: number
  boxIndex: number
  box: Feature<Polygon, GeoJsonProperties> | null
}

/**
 * Rorateに対応したNonMaxSuppression
 * (重複する検出結果の中から最も信頼度の高いものを選択し、他を除去する手法)
 * @param {tf.Tensor2D} boxes 検出した物体のBBox
 * @param {tf.Tensor1D} score 検出した物体のスコア
 * @param {number} maxOutputSize 
 * @param {number} iouThreshold IoUの閾値
 * @param {number} scoreThreshold スコアの閾値
 * @returns {tf.Tensor1D} 重複を削除されたボックスのインデックス
 */
export default function nonMaxSuppressionWithRotate(
    boxes: tf.Tensor2D,
    score: tf.Tensor1D,
    maxOutputSize: number,
    iouThreshold: number = 0.5,
    scoreThreshold: number = 0.3
  ): tf.Tensor1D {
    let candidates: Candidate[] = []
    // Scoreの閾値以下を切り捨て。結果はindexで返すためCandidate型に変換し、元のindexを保持する
    const scoreArray = score.arraySync()
    for (let i = 0; i < scoreArray.length; i++) {
      if (scoreArray[i] > scoreThreshold) {
        candidates.push({ score: scoreArray[i], boxIndex: i, box: null } as Candidate)
      }
    }
  
    // scoreがいい順番に並び変え
    candidates.sort((a, b) => (b.score - a.score))

    // 回転した座標に変換する
    const candidatesTensor = tf.tensor1d(candidates.map(e => e.boxIndex), "int32")
    const rotatedMatrix = rotationMatrix(boxes.gather(candidatesTensor, 0))

    // turfで処理するためのポリゴンに変更
    const polygons = matrix2Polygons(rotatedMatrix)
    polygons.forEach((polygon, index) => {
      candidates[index].box = polygon
    })
  
    // 選択されたボックスのインデックスを格納する配列を初期化
    const selectedindexes: Candidate[] = [];

    // ボックスの重複を削除する処理
    while (candidates.length > 0) {
      const currentCandidate = candidates[0]
      // 残っている候補で一番いいスコアのboxは残す
      selectedindexes.push(candidates[0])
      // maxを超えていたら終わり
      if (selectedindexes.length >= maxOutputSize) {
        break;
      }
      // 一番いいスコアのboxと残っているboxを比較し、IoUの値が閾値より小さいもののみ候補に残す
      candidates.filter(box => box.boxIndex !== currentCandidate.boxIndex)
      candidates = candidates.filter((candidate) => {
        if (candidate.boxIndex === currentCandidate.boxIndex) return false;
        const iou = calculateRotatedIOU(currentCandidate.box!, candidate.box!)
        return iou < iouThreshold
      })
    }
    // 重複を削除した結果のボックスインデックスを返す
    return tf.tensor1d(selectedindexes.map(e => e.boxIndex), "int32")
  }
  
  /**
   * tarf.jsのFeature<Polygon, GeoJsonProperties>を元に
   * 二つのBBOXのIoUを計算する
   * @param {Feature<Polygon, GeoJsonProperties>} polygon_a 
   * @param {Feature<Polygon, GeoJsonProperties>} polygon_b 
   * @returns {number} IoUの数値
   */
  function calculateRotatedIOU(polygon_a: Feature<Polygon, GeoJsonProperties>, polygon_b: Feature<Polygon, GeoJsonProperties>): number {
    const intersectPolygon = turf.intersect(turf.featureCollection([polygon_a, polygon_b]))
    if (!intersectPolygon) {
      return 0
    }
    const unionPolygon = turf.union(turf.featureCollection([polygon_a, polygon_b]))
  
    if (!unionPolygon) {
      return 0
    }
  
    const iou = turf.area(intersectPolygon) / turf.area(unionPolygon)
    return iou
  
  }
  
  /**
   * bboxとradを元に回転した座標を求める
   * @param {tf.Tensor<tf.Rank>} boxes 検出結果
   * @returns {tf.Tensor<tf.Rank>} 回転した座標
   */
  function rotationMatrix(boxes: tf.Tensor<tf.Rank>) {
    const results = tf.tidy(() => {
      const [x, y, w, h, rad] = tf.split(boxes, [1, 1, 1, 1, 1], 1)
      // cosAとsinAを計算 (角度のラジアン部分)
      const cos = tf.cos(rad).squeeze();
      const sin = tf.sin(rad).squeeze();
  
      // x,yを中心としてx1~4, y1~4を求める
      const x1 = w.div(-2).squeeze()
      const x2 = w.div(2).squeeze()
      const y1 = h.div(-2).squeeze()
      const y2 = h.div(2).squeeze()
      // 回転した点を求めるp1~時計回りに進める
      const p1x = x1.mul(cos).sub(y1.mul(sin)).add(x.squeeze())
      const p1y = x1.mul(sin).add(y1.mul(cos)).add(y.squeeze())
      const p2x = x2.mul(cos).sub(y1.mul(sin)).add(x.squeeze())
      const p2y = x2.mul(sin).add(y1.mul(cos)).add(y.squeeze())
      const p3x = x2.mul(cos).sub(y2.mul(sin)).add(x.squeeze())
      const p3y = x2.mul(sin).add(y2.mul(cos)).add(y.squeeze())
      const p4x = x1.mul(cos).sub(y2.mul(sin)).add(x.squeeze())
      const p4y = x1.mul(sin).add(y2.mul(cos)).add(y.squeeze())
  
      return tf.stack([p1x, p1y, p2x, p2y, p3x, p3y, p4x, p4y])
    })
    return results
  }
  
  /**
   * tf.Tensor型のBBox行列をturf.jsのFeature<Polygon, GeoJsonProperties>の配列に変換する
   * @param {tf.Tensor<tf.Rank>} matrix 回転した座標の行列
   * @returns {Feature<Polygon, GeoJsonProperties>[]} tarf.jsのFeature<Polygon, GeoJsonProperties>の配列
   */
  function matrix2Polygons(matrix: tf.Tensor<tf.Rank>): Feature<Polygon, GeoJsonProperties>[] {
  
    const transposedMatrix = tf.tidy(() => {
      return matrix.transpose([1, 0])
  
    })
    const matrixArray = transposedMatrix.arraySync() as []
    return matrixArray.map(e => {
      return turf.polygon([[
        [e[0], e[1]],
        [e[2], e[3]],
        [e[4], e[5]],
        [e[6], e[7]],
        [e[0], e[1]],
      ]])
    })
  }

NMSが準備できたら手順は物体検出とほぼ一緒でラジアンを追加するだけです。

types.ts
/**
 * Oriented Bounding Box
 */
export interface OrientedBbox {
  x: number;
  y: number;
  w: number;
  h: number;
  label: number;
  score: number;
  r: number;
}
result_to_obb.ts
import "@tensorflow/tfjs-backend-cpu";
import "@tensorflow/tfjs-backend-webgl";

import * as tf from "@tensorflow/tfjs";
import nonMaxSuppressionWithRotate from "./non_max_suppression_with_rotate";
import { OrientedBbox } from "./types";

/**
 * OrientedBboxに変換する
 * @param result 
 * @param labelCount 
 * @param minScore 
 * @param targetId 
 * @returns 
 */
export const resultToOrientedBbox = (
    result: tf.Tensor<tf.Rank>,
    labelCount: number,
    maxOutputSize: number,
    iouThreshold: number,
    minScore: number,
    targetId?: number | undefined
  ): OrientedBbox[] => {
    const bboxes = tf.tidy(() => {
      const temp = result.squeeze()
      // x, y, w, hを取り出す
      const x = temp.slice([0, 0], [1, -1]); // x座標
      const y = temp.slice([1, 0], [1, -1]); // y座標
      const w = temp.slice([2, 0], [1, -1]); // 幅
      const h = temp.slice([3, 0], [1, -1]); // 高さ
      const r = temp.slice([(result.shape[1] ?? 0) - 1, 0], [1, -1])// R
  
      const x1 = tf.sub(x, tf.div(w, 2))
      const y1 = tf.sub(y, tf.div(h, 2))
      const x2 = tf.add(x1, w)
      const y2 = tf.add(y1, h)
      const boxes = tf.stack([y1, x1, y2, x2], 2).squeeze();
      const boxesWithR = tf.stack([x, y, w, h, r], 2).squeeze();
  
      const maxScores = targetId === undefined ? temp.slice([4, 0], [labelCount, -1]).max(0) : temp.slice([targetId + 4, 0], [1, -1]).max(0)
      const labelIndexes = targetId === undefined ? temp.slice([4, 0], [labelCount, -1]).argMax(0) : tf.fill(maxScores.shape, targetId)
      const bboxIndexs = nonMaxSuppressionWithRotate(
        boxesWithR.as2D(boxesWithR.shape[0], boxesWithR.shape[1]!),
        maxScores.as1D(), maxOutputSize, iouThreshold, minScore)
      
      const resultBboxes = boxes.gather(bboxIndexs, 0).arraySync() as number[][]
      const resultScores = maxScores.gather(bboxIndexs, 0).arraySync() as number[]
      const resultLables = labelIndexes.gather(bboxIndexs, 0).arraySync() as number[]
  
      const rs = r.squeeze().gather(bboxIndexs, 0).arraySync() as number[]
  
      return resultBboxes.map((bbox, index) => {
        return {
          x: bbox[1],
          y: bbox[0],
          w: bbox[3] - bbox[1],
          h: bbox[2] - bbox[0],
          score: resultScores[index],
          label: resultLables[index],
          r: rs[index]
        }
      }) 
    });
    return bboxes as OrientedBbox[];
  };

セグメンテーション

セグメンテーションの場合は少し特殊で、resultが配列になります。これは前述のようなデータと、画像の重みのデータが出力されるためです。

[
 [1, 116, 8400],
 [1, 160, 160, 32]
]

[1, 116, 8400]の116のデータはbboxの4つと各ラベルのスコアの84、そしてマスクの重みの32で構成されています

[1, 160, 160, 32]はプロトタイプマスクで、セグメンテーション用のマスクを生成するための 基底マスク（特徴マップのようなもの）です。これ自体は直接物体を表すわけではなく、複数の物体のマスクを合成・表現するための「元となる情報」として使われます。

各ラベルのマスクとプロトタイプマスクを合成することで最終的なマスクデータを取得できます。

以下のissueを参考にしましたが、正直よくわかっていないです。実際に表示するとうまく動いていそうです

https://github.com/ultralytics/ultralytics/issues/2953

types.ts
/**
 * セグメンテーションのバウンディングボックス
 */
export interface SegBbox {
  x: number;
  y: number;
  w: number;
  h: number;
  label: number;
  score: number;
  mask: number[][];
}
result_to_seg.ts
import "@tensorflow/tfjs-backend-cpu";
import "@tensorflow/tfjs-backend-webgl";

import * as tf from "@tensorflow/tfjs";
import { SegBbox } from "./types";

/**
 * SegmentationBBOXに変換する
 * @param result 
 * @param labelCount 
 * @param maxOutputSize 
 * @param iouThreshold 
 * @param minScore 
 * @param targetId 
 * @returns 
 */
export const resultToSegBbox = (
    result: tf.Tensor<tf.Rank>[],
    labelCount: number,
    maxOutputSize: number,
    iouThreshold: number,
    minScore: number,
    targetId?: number | undefined
  ): SegBbox[] => {
    const bbox = tf.tidy(() => {
      const temp = result[0].squeeze();
      const x = temp.slice([0, 0], [1, -1]); // x座標
      const y = temp.slice([1, 0], [1, -1]); // y座標
      const w = temp.slice([2, 0], [1, -1]); // 幅
      const h = temp.slice([3, 0], [1, -1]); // 高さ
  
      const x1 = tf.sub(x, tf.div(w, 2))
      const y1 = tf.sub(y, tf.div(h, 2))
      const x2 = tf.add(x1, w)
      const y2 = tf.add(y1, h)
      const boxes = tf.stack([y1, x1, y2, x2], 2).squeeze();
  
      const maxScores = targetId === undefined ? temp.slice([4, 0], [labelCount, -1]).max(0) : temp.slice([targetId + 4, 0], [1, -1]).max(0)
      const labelIndexes = targetId === undefined ? temp.slice([4, 0], [labelCount, -1]).argMax(0) : tf.fill(maxScores.shape, targetId)
      const boxIndexes = tf.image.nonMaxSuppression(
          boxes.as2D(boxes.shape[0], boxes.shape[1]!),
          maxScores.as1D(), maxOutputSize, iouThreshold, minScore); 
  
      const resultBboxes = boxes.gather(boxIndexes, 0).arraySync() as []
      const resultScores = maxScores.gather(boxIndexes, 0).arraySync() as []
      const resultLabels = labelIndexes.gather(boxIndexes, 0).arraySync() as []
  
      // get Mask
      // 予測したボックスのマスクを取り出す
      const vectors = temp.slice([4 + labelCount, 0], [-1, -1]).transpose([1, 0]);
      const resultVectors = vectors.gather(boxIndexes, 0);
      // 画像を一つの配列に変換
      const maskWeight = result[1].squeeze().reshape([160 * 160, 32]);
      // 変換
      const transponsedVectors = resultVectors.transpose([1, 0]);
      // マスクの重みとベクトルの内積を取る
      const dotProduct = tf.matMul(maskWeight, transponsedVectors);
      // シグモイド関数で0から1の範囲に変換
      const probabiltyMap = dotProduct.sigmoid();

      // minScore以上の確率を持つピクセルを取り出す
      const binaryMask = probabiltyMap.greater(minScore);
      const masks = binaryMask
        .transpose([1, 0])
        .reshape([resultBboxes.length, 160, 160])
        .arraySync() as [];
  
      return resultBboxes.map((bbox, index) => {
        return {
          x: bbox[1],
          y: bbox[0],
          w: bbox[3] - bbox[1],
          h: bbox[2] - bbox[0],
          score: resultScores[index],
          label: resultLabels[index],
          mask: masks[index],
        };
      });
    });
    return bbox as SegBbox[];
  };

取得できるマスクデータは160×160のサイズです。これも元の画像に合わせて拡大等する必要があります。

結果を表示

データが取得処理ができたのでcanvasに表示します。

物体検出の表示

https://github.com/yu2728/ObjectDetectionInBrowser/blob/main/src/libs/view/detect.ts

OBBの表示

https://github.com/yu2728/ObjectDetectionInBrowser/blob/main/src/libs/view/oriented.ts

セグメンテーションの表示

https://github.com/yu2728/ObjectDetectionInBrowser/blob/main/src/libs/view/seg.ts