#はじめに
最近pytorchを覚えたので遊んでみます。

やることとしては、点群の位置合わせです。

目的としては、icpは局所解に陥りがちなイメージが個人的にはありますが、
pytorchの最先端な最適化関数なら結構いい線行くんじゃね？
というふわっとした期待の元やってみます。

流れは以下の通りです。

for 点群Bの各ポイント
   1.変換行列$P = [R|t]$を定義する。
     このパラメータの最適化を行うことで位置合わせをする。
   2.点群の最近傍を計算し、最も近いポイントの組を取得する。
   3.変換行列$P$を適用した点群Bのポイントと、
     点群Aの最近傍のポイント、この2つを損失関数で評価する。
   4.最適化処理
最適化された変換行列$P$を適用してもう一度1から実行

以下の流れで検証します。

1.同じ点群を2つ用意して、移動だけ行う（3次元のパラメータの調整）
2.同じ点群を2つ用意して、回転だけ行う（9次元のパラメータの調整）
3.同じ点群を2つ用意して、回転・移動行う（12次元のパラメータの調整）
3.異なる2つの点群を用意して、回転・移動行う（12次元のパラメータの調整）
3.他方にノイズを付与して異なる2つの点群を用意して、回転・移動行う（12次元のパラメータの調整）

#結果

・・・と、書きましたが1の段階で失敗しました。
pytorchは未だ使い慣れていないので、何かが足りてないのかもしれません。

