Python 简版 直方图拉伸

根据需求 进行tiff的直方图拉伸,代码如下

import rasterio
import numpy as np
import matplotlib.pyplot as plt


def process_tiff(input_path, band_number, output_path=None):
    # 读取TIFF文件
    with rasterio.open(input_path) as src:
        # 获取文件中的波段数量
        if band_number > src.count:
            # 如果指定的波段号超过了文件中波段的数量，则报错
            raise ValueError(f"Band number {band_number} exceeds the number of bands in the file.")
        #读取指定波段的数据
        band_data = src.read(band_number)

        # 显示原始直方图
    plt.hist(band_data.ravel(), bins=256, range=(0, 256), color='blue', alpha=0.7, label='Original Histogram')
    # 显示原始直方图
    plt.title('Histogram of TIFF Band')
    # 设置坐标轴标签
    plt.xlabel('Pixel Value')
    # 设置坐标轴标签
    plt.ylabel('Frequency')
    # 设置图例位置
    plt.legend(loc='upper right')
    # 显示图
    plt.show()

    # 简单的直方图拉伸
    min_val = np.percentile(band_data, 1)
    max_val = np.percentile(band_data, 99)
    # 进行直方图拉伸
    band_stretched = (band_data - min_val) * 255 / (max_val - min_val)
    # 进行截断
    band_stretched = np.clip(band_stretched, 0, 255).astype(np.uint8)

    # 显示拉伸后的直方图
    plt.hist(band_stretched.ravel(), bins=256, range=(0, 256), color='red', alpha=0.7, label='Stretched Histogram')
    # 显示拉伸后的直方图
    plt.title('Histogram of Stretched TIFF Band')
    # 设置坐标轴标签
    plt.xlabel('Pixel Value')
    # 设置坐标轴标签
    plt.ylabel('Frequency')
    # 设置图例位置
    plt.legend(loc='upper right')
    # 显示图
    plt.show()

    # 如果有指定输出路径，则将处理后的数据写回TIFF文件
    if output_path:
        # 写入处理后的数据
        with rasterio.open(output_path, 'w', driver='GTiff',
                           height=band_stretched.shape[0], width=band_stretched.shape[1],
                           count=1, dtype=band_stretched.dtype,
                           crs=src.crs, transform=src.transform) as dst:
            # 写入处理后的数据
            dst.write(band_stretched, 1)

        # 使用示例


input_file = '输入的tiff路径'
output_file = '输出的tiff路径'
band_to_process = 1
process_tiff(input_file, band_to_process, output_file)