import rasterio
from rasterio.merge import merge
import os

inPath = '../../data/tif_crop/'

parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), inPath))
tif_files = os.listdir(parent_dir)

# 读取所有 TIF 文件并放入列表
src_files_to_mosaic = []
for tif in tif_files:
    src = rasterio.open(parent_dir + '/' + tif)
    src_files_to_mosaic.append(src)

# 拼接 TIF 文件
mosaic, out_trans = merge(src_files_to_mosaic)

# 获取拼接后的元数据
out_meta = src_files_to_mosaic[0].meta.copy()

# 更新拼接后文件的元数据
out_meta.update({
    "driver": "GTiff",
    "count": mosaic.shape[0],  # 波段数
    "crs": src_files_to_mosaic[0].crs,  # 坐标参考系统
    "transform": out_trans,  # 转换矩阵
    "width": mosaic.shape[2],  # 拼接后的图像宽度
    "height": mosaic.shape[1]  # 拼接后的图像高度
})

# 保存拼接后的结果为新文件
outPath = '/home/project/paddlers/data/tif_merge/tif_mosaic.tif'
with rasterio.open(outPath, "w", **out_meta) as dest:
    dest.write(mosaic)

# 关闭所有文件
for src in src_files_to_mosaic:
    src.close()

print(f"拼接后的文件保存为: {outPath}")