YOLOv5-7.0 インスタンスのセグメンテーション + TensorRT デプロイメント

1: はじめに

YOLOv5 を組み合わせてセグメンテーションタスクを実行し、TensorRT をデプロイすることは、挑戦的でエキサイティングなタスクです。セグメンテーションタスクでは、モデルがターゲットの存在を検出するだけでなく、ターゲットの境界と輪郭を正確に理解し、対応するカテゴリラベルを各ピクセルに割り当てる必要があるため、コンピュータはターゲットをより深く理解して解釈できます。画像。高性能の深層学習推論エンジンとして、TensorRT はモデル推論プロセスを大幅に加速し、リアルタイムアプリケーションに強力なサポートを提供します。

この記事では、YOLOv5 とセグメンテーションタスクを組み合わせて、オブジェクト検出とピクセルレベルのセマンティックセグメンテーションを同時に実現する方法を検討します。モデル融合のテクノロジーと手順を詳細に紹介し、TensorRT を使用してモデルを最適化し、組み込みデバイスやエッジコンピューティング環境で効率的な展開を実現する方法について詳しく説明します。実験結果とパフォーマンスメトリクスを詳細に説明することで、このアプローチの有効性と可能性を実証し、YOLOv5、セグメンテーションタスク、TensorRT デプロイメントの組み合わせについての包括的な理解を読者に提供します。

2: パイソン

pycharm を開き、ターミナルに pip install labelme と入力します。
その後のラベル付け作業を容易にするために、C ドライブ -> ユーザー -> ユーザー名 -> .labelmerc ファイルを開く必要があります。ファイルを開いた後、最初の行のauto_save を true に変更します。ラベル付けボックスを容易にするために、次のようにします。ラベリングの変更を容易にするために、create_polygon を W に変更する必要があります。ボックスで edit_polygon を J に変更します。
ダウンロードが完了したら、pycharm ターミナルで labelme と入力し、データセットのフォルダーを開いてラベルを付けます。ここでは画像のデモはありません。

マークを付けた後、json ファイルを txt ファイルに変換し、必要なコードを以下に置く必要があります。

import json
import os
import argparse
from tqdm import tqdm


def convert_label_json(json_dir, save_dir, classes):
    json_paths = os.listdir(json_dir)
    classes = classes.split(',')

    for json_path in tqdm(json_paths):
        # for json_path in json_paths:
        path = os.path.join(json_dir, json_path)
        with open(path, 'r') as load_f:
            json_dict = json.load(load_f)
        h, w = json_dict['imageHeight'], json_dict['imageWidth']

        # save txt path
        txt_path = os.path.join(save_dir, json_path.replace('json', 'txt'))
        txt_file = open(txt_path, 'w')

        for shape_dict in json_dict['shapes']:
            label = shape_dict['label']
            label_index = classes.index(label)
            points = shape_dict['points']

            points_nor_list = []

            for point in points:
                points_nor_list.append(point[0] / w)
                points_nor_list.append(point[1] / h)

            points_nor_list = list(map(lambda x: str(x), points_nor_list))
            points_nor_str = ' '.join(points_nor_list)

            label_str = str(label_index) + ' ' + points_nor_str + '\n'
            txt_file.writelines(label_str)


if __name__ == "__main__":
    """
    python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_ringsts --classes "cat,dogs"
    """
    parser = argparse.ArgumentParser(description='json convert to txt params')
    parser.add_argument('--json-dir', type=str, default=r'json', help='json path dir')
    parser.add_argument('--save-dir', type=str, default=r'txt',help='txt save dir')
    parser.add_argument('--classes', type=str,default="1", help='classes')
    args = parser.parse_args()
    json_dir = args.json_dir
    save_dir = args.save_dir
    classes = args.classes
    convert_label_json(json_dir, save_dir, classes)

txt ファイルに変換した後、データセットを分割してトレーニングを実行します (スキルがある場合はこのステップで十分です。ここでは説明しません)

