MaskedFusion代码复现（包含可视化代码）

环境配置：

• Ubuntu 18.04
• Python 3.6
• Pytorch 1.２.０
• CUDA 10.0
• cuDNN v7.6.２

一、克隆项目：

git clone https://github.com/kroglice/MaskedFusion.git

二、创建虚拟环境

# 创建虚拟环境
conda create -n MF python=3.6
# 激活虚拟环境
conda activate MF
# 安装依赖
pip install pillow==6.2.1 opencv-python==4.1.2.30 scipy==1.4.1 numpy==1.18.0 pyyaml==5.2
conda install pytorch==1.2.0 torchvision==0.4.0 cudatoolkit=10.0 -c pytorch

三、编译knn

源码中的lib/knn需要根据pytorch和CUDA版本重新编译

cd cd MaskedFusion/lib/knn
python setup.py build
python setup.py install
cd dist
unzip knn_pytorch-0.1-py3.6-linux-x86_64.egg
cp knn_pytorch/knn_pytorch.py ../knn_pytorch/knn_pytorch.py
cp knn_pytorch/knn_pytorch.cpython-36m-x86_64-linux-gnu.so ../knn_pytorch/knn_pytorch.cpython-36m-x86_64-linux-gnu.so
cd ../../..

四、数据集软连接

# ln -s 解压缩后的数据集文件夹的路径 代码里的数据集文件夹路径
ln -s /home/cyp/project/datasets/linemod/Linemod_preprocessed /home/cyp/project/MaskedFusion/datasets/linemod/Linemod_preprocessed 
ln -s /home/cyp/project/datasets/ycb/YCB_Video_Dataset /home/cyp/project/MaskedFusion/datasets/ycb/YCB_Video_Dataset 

注：Linemod_preprocessed文件夹无需创建

五、训练

sh ./experiments/scripts/train_linemod.sh
sh ./experiments/scripts/train_ycb.sh

六、评估

sh ./experiments/scripts/eval_ycb.sh
sh ./experiments/scripts/eval_linemod.sh

七、结果可视化

## 可视化代码

infer_linemod.py

import _init_paths
import argparse
import copy
import time

import cv2
import numpy as np
import open3d as o3d
import torch.utils.data
from torch.autograd import Variable

from datasets.linemod.dataset import PoseDataset as PoseDataset_linemod
from lib.network import PoseNet, PoseRefineNet
from lib.transformations import quaternion_matrix, quaternion_from_matrix


def show_3D(image, K, target, color, name):
    d2_list = []
    for i in range(target.shape[0]):
        d2 = np.transpose(np.matmul(K, np.transpose(target[i])))
        d2 = d2 / d2[2]
        d2_list.append(d2[:2])
    d3 = np.array(d2_list, dtype=np.int)
    for i in range(d3.shape[0]):
        image = cv2.circle(image, center=(d3[i][0], d3[i][1]), radius=1, color=color)

    cv2.imwrite(opt.out_root + '/visual/{0}.jpg'.format(name), image)


# 1.创建解析器
parser = argparse.ArgumentParser()

# 2.添加参数
# 验证数据根目录
parser.add_argument('--dataset_root', type=str, default='', help='dataset root dir')
# 推理结果保存目录
parser.add_argument('--out_root', type=str, default='', help='out root dir')
# 主干网络的模型，也就是PoseNet网络模型
parser.add_argument('--model', type=str, default='', help='resume PoseNet model')
# 姿态提炼网络模型，及PoseRefineNet网络模型
parser.add_argument('--refine_model', type=str, default='', help='resume PoseRefineNet model')
opt = parser.parse_args()
# 3.解析参数
opt = parser.parse_args()

# 测试目标物体的总数目
num_objects = 13
objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
color_list = [(80, 196, 143), (38, 204, 216), (54, 133, 254), (153, 119, 239),
              (245, 97, 111), (247, 177, 63), (249, 226, 100), (244, 122, 117),
              (0, 157, 178), (2, 75, 81), (7, 128, 207), (118, 80, 5), (0, 0, 0)]

num_points = 1000  # 根据当前需要测试帧的RGB-D，生成的点云数据，随机选择其中500个点云。
iteration = 4
bs = 1

# 相机矩阵
cam_cx = 325.26110
cam_cy = 242.04899
cam_fx = 572.41140
cam_fy = 573.57043
K = np.array([[cam_fx, 0.0, cam_cx],
              [0.0, cam_fy, cam_cy],
              [0.0, 0.0, 1.0]], dtype=np.float)

# PoseNet网络模型得类构建，此时还没有进行前向传播
estimator = PoseNet(num_points=num_points, num_obj=num_objects)
estimator.cuda()
refiner = PoseRefineNet(num_points=num_points, num_obj=num_objects)
refiner.cuda()

# PoseNet模型参数加载-加载模型
estimator.load_state_dict(torch.load(opt.model))
# PoseRefineNet模型参数加载-加载模型
refiner.load_state_dict(torch.load(opt.refine_model))
estimator.eval()
refiner.eval()

