ROPNet プロジェクトのトレーニングモデル net40 データセットに基づく 3D 点群の構成

プロジェクトアドレス: https://github.com/zhulf0804/ROPNet in MVP 登録チャレンジ (ICCV ワークショップ 2021) (ICCV Workshop 2021) で準優勝。プロジェクトはwin10環境でも実行可能です。
紙のアドレス: https://arxiv.org/abs/2107.02583

ネットワークの紹介: 特徴的な特徴を持つ代表的な重複点を登録に利用し、パーツ間の登録をパーツ完全な登録に変換する新しいディープラーニングモデル。コンテキストガイダンスモジュールは、pointnet によって出力された特徴に基づいて設計され、エンコーダーを使用してグローバル特徴を抽出し、ポイントオーバーラップスコアを予測します。代表的な重なり合う点をより適切に見つけるために、抽出されたグローバルフィーチャを大まかな位置合わせに使用します。次に、ポイントフィーチャを強化し、ポイントオーバーラップスコアとフィーチャマッチングに基づいて非代表的なポイントを削除するトランスフォーマーが導入されます。類似性行列は部分から完全モードで構築され、最後に重み付けされたサポートベクトルの差を使用して変換行列が推定されます。
ここに画像の説明を挿入します
導入効果: データの観点から見ると、ROPNet と RPMNet は崖のような主導的な地位を維持しています

1. 動作環境のインストール

1.1 プロジェクトのダウンロード

開きhttps://github.com/zhulf0804/ROPNet、[Download ZIP] をクリックして、指定したディレクトリにコードを抽出します。
ここに画像の説明を挿入します

1.2 依存関係のインストール

pytorchがインストールされている環境のターミナルで、ROPNet-master/srcディレクトリに入り、以下のインストールコマンドを実行します。すでにtorch環境とopen3dパッケージをインストールしている場合は、再度インストールする必要はありません。

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

pip install open3d

1.3 モデルとデータのダウンロード

modelnet40 データセットこちら [435M]
データセットをダウンロードして、次の場所に保存します。パス。
ここに画像の説明を挿入します

公式ウェブサイトの事前トレーニング済みモデル、なし。
サードパーティの事前トレーニング済みモデル: ROPNet プロジェクトを使用して、modelnet40 データセットでトレーニングされたモデル

2. キーコード

2.1 データローダー

作成者が提供するデータローダーはhttps://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip データのみをロードできます。 set の場合、返される tgt_cloud および src_cloud は基本的に点群サンプルに基づいています。 其中的self.label2cat, self.cat2label， self.symmetric_labels等对象代码实际上是没有任何作用的。

import copy
import h5py
import math
import numpy as np
import os
import torch

from torch.utils.data import Dataset
import sys

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOR_DIR = os.path.dirname(BASE_DIR)
sys.path.append(ROOR_DIR)
from utils import  random_select_points, shift_point_cloud, jitter_point_cloud, \
    generate_random_rotation_matrix, generate_random_tranlation_vector, \
    transform, random_crop, shuffle_pc, random_scale_point_cloud, flip_pc
    


half1 = ['airplane', 'bathtub', 'bed', 'bench', 'bookshelf', 'bottle', 'bowl',
         'car', 'chair', 'cone', 'cup', 'curtain', 'desk', 'door', 'dresser',
         'flower_pot', 'glass_box', 'guitar', 'keyboard', 'lamp']
half1_symmetric = ['bottle', 'bowl', 'cone', 'cup', 'flower_pot', 'lamp']

half2 = ['laptop', 'mantel', 'monitor', 'night_stand', 'person', 'piano',
         'plant', 'radio', 'range_hood', 'sink', 'sofa', 'stairs', 'stool',
         'table', 'tent', 'toilet', 'tv_stand', 'vase', 'wardrobe', 'xbox']
half2_symmetric = ['tent', 'vase']


