112 lines
3.9 KiB
Python
112 lines
3.9 KiB
Python
import numpy as np
|
||
from osgeo import gdal
|
||
|
||
|
||
def read_tif(fileName):
|
||
dataset = gdal.Open(fileName)
|
||
|
||
im_width = dataset.RasterXSize # 栅格矩阵的列数
|
||
im_height = dataset.RasterYSize # 栅格矩阵的行数
|
||
im_bands = dataset.RasterCount # 波段数
|
||
im_data = dataset.ReadAsArray().astype(np.float32) # 获取数据
|
||
if len(im_data.shape) == 2:
|
||
im_data = im_data[np.newaxis, :]
|
||
im_geotrans = dataset.GetGeoTransform() # 获取仿射矩阵信息
|
||
im_proj = dataset.GetProjection() # 获取投影信息
|
||
|
||
return im_data, im_width, im_height, im_bands, im_geotrans, im_proj
|
||
|
||
|
||
def write_tif(im_data, im_width, im_height, path, im_geotrans, im_proj):
|
||
if 'int8' in im_data.dtype.name:
|
||
datatype = gdal.GDT_Byte
|
||
elif 'int16' in im_data.dtype.name:
|
||
datatype = gdal.GDT_UInt16
|
||
else:
|
||
datatype = gdal.GDT_Float32
|
||
|
||
if len(im_data.shape) == 3:
|
||
im_bands, im_height, im_width = im_data.shape
|
||
else:
|
||
im_bands, (im_height, im_width) = 1, im_data.shape
|
||
# 创建文件
|
||
driver = gdal.GetDriverByName("GTiff")
|
||
dataset = driver.Create(path, im_width, im_height, im_bands, datatype)
|
||
if dataset != None and im_geotrans != None and im_proj != None:
|
||
dataset.SetGeoTransform(im_geotrans) # 写入仿射变换参数
|
||
dataset.SetProjection(im_proj) # 写入投影
|
||
for i in range(im_bands):
|
||
dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
|
||
del dataset
|
||
|
||
|
||
fileName = 'E:/RSdata/wlk_tif/wlk_right/wlk_right_cj.tif'
|
||
im_data, im_width, im_height, im_bands, im_geotrans, im_proj = read_tif(
|
||
fileName)
|
||
|
||
|
||
mask_fileName = 'E:/RSdata/wlk_tif/wlk_right/label_right_cj.tif'
|
||
mask_im_data, mask_im_width, mask_im_height, mask_im_bands, mask_im_geotrans, mask_im_proj = read_tif(
|
||
mask_fileName)
|
||
mask_im_data = np.int8(mask_im_data)
|
||
|
||
|
||
# geotiff归一化
|
||
lower_percent = 2
|
||
upper_percent = 98
|
||
for i in range(im_bands):
|
||
arr = im_data[i, :, :]
|
||
lower = np.percentile(arr, lower_percent)
|
||
upper = np.percentile(arr, upper_percent)
|
||
stretched = np.clip((arr - lower) / (upper - lower), 0, 1)
|
||
im_data[i] = (stretched * 255).astype(np.uint8)
|
||
|
||
|
||
# 计算大图每个波段的均值和方差，train.py里transform会用到
|
||
im_data = im_data/255
|
||
for i in range(im_bands):
|
||
pixels = im_data[i, :, :].ravel()
|
||
print("波段{} mean: {:.4f}, std: {:.4f}".format(
|
||
i, np.mean(pixels), np.std(pixels)))
|
||
im_data = im_data*255
|
||
|
||
|
||
# 切成小图
|
||
a = 0
|
||
size = 448
|
||
for i in range(0, int(mask_im_height / size)):
|
||
for j in range(0, int(mask_im_width / size)):
|
||
im_cut = im_data[:, i * size:i * size + size, j * size:j * size + size]
|
||
mask_cut = mask_im_data[:, i * size:i *
|
||
size + size, j * size:j * size + size]
|
||
|
||
labelfla = np.array(mask_cut).flatten()
|
||
if np.all(labelfla == 15): # 15为NoData
|
||
print("Skip!!!")
|
||
else:
|
||
# 取5、4、3波段，注意顺序
|
||
rgb_cut = np.stack([
|
||
im_cut[4], # 第5波段
|
||
im_cut[3], # 第4波段
|
||
im_cut[2], # 第3波段
|
||
], axis=0) # shape: (3, H, W)
|
||
# 转为 (H, W, 3)
|
||
rgb_cut = np.transpose(rgb_cut, (1, 2, 0))
|
||
# 归一化到0-255并转uint8
|
||
rgb_cut = np.clip(rgb_cut, 0, 255)
|
||
rgb_cut = rgb_cut.astype(np.uint8)
|
||
# 保存为jpg
|
||
from PIL import Image
|
||
rgb_img = Image.fromarray(rgb_cut)
|
||
rgb_img.save(
|
||
f'E:/RSdata/wlk_right_448/dataset_5m_jpg/img_{a}.jpg')
|
||
|
||
# mask只取第一个波段（如果是单通道），转uint8保存为png
|
||
mask_arr = mask_cut[0] if mask_cut.shape[0] == 1 else mask_cut
|
||
mask_arr = np.clip(mask_arr, 0, 255).astype(np.uint8)
|
||
mask_img = Image.fromarray(mask_arr)
|
||
mask_img.save(f'E:/RSdata/wlk_right_448/mask_png/mask_{a}.png')
|
||
|
||
print(f'img_{a}.jpg and mask_{a}.png saved')
|
||
a = a+1
|