姿态估計2-06：PVNet(6D姿态估計)-源碼無死角解析（2）-資料讀取，預處理

以下連結是個人關于PVNet(6D姿态估計) 所有見解，如有錯誤歡迎大家指出，我會第一時間糾正。有興趣的朋友可以加微信：a944284742互相讨論技術。若是幫助到了你什麼，一定要記得點贊！因為這是對我最大的鼓勵。 文 末 附 帶 \color{blue}{文末附帶} 文末附帶 公 衆 号 − \color{blue}{公衆号 -} 公衆号− 海 量 資 源 。 \color{blue}{ 海量資源}。 海量資源。

姿态估計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 W,H W,H 是相等的。