姿态估计2-06：PVNet(6D姿态估计)-源码无死角解析（2）-数据读取，预处理

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

代码注释

根据上篇博客的介绍，我们可以看到 train_net.py 中的如下代码：

    # 创建训练以及评估数据集迭代器
    train_loader = make_data_loader(cfg, is_train=True, max_iter=cfg.ep_iter)
    val_loader = make_data_loader(cfg, is_train=False)

本人的讲解都是基于 linemod 数据集，所以上述迭代器的创建是在data/lib/datasets/linemod/pvnet.py中实现，该代码的注释如下：

import torch.utils.data as data
from pycocotools.coco import COCO
import numpy as np
import os
from PIL import Image
from lib.utils.pvnet import pvnet_data_utils, pvnet_linemod_utils, visualize_utils
from lib.utils.linemod import linemod_config
from lib.datasets.augmentation import crop_or_padding_to_fixed_size, rotate_instance, crop_resize_instance_v1
import random
import torch
from lib.config import cfg


class Dataset(data.Dataset):

    def __init__(self, ann_file, data_root, split, transforms=None):
        super(Dataset, self).__init__()
        # data_root = data/linemod/cat/JPEGImages, 训练或者测试图片路径
        self.data_root = data_root
        # 表示该数据集是分割出来的训练集还是测试集
        self.split = split
        # 加载COCO注释文件ann_file = data/linemod/cat/train.json
        self.coco = COCO(ann_file)
        #  获得图片的所有ids
        self.img_ids = np.array(sorted(self.coco.getImgIds()))
        # 用于数据预处理，
        self._transforms = transforms
        # 赋值配置信息
        self.cfg = cfg

    def read_data(self, img_id):
        # 根据图片的id，获得对应图像注释的ids
        ann_ids = self.coco.getAnnIds(imgIds=img_id)
        # 因为对于每张图片，只有一个ann_ids，所以后面加了[0]
        anno = self.coco.loadAnns(ann_ids)[0]
        # 获得图像的路径
        path = self.coco.loadImgs(int(img_id))[0]['file_name']
        inp = Image.open(path)
        # 获去图片标出的2d关键点以及3d关键点坐标
        kpt_2d = np.concatenate([anno['fps_2d'], [anno['center_2d']]], axis=0)
        # 根据注释的cls，获得类别的idx标签
        cls_idx = linemod_config.linemod_cls_names.index(anno['cls']) + 1
        # 数据中存在三种type[real, fuse, render]
        mask = pvnet_data_utils.read_linemod_mask(anno['mask_path'], anno['type'], cls_idx)

        return inp, kpt_2d, mask

    def __getitem__(self, index_tuple):
        # 获得图片索引，以及高宽(随机获得，用于数据增强)，
        index, height, width = index_tuple
        # 获得图片id
        img_id = self.img_ids[index]
        # 获得图片像素，2d关键点以及 mask
        img, kpt_2d, mask = self.read_data(img_id)
        # 如果为训练则进行数据增强
        if self.split == 'train':
            inp, kpt_2d, mask = self.augment(img, mask, kpt_2d, height, width)
        else:
            inp = img

        # 如果设置了transforms则执行
        if self._transforms is not None:
            inp, kpt_2d, mask = self._transforms(inp, kpt_2d, mask)

        # 根据mask计算出属于mask中每个像素到对应关键点的vector, 即论文中的Vk(p)
        vertex = pvnet_data_utils.compute_vertex(mask, kpt_2d).transpose(2, 0, 1)


        ret = {'inp': inp, 'mask': mask.astype(np.uint8), 'vertex': vertex, 'img_id': img_id, 'meta': {}}
        # visualize_utils.visualize_linemod_ann(torch.tensor(inp), kpt_2d, mask, True)

        return ret

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

    def augment(self, img, mask, kpt_2d, height, width):
        """
        数据增强
        :param img: 像素
        :param mask: 图像掩码
        :param kpt_2d: 二维关键点坐标，最后一个为中心坐标
        :return:
        """
        # add one column to kpt_2d for convenience to calculate
        # kpt_2d[9,2]-->hcoords[9,3],增加一列，并且该列全部为1
        hcoords = np.concatenate((kpt_2d, np.ones((9, 1))), axis=-1)
        # 对像素的格式进行转换
        img = np.asarray(img).astype(np.uint8)
        # 计算是否存属于前景物体，即foregroun>0
        foreground = np.sum(mask)
        # randomly mask out to add occlusion
        # foreground>0,表示存在前景物体
        if foreground > 0:
            # 进行随机[self.cfg.train.rotate_min, self.cfg.train.rotate_max]之间的角度旋转，
            # img, mask, mask 都会进行旋转
            img, mask, hcoords = rotate_instance(img, mask, hcoords, self.cfg.train.rotate_min, self.cfg.train.rotate_max)
            # 随机进行crop, resize操作
            img, mask, hcoords = crop_resize_instance_v1(img, mask, hcoords, height, width,
                                                         self.cfg.train.overlap_ratio,
                                                         self.cfg.train.resize_ratio_min,
                                                         self.cfg.train.resize_ratio_max)
        # 如果只有背景，没有前景。则进行crop，padding操作，把图片变换到固定大小
        else:
            img, mask = crop_or_padding_to_fixed_size(img, mask, height, width)
        kpt_2d = hcoords[:, :2]

        return img, kpt_2d, mask

代码讲解

总的说，就是获得训练图像的像素，以及对应的标签，其中标签包含了两个部分，分别为vertex（向量域），以及语义分割得到的mask掩码。vertex与mask的 $W,H$ 是相等的。