姿态估计0-10：DenseFusion(6D姿态估计)-源码解析（6）-测试代码注释与详解

以下链接是个人关于DenseFusion(6D姿态估计) 所有见解，如有错误欢迎大家指出，我会第一时间纠正。有兴趣的朋友可以加微信：a944284742相互讨论技术。若是帮助到了你什么，一定要记得点赞！因为这是对我最大的鼓励。
姿态估计0-00：DenseFusion(6D姿态估计)-目录-史上最新无死角讲解https://blog.csdn.net/weixin_43013761/article/details/103053585

代码引导

在前面的章节，已经讲解了整个网络的数据预处理，迭代训练，网络框架，以及loss定义等，现在我们就剩下最后的网络测试没有讲解了，那么我们该章节就开始把，如果帅气貌美（我不知道是小哥哥还是小姐姐）的你已经看到了这里，要给我个赞啊，因为我们都一样好看嘛！

代码注释

通过前面的章节，我相信大家对整个框架流程已经十分透彻了，所以这里我直接给出测试代码注释tools/eval_linemod.py：

import _init_paths
import argparse
import os
import random
import numpy as np
import yaml
import copy
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from datasets.linemod.dataset import PoseDataset as PoseDataset_linemod
from lib.network import PoseNet, PoseRefineNet
from lib.loss import Loss
from lib.loss_refiner import Loss_refine
from lib.transformations import euler_matrix, quaternion_matrix, quaternion_from_matrix
from lib.knn.__init__ import KNearestNeighbor

parser = argparse.ArgumentParser()
# 验证数据根目录
parser.add_argument('--dataset_root', type=str, default = '', help='dataset 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()

# 测试目标物体的总数目
num_objects = 13
objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
# 根据当前需要测试帧的RGB-D，生成的点云数据，随机选择其中500个点云。
num_points = 500
# 提炼网络迭代的次数
iteration = 4
# batch_size
bs = 1
# 数据集配置网络的目录，该目录存放了models_info.yml文件，保存的x,y,轴的起始位置，以及其对应的半径（欧氏距离，作用后面有注释）
dataset_config_dir = 'datasets/linemod/dataset_config'
# 评估结果得输出目录
output_result_dir = 'experiments/eval_result/linemod'
# 和训练一样，主要为了寻找到那个最小得距离
knn = KNearestNeighbor(1)

# 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('eval', num_points, False, opt.dataset_root, 0.0, True)
# 转化为torch迭代器形式
testdataloader = torch.utils.data.DataLoader(testdataset, batch_size=1, shuffle=False, num_workers=10)

# 对称物体索引的序列号
sym_list = testdataset.get_sym_list()
# 这里是获得测试数据中，点云最多的数目，实际为500
num_points_mesh = testdataset.get_num_points_mesh()
# 计算点云距离loss的类，此时只创建了一个类对象
criterion = Loss(num_points_mesh, sym_list)
criterion_refine = Loss_refine(num_points_mesh, sym_list)

# 存储每个模型的半径标准，该值主要和计算出来的dis（所有点云的平均距离）进行对比，如果小于该值，则数目达标了，否则该张图片的测试没有达标
diameter = []
meta_file = open('{0}/models_info.yml'.format(dataset_config_dir), 'r')
meta = yaml.load(meta_file)

# 除以1000应该是归一化，至于乘以0.1或许是单位转换（瞎猜的，大家可以去深究一下，知道答案了一定要告诉我啊）
for obj in objlist:
    diameter.append(meta[obj]['diameter'] / 1000.0 * 0.1)
print(diameter)

# 记录每个目标物体合格的数目，也就是测试出来的dis要小于models_info.yml文件对应的diameter参数，初始全部为0
success_count = [0 for i in range(num_objects)]
print(success_count)
# 该处是记录每个目标物体测试的总数目
num_count = [0 for i in range(num_objects)]

# 用于记录的日志
fw = open('{0}/eval_result_logs.txt'.format(output_result_dir), 'w')

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

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

    # 通过PoseNet，预测当前帧的poses
    pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)
    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_t = my_t_final

    # Here 'my_pred' is the final pose estimation result after refinement ('my_r': quaternion, '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

    # 获得目标点云数据，其也是由model_points转换而来，但是其转换的pose是最标准的，
    target = target[0].cpu().detach().numpy()

    # 如果是对称的物体，需要做额外的处理，使用KNN找到最小的那个距离
    if idx[0].item() in sym_list:
        pred = torch.from_numpy(pred.astype(np.float32)).cuda().transpose(1, 0).contiguous()
        target = torch.from_numpy(target.astype(np.float32)).cuda().transpose(1, 0).contiguous()
        inds = knn(target.unsqueeze(0), pred.unsqueeze(0))
        target = torch.index_select(target, 1, inds.view(-1) - 1)
        dis = torch.mean(torch.norm((pred.transpose(1, 0) - target.transpose(1, 0)), dim=1), dim=0).item()
    else:
        dis = np.mean(np.linalg.norm(pred - target, axis=1))

    # 然后判断出该处理是否小于models_info.yml文件中保存的半径（欧氏距离）
    # 如果小于，表示就是合格的
    if dis < diameter[idx[0].item()]:
        success_count[idx[0].item()] += 1
        print('No.{0} Pass! Distance: {1}'.format(i, dis))
        fw.write('No.{0} Pass! Distance: {1}\n'.format(i, dis))
    else:
        print('No.{0} NOT Pass! Distance: {1}'.format(i, dis))
        fw.write('No.{0} NOT Pass! Distance: {1}\n'.format(i, dis))
    num_count[idx[0].item()] += 1

for i in range(num_objects):
    print('Object {0} success rate: {1}'.format(objlist[i], float(success_count[i]) / num_count[i]))
    fw.write('Object {0} success rate: {1}\n'.format(objlist[i], float(success_count[i]) / num_count[i]))
print('ALL success rate: {0}'.format(float(sum(success_count)) / sum(num_count)))
fw.write('ALL success rate: {0}\n'.format(float(sum(success_count)) / sum(num_count)))
fw.close()

本来想讲解一下的，我相信大家浏览一下就知道了，其实就是训练过程中的前向传播。拜拜，祝小哥哥越来越帅，小姐姐越来越美。记得要喝热水哈！