将cornell数据集中PCD文件转换为深度图

Convert the PCD files to depth images by running

将.txt文件--->d.tiff文件

直接贴代码：

！！此代码需在终端运行，命令最下面有，自己可以看着修改一下（pycharm下端终端和CMD都可，记得要cd 进文件所在目录！！！）

import argparse
import glob
import os

import numpy as np
from imageio import imsave


import warnings

import cv2
import matplotlib.pyplot as plt
import numpy as np
from imageio import imread
from skimage.transform import rotate, resize

warnings.filterwarnings("ignore", category=UserWarning)


class Image:
    """
    Wrapper around an image with some convenient functions.
    """

    def __init__(self, img):
        self.img = img

    def __getattr__(self, attr):
        # Pass along any other methods to the underlying ndarray
        return getattr(self.img, attr)

    @classmethod
    def from_file(cls, fname):
        return cls(imread(fname))

    def copy(self):
        """
        :return: Copy of self.
        """
        return self.__class__(self.img.copy())

    def crop(self, top_left, bottom_right, resize=None):
        """
        Crop the image to a bounding box given by top left and bottom right pixels.
        :param top_left: tuple, top left pixel.
        :param bottom_right: tuple, bottom right pixel
        :param resize: If specified, resize the cropped image to this size
        """
        self.img = self.img[top_left[0]:bottom_right[0], top_left[1]:bottom_right[1]]
        if resize is not None:
            self.resize(resize)

    def normalise(self):
        """
        Normalise the image by converting to float [0,1] following the structure of the thesis 
        https://ro.ecu.edu.au/cgi/viewcontent.cgi?article=3172&context=theses
        """
        self.img = self.img.astype(np.float32) / 255.0
        #self.img -= self.img.mean()

    def resize(self, shape):
        """
        Resize image to shape.
        :param shape: New shape.
        """
        if self.img.shape == shape:
            return
        self.img = resize(self.img, shape, preserve_range=True).astype(self.img.dtype)

    def rotate(self, angle, center=None):
        """
        Rotate the image.
        :param angle: Angle (in radians) to rotate by.
        :param center: Center pixel to rotate if specified, otherwise image center is used.
        """
        if center is not None:
            center = (center[1], center[0])
        self.img = rotate(self.img, angle / np.pi * 180, center=center, mode='symmetric', preserve_range=True).astype(
            self.img.dtype)

    def show(self, ax=None, **kwargs):
        """
        Plot the image
        :param ax: Existing matplotlib axis (optional)
        :param kwargs: kwargs to imshow
        """
        if ax:
            ax.imshow(self.img, **kwargs)
        else:
            plt.imshow(self.img, **kwargs)
            plt.show()

    def zoom(self, factor):
        """
        "Zoom" the image by cropping and resizing.
        :param factor: Factor to zoom by. e.g. 0.5 will keep the center 50% of the image.
        """
        sr = int(self.img.shape[0] * (1 - factor)) // 2
        sc = int(self.img.shape[1] * (1 - factor)) // 2
        orig_shape = self.img.shape
        self.img = self.img[sr:self.img.shape[0] - sr, sc: self.img.shape[1] - sc].copy()
        self.img = resize(self.img, orig_shape, mode='symmetric', preserve_range=True).astype(self.img.dtype)


class DepthImage(Image):
    def __init__(self, img):
        super().__init__(img)

    @classmethod
    def from_pcd(cls, pcd_filename, shape, default_filler=0, index=None):
        """
            Create a depth image from an unstructured PCD file.
            If index isn't specified, use euclidean distance, otherwise choose x/y/z=0/1/2
        """
        img = np.zeros(shape)
        if default_filler != 0:
            img += default_filler

        with open(pcd_filename) as f:
            for l in f.readlines():
                ls = l.split()

                if len(ls) != 5:
                    # Not a point line in the file.
                    continue
                try:
                    # Not a number, carry on.
                    float(ls[0])
                except ValueError:
                    continue

                i = int(ls[4])
                r = i // shape[1]
                c = i % shape[1]

                if index is None:
                    x = float(ls[0])
                    y = float(ls[1])
                    z = float(ls[2])

                    #img[r, c] = np.sqrt(x ** 2 + y ** 2 + z ** 2)
                    img[r, c] = z
                else:
                    img[r, c] = float(ls[index])

        return cls(img / 1000.0)

    @classmethod
    def from_tiff(cls, fname):
        return cls(imread(fname))

    def inpaint(self, missing_value=0):
        """
        Inpaint missing values in depth image.
        :param missing_value: Value to fill in teh depth image.
        """
        # cv2 inpainting doesn't handle the border properly
        # https://stackoverflow.com/questions/25974033/inpainting-depth-map-still-a-black-image-border
        self.img = cv2.copyMakeBorder(self.img, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
        mask = (self.img == missing_value).astype(np.uint8)

        # Scale to keep as float, but has to be in bounds -1:1 to keep opencv happy.
        scale = np.abs(self.img).max()
        self.img = self.img.astype(np.float32) / scale  # Has to be float32, 64 not supported.
        self.img = cv2.inpaint(self.img, mask, 1, cv2.INPAINT_NS)

        # Back to original size and value range.
        self.img = self.img[1:-1, 1:-1]
        self.img = self.img * scale

    def normalise(self):
        """
        Normalise by subtracting the mean and clippint [-1, 1]. Then re-normalize between 0 and 1
        """
        self.img = np.clip((self.img - self.img.mean()), -1, 1)
        self.img = cv2.normalize(self.img,None,alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)



# ---------------------------------------------------------------------------------------#
# 我的路径为：E:\PycharmProjects\grasping\utils\dataset_processing\generate_cornell_depth.py
# 首先 ：cd E:\PycharmProjects\grasping\utils
# 然后输入命令：python .\dataset_processing\generate_cornell_depth.py E:\PycharmProjects\grasping\Dataset
# 在同目录下要有Dataset文件夹
# Dataset文件夹结构：
#          --Dataset
#          -------01 02...
#          --------------pcd0100.txt pcd0100cneg.txt  pcd0111cpos.txt  ...
# ---------------------------------------------------------------------------------------#


if __name__ == '__main__':
    # 创建解析器
    parser = argparse.ArgumentParser(description='Generate depth images from Cornell PCD files.')
    # 添加参数
    parser.add_argument('path', type=str, help='Path to Cornell Grasping Dataset')
    # 解析参数
    args = parser.parse_args()

    pcds = glob.glob(os.path.join(args.path, '*', 'pcd*[0-9].txt'))
    pcds.sort()
    print(args.path)

    for pcd in pcds:
        di = DepthImage.from_pcd(pcd, (480, 640))
        di.inpaint()

        of_name = pcd.replace('.txt', 'd.tiff')
        print(of_name)
        imsave(of_name, di.img.astype(np.float32))

我的路径为：E:\PycharmProjects\grasping\utils\dataset_processing\generate_cornell_depth.py

首先 ：cd E:\PycharmProjects\grasping\utils

然后输入命令：python .\dataset_processing\generate_cornell_depth.py E:\PycharmProjects\grasping\Dataset

在同目录下要有Dataset文件夹

# Dataset文件夹结构：

#          --Dataset

#          -------01 02...

#          --------------pcd0100.txt pcd0100cneg.txt  pcd0111cpos.txt  ...