tetst.py

    import copy
    import numpy as np
    import open3d as o3d
    import random
    import math
    import torch.nn as nn
    import torch.nn.functional as F
    import torch
    import torch.optim as optim
    import matplotlib.pyplot as plt

    epoch = 1000

    def getPLYfromNumpy(nplist):
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(nplist)
        return pcd


    def write_point_cloud(path, pcl):
        assert o3d.io.write_point_cloud(path, pcl), "write pcl error : " + path


    def read_point_cloud(path):
        pcl = o3d.io.read_point_cloud(path)
        if len(pcl.points) == 0: assert False, "load pcl error : " + path
        return pcl

    def Register_EachPoint_RT(pclA, pclB,testP,criterion,optimizer):
        history = {
            'train_loss': [],
            'test_acc': [],
        }

        transP = torch.tensor(
            [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]],
            requires_grad=True)
        params = [transP]
        optimizer = optimizer(params)

        kd_tree_A = o3d.geometry.KDTreeFlann(pclA)
        cnt = 0
        # ポイント単位でやってみる
        for j in range(epoch):
            for didx in range(len(pclB.points)):
                cnt += 1
                optimizer.zero_grad()
                # 最近傍の計算
                [_, Aidx1, _] = kd_tree_A.search_knn_vector_3d(pclB.points[didx], 1)
                ptA_sample = pclA.points[Aidx1[0]]
                ptB_sample = pclB.points[didx]

                # 　同次座標
                ptA_sample = np.array([ptA_sample[0], ptA_sample[1], ptA_sample[2], 1])
                ptB_sample = np.array([ptB_sample[0], ptB_sample[1], ptB_sample[2], 1])
                ptA_sample = ptA_sample.reshape(4, 1)
                ptB_sample = ptB_sample.reshape(4, 1)

                A_tor = torch.tensor(ptA_sample.tolist(), requires_grad=False)
                B_tor = torch.tensor(ptB_sample.tolist(), requires_grad=False)
                answer = A_tor
                output = torch.mm(transP, B_tor)

                loss = criterion(answer, output)

                loss.backward()
                optimizer.step()

                # print( j, cnt, " : 誤差 = ", loss.item(),"\n",transP)
                ls = np.linalg.norm(testP - transP.to('cpu').detach().numpy().copy())

                history['train_loss'].append(loss)
                history['test_acc'].append(ls)

        print(" : 誤差 = ", loss.item(),"\t 正解の変換行列との誤差 = ",ls)

        plt.figure()
        plt.plot(range(1, cnt + 1), history['train_loss'], label='train_loss')
        plt.xlabel('train_loss')
        plt.legend()
        plt.savefig('train_loss.png')

        plt.figure()
        plt.plot(range(1, cnt + 1), history['test_acc'], label='test_acc')
        plt.xlabel('test_acc')
        plt.legend()
        plt.savefig('test_acc.png')

        return transP

    def Register_EachPoint_T(pclA, pclB,testP,criterion,optimizer):
        history = {
            'train_loss': [],
            'test_acc': [],
        }

        transP = torch.tensor([[0.0], [0.0], [0.0]],requires_grad=True)
        params = [transP]
        optimizer = optimizer(params)

        kd_tree_A = o3d.geometry.KDTreeFlann(pclA)
        cnt = 0

        # ポイント単位でやってみる
        for j in range(epoch):
            for didx in range(len(pclB.points)):
                cnt += 1
                optimizer.zero_grad()
                # 点群Bの各ポイントに対して最も近い点群Aのポイントを取得
                [_, Aidx1, _] = kd_tree_A.search_knn_vector_3d(pclB.points[didx], 1)
                ptA_sample = pclA.points[Aidx1[0]]
                ptB_sample = pclB.points[didx]

                # 　同次座標
                ptA_sample = np.array([ptA_sample[0], ptA_sample[1], ptA_sample[2]])
                ptB_sample = np.array([ptB_sample[0], ptB_sample[1], ptB_sample[2]])
                ptA_sample = ptA_sample.reshape(3, 1)
                ptB_sample = ptB_sample.reshape(3, 1)

                # 各ポイントをTensor化
                A_tor = torch.tensor(ptA_sample.tolist(), requires_grad=False)
                B_tor = torch.tensor(ptB_sample.tolist(), requires_grad=False)

                # 点群Bを調整して、点群Aに合わせていく。
                answer = A_tor
                output = (B_tor + transP)

                #　損失計算　-> 最適化
                loss = criterion(answer, output)
                loss.backward()
                optimizer.step()

                # 正解の変換行列との比較。（0が望ましい）
                ls = np.linalg.norm(testP - transP.to('cpu').detach().numpy().copy())

                history['train_loss'].append(loss)
                history['test_acc'].append(ls)

            print(" : 誤差 = ", loss.item(), "\t 正解の変換行列との誤差 = ", ls)

            # 調整結果を反映 -> 次のループで再び最近傍計算
            nptransP = transP.to('cpu').detach().numpy().copy().reshape(1,3)
            pclB = getPLYfromNumpy(pclB.points + nptransP)


        plt.figure()
        plt.plot(range(1, cnt + 1), history['train_loss'], label='train_loss')
        plt.xlabel('train_loss')
        plt.legend()
        plt.savefig('train_loss.png')

        plt.figure()
        plt.plot(range(1, cnt + 1), history['test_acc'], label='test_acc')
        plt.xlabel('test_acc')
        plt.legend()
        plt.savefig('test_acc.png')

        return transP

    POINT_NUM = 1024

    # http://graphics.stanford.edu/data/3Dscanrep/
    pclA = read_point_cloud("bun000.ply")
    A = np.array(pclA.points)
    A = np.array(random.sample(A.tolist(), POINT_NUM))

    # 難易度が少し上の点群。たぶんこれはまだできない・・・
    # pclB = read_point_cloud("bun045.ply")
    # B = np.array(pclB.points)
    # B = np.array(random.sample(B.tolist(), POINT_NUM))
    # # ノイズ付与
    # B += np.random.randn(POINT_NUM, 3) * 0.005
    # # 点群のUnorder性（順番バラバラ）付与
    # np.random.shuffle(B)
    # pclB_sample = getPLYfromNumpy(B)

    pclA_sample = getPLYfromNumpy(A)
    T_Projection = np.array([[1, 0, 0, 0.5],
                       [0, 1, 0, 0],
                       [0, 0, 1, 0],
                       [0, 0, 0, 1]])
    T_translation = np.array([[T_Projection[0][3]], [T_Projection[1][3]], [T_Projection[2][3]]])
    pclA_trans_sample = getPLYfromNumpy(A).transform(T_Projection)

    write_point_cloud("A_before.ply", pclA_sample)
    write_point_cloud("A_rot_before.ply", pclA_trans_sample)


    def testEstimateT(pclA_sample,pclA_trans_sample,T_translation):
        optimizer = optim.Adam
        # MSELoss
        transP = Register_EachPoint_T(pclA_sample, pclA_trans_sample, T_translation, nn.MSELoss(),optimizer)
        T_res = np.array([[1, 0, 0, transP[0]],
                          [0, 1, 0, transP[1]],
                          [0, 0, 1, transP[2]],
                          [0, 0, 0, 1]])
        pclA_res = copy.copy(pclA_trans_sample)
        pclA_res = pclA_res.transform(T_res)
        write_point_cloud("TOnlytest_A_rot_after_MSELoss.ply", pclA_res)

        # # L1Loss
        # transP = Register_EachPoint_T(pclA_sample, pclA_trans_sample, T_translation, nn.L1Loss(),optimizer)
        # T_res = np.array([[1, 0, 0, transP[0]],
        #                   [0, 1, 0, transP[1]],
        #                   [0, 0, 1, transP[2]],
        #                   [0, 0, 0, 1]])
        # pclA_res = copy.copy(pclA_trans_sample)
        # pclA_res = pclA_res.transform(T_res)
        # write_point_cloud("TOnlytest_A_rot_after_L1Loss.ply", pclA_res)

    def testEstimateRT(pclA_sample,pclA_trans_sample,T_Projection):
        optimizer = optim.Adam
        # MSELoss
        transP = Register_EachPoint_RT(pclA_sample, pclA_trans_sample, T_Projection, nn.MSELoss(),optimizer)
        transP = transP.to('cpu').detach().numpy().copy()
        pclA_res = copy.copy(pclA_trans_sample)
        pclA_res = pclA_res.transform(transP)
        write_point_cloud("RTtest_A_rot_after_MSELoss.ply", pclA_res)

    testEstimateT(pclA_sample, pclA_trans_sample, T_translation)
    # testEstimateRT(pclA_sample, pclA_trans_sample, T_Projection)


    exit()

損失関数の出力結果を見てみます。
これだけ見ると、あっという間に0に収束してるように見えます。

そして最適化により出力された変換行列と、正解の変換行列を比較したものがこれ。

・・・最初ちょっとだけ下がったかと思えば、あとはずっと腹ばいです。0.5と大きめの誤差です。
点群も出力しましたがわずかーーーーに近寄っただけだったので割愛します。

#まとめ
外れ値もなく3次元しかないのにどうしてぇ・・・という感じです。
pytorchの使い方を単純に間違えている気もします。
分かる方いたらぜひご連絡を。

次はPointNetのT-netを拡張させたもの、PointNetLKについて勉強します。
https://github.com/wentaoyuan/it-net
https://www.slideshare.net/naoyachiba18/pointnetlk-robust-efficient-point-cloud-registration-using-pointnet-167874587