pngWriter.py 文件源码

def save(self, data):
        """Takes in an array of CYX pixel values and writes them to a png

        :param data: a CYX or YX array with C being the rgb channels for each pixel value
        """

        # check for rgb, rgba, or r
        if len(data.shape) == 3:
            assert data.shape[0] in [4, 3, 2, 1]
            # if three dimensions, transpose to YXC (imsave() needs it in these axes)
            data = np.transpose(data, (1, 2, 0))
            # if there's only one channel, repeat across the next two channels
            if data.shape[2] == 1:
                data = np.repeat(data, repeats=3, axis=2)
            elif data.shape[2] == 2:
                data = np.pad(data, ((0, 0), (0, 0), (0, 1)), 'constant')
        elif len(data.shape) != 2:
            raise ValueError("Data was not of dimensions CYX or YX")

        imsave(self.file_path, data, format="png")