class ModelNet40(Dataset):
    def __init__(self, root, split, npts, p_keep, noise, unseen, ao=False,
                 normal=False):
        super(ModelNet40, self).__init__()
        self.single = False # for specific-class visualization
        assert split in ['train', 'val', 'test']
        self.split = split
        self.npts = npts
        self.p_keep = p_keep
        self.noise = noise
        self.unseen = unseen
        self.ao = ao # Asymmetric Objects
        self.normal = normal
        self.half = half1 if split in 'train' else half2
        self.symmetric = half1_symmetric + half2_symmetric
        self.label2cat, self.cat2label = self.label2category(
            os.path.join(root, 'shape_names.txt'))
        self.half_labels = [self.cat2label[cat] for cat in self.half]
        self.symmetric_labels = [self.cat2label[cat] for cat in self.symmetric]
        files = [os.path.join(root, 'ply_data_train{}.h5'.format(i))
                 for i in range(5)]
        if split == 'test':
            files = [os.path.join(root, 'ply_data_test{}.h5'.format(i))
                     for i in range(2)]
        self.data, self.labels = self.decode_h5(files)
        print(f'split: {
      
      self.split}, unique_ids: {
      
      len(np.unique(self.labels))}')

        if self.split == 'train':
            self.Rs = [generate_random_rotation_matrix() for _ in range(len(self.data))]
            self.ts = [generate_random_tranlation_vector() for _ in range(len(self.data))]

    def label2category(self, file):
        with open(file, 'r') as f:
            label2cat = [category.strip() for category in f.readlines()]
            cat2label = {
    
    label2cat[i]: i for i in range(len(label2cat))}
        return label2cat, cat2label

    def decode_h5(self, files):
        points, normal, label = [], [], []
        for file in files:
            f = h5py.File(file, 'r')
            cur_points = f['data'][:].astype(np.float32)
            cur_normal = f['normal'][:].astype(np.float32)
            cur_label = f['label'][:].flatten().astype(np.int32)
            if self.unseen:
                idx = np.isin(cur_label, self.half_labels)
                cur_points = cur_points[idx]
                cur_normal = cur_normal[idx]
                cur_label = cur_label[idx]
            if self.ao and self.split in ['val', 'test']:
                idx = ~np.isin(cur_label, self.symmetric_labels)
                cur_points = cur_points[idx]
                cur_normal = cur_normal[idx]
                cur_label = cur_label[idx]
            if self.single:
                idx = np.isin(cur_label, [8])
                cur_points = cur_points[idx]
                cur_normal = cur_normal[idx]
                cur_label = cur_label[idx]
            points.append(cur_points)
            normal.append(cur_normal)
            label.append(cur_label)
        points = np.concatenate(points, axis=0)
        normal = np.concatenate(normal, axis=0)
        data = np.concatenate([points, normal], axis=-1).astype(np.float32)
        label = np.concatenate(label, axis=0)
        return data, label

    def compose(self, item, p_keep):
        tgt_cloud = self.data[item, ...]
        if self.split != 'train':
            np.random.seed(item)
            R, t = generate_random_rotation_matrix(), generate_random_tranlation_vector()
        else:
            tgt_cloud = flip_pc(tgt_cloud)
            R, t = generate_random_rotation_matrix(), generate_random_tranlation_vector()

        src_cloud = random_crop(copy.deepcopy(tgt_cloud), p_keep=p_keep[0])
        src_size = math.ceil(self.npts * p_keep[0])
        tgt_size = self.npts
        if len(p_keep) > 1:
            tgt_cloud = random_crop(copy.deepcopy(tgt_cloud),
                                    p_keep=p_keep[1])
            tgt_size = math.ceil(self.npts * p_keep[1])

        src_cloud_points = transform(src_cloud[:, :3], R, t)
        src_cloud_normal = transform(src_cloud[:, 3:], R)
        src_cloud = np.concatenate([src_cloud_points, src_cloud_normal],
                                   axis=-1)
        src_cloud = random_select_points(src_cloud, m=src_size)
        tgt_cloud = random_select_points(tgt_cloud, m=tgt_size)

        if self.split == 'train' or self.noise:
            src_cloud[:, :3] = jitter_point_cloud(src_cloud[:, :3])
            tgt_cloud[:, :3] = jitter_point_cloud(tgt_cloud[:, :3])
        tgt_cloud, src_cloud = shuffle_pc(tgt_cloud), shuffle_pc(
            src_cloud)
        return src_cloud, tgt_cloud, R, t

    def __getitem__(self, item):
        src_cloud, tgt_cloud, R, t = self.compose(item=item,
                                                  p_keep=self.p_keep)
        if not self.normal:
            tgt_cloud, src_cloud = tgt_cloud[:, :3], src_cloud[:, :3]
        return tgt_cloud, src_cloud, R, t

    def __len__(self):
        return len(self.data)

2.2 モデル設計

モデルの設計は次のとおりです。
ここに画像の説明を挿入します

2.3 損失設計

これには主に Init_loss、Refine_loss、Ol_loss が含まれます。
Init_loss は $予測点群_0 の計算に使用されます$ ターゲット点群による MSE または MAE 損失。
Refine_loss は $予測点群_{[ 1:]}$ ターゲット点群の場合重み付けされた前損失
Ol_loss は、2 つの入力点群の出力の重複スコアを計算するために使用されます。 2 つの点群の対応する点の重複スコアが同じであること。
ここに画像の説明を挿入します