# 以评估得模式加载数据
testdataset = PoseDataset_linemod('test', num_points, False, opt.dataset_root, 0.0, True)
# 转化为torch迭代器形式
testdataloader = torch.utils.data.DataLoader(testdataset, batch_size=1, shuffle=False, num_workers=8)

item_count = 0
ori_img_path = []
for obj in objlist:
    input_file = open('{0}/data/{1}/test.txt'.format(opt.dataset_root, '%02d' % obj))
    while 1:
        # 记录处理的数目
        item_count += 1
        input_line = input_file.readline()  # 0000
        # test模式下，图片序列为10的倍数则continue
        if item_count % 10 != 0:
            continue
        # 文件读取完成
        if not input_line:
            break
        if input_line[-1:] == '\n':
            input_line = input_line[:-1]
        # 把RGB图像的路径加载到self.list_rgb列表中
        ori_img_path.append('{0}/data/{1}/rgb/{2}.png'.format(opt.dataset_root, '%02d' % obj, input_line))

for i, data in enumerate(testdataloader, 0):
    # 先获得一个样本需要的数据
    points, choose, img, target, model_points, idx, mask = data
    if len(points.size()) == 2:
        print('No.{0} NOT Pass! Lost detection!'.format(i))
        continue

    points, choose, img, target, model_points, idx, mask = Variable(points).cuda(), \
                                                           Variable(choose).cuda(), \
                                                           Variable(img).cuda(), \
                                                           Variable(target).cuda(), \
                                                           Variable(model_points).cuda(), \
                                                           Variable(idx).cuda(), \
                                                           Variable(mask).cuda()

    inner_time = time.time()
    # 通过PoseNet，预测当前帧的poses
    pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx, mask)
    pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, num_points, 1)
    pred_c = pred_c.view(bs, num_points)
    how_max, which_max = torch.max(pred_c, 1)
    pred_t = pred_t.view(bs * num_points, 1, 3)

    # 从所有的poses中找到最好的那一个，及置信度最高的那个
    my_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
    my_t = (points.view(bs * num_points, 1, 3) + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
    my_pred = np.append(my_r, my_t)

    # 进行PoseRefineNet网络的循环迭代
    for ite in range(0, iteration):
        # 前面得到的结果已经不是torch格式的数据了，所以这里进行一个转换
        T = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(num_points,
                                                                                         1).contiguous().view(1,
                                                                                                              num_points,
                                                                                                              3)

        # 把姿态旋转参数r转换为矩阵形式
        my_mat = quaternion_matrix(my_r)
        R = Variable(torch.from_numpy(my_mat[:3, :3].astype(np.float32))).cuda().view(1, 3, 3)
        my_mat[0:3, 3] = my_t

        # 把points进行一个逆操作得到new_points
        new_points = torch.bmm((points - T), R).contiguous()
        # 进行提炼，得到新的姿态pose
        pred_r, pred_t = refiner(new_points, emb, idx)

        pred_r = pred_r.view(1, 1, -1)
        pred_r = pred_r / (torch.norm(pred_r, dim=2).view(1, 1, 1))
        my_r_2 = pred_r.view(-1).cpu().data.numpy()
        my_t_2 = pred_t.view(-1).cpu().data.numpy()
        # 转化为四元数的矩阵
        my_mat_2 = quaternion_matrix(my_r_2)
        # 获得偏移参数（矩阵）
        my_mat_2[0:3, 3] = my_t_2

        # my_mat第一次迭代是主干网路初始的预测结果，之后是上一次迭代之后最后输出的结果
        # my_mat_2是当前迭代refiner预测出来的结果，矩阵相乘之后，把其当作该次迭代最后的结果
        my_mat_final = np.dot(my_mat, my_mat_2)
        my_r_final = copy.deepcopy(my_mat_final)
        my_r_final[0:3, 3] = 0

        # 同样转化为四元数矩阵
        my_r_final = quaternion_from_matrix(my_r_final, True)
        my_t_final = np.array([my_mat_final[0][3], my_mat_final[1][3], my_mat_final[2][3]])

        # 把当前最后的估算结果，赋值给送入下次迭代的pose
        my_pred = np.append(my_r_final, my_t_final)
        my_r = my_r_final  # 'my_r': quaternion
        my_t = my_t_final  # 'my_t': translation

    model_points = model_points[0].cpu().detach().numpy()
    # 转化为四元数矩阵，取出其中的旋转矩阵，这里的my_r是最后预测的结果
    my_r = quaternion_matrix(my_r)[:3, :3]

    # model_points经过姿态转化之后的点云数据
    pred = np.dot(model_points, my_r.T) + my_t

    # 获得测试图片
    ori_img = cv2.imread(ori_img_path[i])
    # 保存点云
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pred)
    o3d.io.write_point_cloud(opt.out_root + '/ply/{0}.ply'.format(i), pcd)
    # 可视化
    show_3D(ori_img, K, target=pred, color=color_list[idx.item()], name='{}'.format(i))