トレーニングした best.pt を gen_wts.py を通じて wts ファイルに変換します。便宜上、best.pt をディレクトリに置き、ターミナルに次のように入力します: python gen_wts.py -w best.pt

gen_wts.pyのコードは以下の通りです

import sys
import argparse
import os
import struct
import torch
from utils.torch_utils import select_device


def parse_args():
    parser = argparse.ArgumentParser(description='Convert .pt file to .wts')
    parser.add_argument('-w', '--weights', required=True,
                        help='Input weights (.pt) file path (required)')
    parser.add_argument(
        '-o', '--output', help='Output (.wts) file path (optional)')
    parser.add_argument(
        '-t', '--type', type=str, default='detect', choices=['detect', 'cls'],
        help='determines the model is detection/classification')
    args = parser.parse_args()
    if not os.path.isfile(args.weights):
        raise SystemExit('Invalid input file')
    if not args.output:
        args.output = os.path.splitext(args.weights)[0] + '.wts'
    elif os.path.isdir(args.output):
        args.output = os.path.join(
            args.output,
            os.path.splitext(os.path.basename(args.weights))[0] + '.wts')
    return args.weights, args.output, args.type


pt_file, wts_file, m_type = parse_args()
print(f'Generating .wts for {m_type} model')

# Initialize
device = select_device('cpu')
# Load model
print(f'Loading {pt_file}')
model = torch.load(pt_file, map_location=device)  # load to FP32
model = model['ema' if model.get('ema') else 'model'].float()

if m_type == "detect":
    # update anchor_grid info
    anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
    # model.model[-1].anchor_grid = anchor_grid
    delattr(model.model[-1], 'anchor_grid')  # model.model[-1] is detect layer
    # The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.
    model.model[-1].register_buffer("anchor_grid", anchor_grid)
    model.model[-1].register_buffer("strides", model.model[-1].stride)

model.to(device).eval()

print(f'Writing into {wts_file}')
with open(wts_file, 'w') as f:
    f.write('{}\n'.format(len(model.state_dict().keys())))
    for k, v in model.state_dict().items():
        vr = v.reshape(-1).cpu().numpy()
        f.write('{} {} '.format(k, len(vr)))
        for vv in vr:
            f.write(' ')
            f.write(struct.pack('>f', float(vv)).hex())
        f.write('\n')

3. TensorRT

YOLOv5-7.0wang に対応する tensorrt 分割バージョンをダウンロードします-xinyu/tensorrtx: TensorRT ネットワーク定義 API を使用した一般的な深層学習ネットワークの実装 (github.com)
cmakeを使って解凍しますが、面倒なら自分で設定すれば大丈夫です。
TensorRT の構成については以前の記事で書きましたが、よくわからない場合は、 Windows YOLOv5-TensorRT の展開_Windows の tensorrt 展開_Mr Dinosaur のブログ - CSDN ブログを参照してください。
config.h を開いて、検出カテゴリと画像サイズを変更します。
yolov5_seg.cpp を開き、main 関数を見つけてファイルパスを変更します。
ソースコードの実行に問題がある場合 (私もそうです)、もちろん、ソースコードを自分で変更してエンジンファイルを生成することもできます。エンジンファイルが生成された後、分割テストを実行できます。
セグメンテーション結果の推論速度は平均的で、検出よりも遅くなります。

4. まとめ

デプロイの場合、基本的にはこれらの操作です。インターフェイスをカプセル化して、後で簡単に呼び出すことができます。必要に応じて後で更新します。
セグメンテーション速度は検出速度よりも約 10ms 遅いため、速度要件がある場合はよく考える必要があります。
セグメンテーションは大きなターゲットに適していますが、小さなターゲットを検出したい場合は、ターゲット検出を使用します。

半年以上更新してませんでした、ファンの皆様には申し訳ありませんが、随時更新していきます！