具体的な実装コードは以下のとおりです。


import math
import torch
import torch.nn as nn
from utils import square_dists


def Init_loss(gt_transformed_src, pred_transformed_src, loss_type='mae'):

    losses = {
    
    }
    num_iter = 1
    if loss_type == 'mse':
        criterion = nn.MSELoss(reduction='mean')
        for i in range(num_iter):
            losses['mse_{}'.format(i)] = criterion(pred_transformed_src[i],
                                                   gt_transformed_src)
    elif loss_type == 'mae':
        criterion = nn.L1Loss(reduction='mean')
        for i in range(num_iter):
            losses['mae_{}'.format(i)] = criterion(pred_transformed_src[i],
                                                   gt_transformed_src)
    else:
        raise NotImplementedError

    total_losses = []
    for k in losses:
        total_losses.append(losses[k])
    losses = torch.sum(torch.stack(total_losses), dim=0)
    return losses


def Refine_loss(gt_transformed_src, pred_transformed_src, weights=None, loss_type='mae'):
    losses = {
    
    }
    num_iter = len(pred_transformed_src)
    for i in range(num_iter):
        if weights is None:
            losses['mae_{}'.format(i)] = torch.mean(
                torch.abs(pred_transformed_src[i] - gt_transformed_src))
        else:
            losses['mae_{}'.format(i)] = torch.mean(torch.sum(
                weights * torch.mean(torch.abs(pred_transformed_src[i] -
                                               gt_transformed_src), dim=-1)
                / (torch.sum(weights, dim=-1, keepdim=True) + 1e-8), dim=-1))

    total_losses = []
    for k in losses:
        total_losses.append(losses[k])
    losses = torch.sum(torch.stack(total_losses), dim=0)

    return losses


def Ol_loss(x_ol, y_ol, dists):
    CELoss = nn.CrossEntropyLoss()
    x_ol_gt = (torch.min(dists, dim=-1)[0] < 0.05 * 0.05).long() # (B, N)
    y_ol_gt = (torch.min(dists, dim=1)[0] < 0.05 * 0.05).long() # (B, M)
    x_ol_loss = CELoss(x_ol, x_ol_gt)
    y_ol_loss = CELoss(y_ol, y_ol_gt)
    ol_loss = (x_ol_loss + y_ol_loss) / 2
    return ol_loss


def cal_loss(gt_transformed_src, pred_transformed_src, dists, x_ol, y_ol):
    losses = {
    
    }
    losses['init'] = Init_loss(gt_transformed_src,
                               pred_transformed_src[0:1])
    if x_ol is not None:
        losses['ol'] = Ol_loss(x_ol, y_ol, dists)
    losses['refine'] = Refine_loss(gt_transformed_src,
                                   pred_transformed_src[1:],
                                   weights=None)
    alpha, beta, gamma = 1, 0.1, 1
    if x_ol is not None:
        losses['total'] = losses['init'] + beta * losses['ol'] + gamma * losses['refine']
    else:
        losses['total'] = losses['init'] + losses['refine']
    return losses

3. トレーニングと予測

まず src ディレクトリに入り、src ディレクトリ内の modelnet40_ply_hdf5_2048.zip を解凍します。
ここに画像の説明を挿入します

3.1 トレーニング

トレーニングコマンドとトレーニング出力は次のとおりです。

python train.py --root modelnet40_ply_hdf5_2048/ --noise --unseen

python 画像の説明を追加してください
トレーニングプロセス中に、work_dirs\models\checkpoints ディレクトリに 2 つのモデルファイルが生成されます
ここに画像の説明を挿入します

3.2 検証

トレーニングコマンドとトレーニング出力は次のとおりです。

python eval.py --root modelnet40_ply_hdf5_2048/  --unseen --noise  --cuda --checkpoint work_dirs/models/checkpoints/min_rot_error.pth

画像の説明を追加してください

3.3 テスト

トレーニングデータをテストするコマンドは次のとおりです。

python vis.py --root modelnet40_ply_hdf5_2048/  --unseen --noise  --checkpoint work_dirs/models/checkpoints/min_rot_error.pth

具体的な登録効果は次のとおりです。緑色の点群が入力点群、赤色の点群が参照点群、青色の点群が登録点群です。青い点群は基本的に赤い点群と一致していることがわかり、位置合わせ効果が非常に完全であると判断できます。
ここに画像の説明を挿入します

3.4 独自のデータセットを処理する

このプロジェクトに基づいて独自のデータをトレーニングおよび処理するためのチュートリアルは後で説明します。