python 保存为tiff文件, 图像转tiff, tif文件添加投影

前言

将任意格式的图像转换成tiff格式
如指定了tiff_sample则采用与tiff_sample相同的投影和变换方法。

代码

save_tiff.py

import sys
import imageio
import numpy as np
from osgeo import gdal


def save_tiff(out_path, in_data, tiff_sample=None, dtype=None):
    """
    保存tiff文件
    :param out_path: 保存tif文件路径
    :param in_data: 原始数据路径 或 数据array(h,w), array(c,h,w) 或 array(h,w,c) 要求c<h,w
    :param tiff_sample: tif样本路径, 用于提供投影和变换
    :param dtype: gdal数据类型, 默认自动根据输入识别
    :return:
    """
    # 数据
    if isinstance(in_data, str):
        in_data = imageio.imread(in_data)

    # 调整维度
    if len(in_data.shape) == 2:
        in_data = in_data.reshape(1, *in_data.shape)
        dim = 1
    elif len(in_data.shape) == 3:
        if in_data.shape[0] > in_data.shape[-1]:
            in_data = in_data.transpose(2, 0, 1)
        dim = in_data.shape[0]
    else:
        raise NotImplementedError

    # 检测类型
    if dtype is None:
        if in_data.dtype == np.dtype('uint8'):
            dtype = gdal.GDT_Byte
        else:
            dtype = gdal.GDT_Float32

    # 参照文件
    if isinstance(tiff_sample, str):
        tiff_sample = gdal.Open(tiff_sample)

    # 创建TIFF文件
    tiff_driver = gdal.GetDriverByName('GTiff').Create(out_path, *[i for i in in_data.shape[::-1]], dtype)
    if tiff_sample:
        tiff_driver.SetProjection(tiff_sample.GetProjection())  # 投影
        tiff_driver.SetGeoTransform(tiff_sample.GetGeoTransform())  # 变换
    # 写入数据
    for i in range(dim):
        tiff_driver.GetRasterBand(i + 1).WriteArray(in_data[i])
    tiff_driver.FlushCache()
    return


if __name__ == '__main__':
    save_tiff(sys.argv[1], sys.argv[2], sys.argv[3])

使用范例：

python save_tiff.py "C:\output.tiff" "C:\input.png" "C:\sample.tiff"