# -*- coding: utf-8 -*- """ 一张大图端到端：瓦片预测 → 拼接 → 矢量化 """ import os import json import time import torch import numpy as np import cv2 from pathlib import Path from tqdm import tqdm import rasterio from rasterio.windows import Window from osgeo import gdal, ogr, osr import warnings warnings.filterwarnings('ignore') import xml.etree.ElementTree as ET def parse_xml(xml_path: str): """返回 (影像绝对路径, 输出目录绝对路径)""" root = ET.parse(xml_path).getroot() img_node = root.find(".//Input/File_Names/File_Name") out_node = root.find(".//Output/File_Name") if img_node is None or out_node is None: raise RuntimeError("XML 中缺少 Input//File_Name 或 Output//File_Name") img_path = os.path.abspath(img_node.text.strip()) out_path = os.path.abspath(out_node.text.strip()) return img_path, out_path # ---------- 模型加载 ---------- def load_model(model_path, config, device): from dinov3_deeplabv3plus import DinoV3DeepLabV3Plus print("加载模型...") model = DinoV3DeepLabV3Plus( num_classes=config['num_classes'], backbone_name=config['backbone_name'], pretrained=False, weights=config['backbone_weights'], use_4channel=config['use_4channel'] ).to(device) ckpt = torch.load(model_path, map_location=device) model.load_state_dict(ckpt.get('model_state_dict', ckpt)) model.eval() return model # ---------- 瓦片预测 ---------- def predict_tile(model, tile, device, cfg, threshold): with torch.no_grad(): out = model(tile.to(device)) if isinstance(out, dict): out = out['out'] prob = torch.softmax(out, dim=1)[0, 1].cpu().numpy() mask = (prob > threshold).astype(np.uint8) return mask, prob # ---------- 预处理 ---------- def preprocess_tile(tile_arr, cfg): # tile_arr: HWC float32 img = cv2.resize(tile_arr, (cfg['image_size'], cfg['image_size']), interpolation=cv2.INTER_LINEAR) if img.max() > 1: img = img / 65535.0 # 标准化 mean = np.array([0.430, 0.411, 0.296, 0.350]) if cfg['use_4channel'] else np.array([0.430, 0.411, 0.296]) std = np.array([0.213, 0.156, 0.143, 0.180]) if cfg['use_4channel'] else np.array([0.213, 0.156, 0.143]) for i in range(img.shape[2]): img[:, :, i] = (img[:, :, i] - mean[i]) / std[i] tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0) return tensor # ---------- 主流程 ---------- def run_one_shot(cfg: dict): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 0. 路径准备 tif_path = Path(cfg['tif_path']) base_name = tif_path.stem out_dir = tif_path.parent / f"{base_name}_1" out_dir.mkdir(exist_ok=True) # 1. 加载模型 & 配置 with open(cfg['config_path'], 'r', encoding='utf-8') as f: config = json.load(f) model = load_model(cfg['model_path'], config, device) # 2. 打开大图 with rasterio.open(tif_path) as src: h, w, bands = src.height, src.width, src.count transform, crs = src.transform, src.crs print(f"图像尺寸: {w}×{h} 波段: {bands}") # 3. 计算瓦片信息 tile_size, overlap = cfg['tile_size'], cfg['overlap'] stride = tile_size - overlap tiles_x = (w - tile_size) // stride + 1 + ((w - tile_size) % stride != 0) tiles_y = (h - tile_size) // stride + 1 + ((h - tile_size) % stride != 0) total = tiles_x * tiles_y print(f"瓦片数量: {total} ({tiles_x}×{tiles_y})") # 4. 创建输出内存文件 pred_driver = gdal.GetDriverByName('MEM') prob_driver = gdal.GetDriverByName('MEM') pred_ds = pred_driver.Create('', w, h, 1, gdal.GDT_Byte) prob_ds = prob_driver.Create('', w, h, 1, gdal.GDT_Float32) for ds in [pred_ds, prob_ds]: ds.SetGeoTransform(transform.to_gdal()) ds.SetProjection(crs.to_wkt()) pred_band = pred_ds.GetRasterBand(1) prob_band = prob_ds.GetRasterBand(1) count = np.zeros((h, w), dtype=np.uint16) # 5. 分块预测 tile_id = 0 for y in range(0, h - tile_size + 1, stride): for x in range(0, w - tile_size + 1, stride): actual_x, actual_y = min(x, w - tile_size), min(y, h - tile_size) window = Window(actual_x, actual_y, tile_size, tile_size) tile = src.read(window=window) # CHW tile = np.transpose(tile, (1, 2, 0)) # HWC if tile.shape[2] >= 4: tile = tile[:, :, :4] else: tile = tile[:, :, [2, 1, 0]] if tile.shape[2] >= 3 else tile[:, :, :3] tensor = preprocess_tile(tile.astype(np.float32), config) mask, prob = predict_tile(model, tensor, device, config, cfg['threshold']) # 写入内存 pred_band.WriteArray(mask, actual_x, actual_y) prob_band.WriteArray(prob, actual_x, actual_y) count[actual_y:actual_y+tile_size, actual_x:actual_x+tile_size] += 1 tile_id += 1 if tile_id % 500 == 0: print(f" 已预测 {tile_id}/{total}") # 6. 平均化重叠区域 count[count == 0] = 1 pred_final = (pred_band.ReadAsArray().astype(np.float32) / count).round().astype(np.uint8) prob_final = prob_band.ReadAsArray() / count # 7. 保存 GeoTIFF gtiff_path = out_dir / f"{base_name}_prediction.tif" drv = gdal.GetDriverByName('GTiff') ds_out = drv.Create(str(gtiff_path), w, h, 2, gdal.GDT_Byte, options=['COMPRESS=DEFLATE', 'TILED=YES']) ds_out.SetGeoTransform(transform.to_gdal()) ds_out.SetProjection(crs.to_wkt()) # band1: 二值掩膜 b1 = ds_out.GetRasterBand(1) b1.WriteArray(pred_final * 255) b1.SetNoDataValue(0) # band2: 概率 [0-255] b2 = ds_out.GetRasterBand(2) b2.WriteArray((prob_final * 255).astype(np.uint8)) ds_out = None print(f"✅ GeoTIFF 已保存: {gtiff_path}") # 8. 矢量化 shp_path = out_dir / f"{base_name}_sargassum.shp" vectorize(gtiff_path, shp_path, cfg['sieve_pixels'], cfg['min_area']) print(f"✅ 矢量文件已保存: {shp_path}") print("🎉 全部完成！") # ---------- 矢量化 ---------- def vectorize(tif_path: Path, shp_path: Path, sieve_pixels: int, min_area_m2: float): ds = gdal.Open(str(tif_path), gdal.GA_ReadOnly) band = ds.GetRasterBand(1) data = band.ReadAsArray() geo = ds.GetGeoTransform() proj = ds.GetProjection() # 碎斑过滤 if sieve_pixels > 0: kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (sieve_pixels//2, sieve_pixels//2)) data = cv2.morphologyEx(data, cv2.MORPH_OPEN, kernel) mask = (data == 255).astype(np.uint8) contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) driver = ogr.GetDriverByName('ESRI Shapefile') if shp_path.exists(): driver.DeleteDataSource(str(shp_path)) out_ds = driver.CreateDataSource(str(shp_path)) srs = osr.SpatialReference() srs.ImportFromWkt(proj) layer = out_ds.CreateLayer("sargassum", srs, ogr.wkbPolygon) fd = ogr.FieldDefn('DN', ogr.OFTInteger) layer.CreateField(fd) pixel_area = abs(geo[1] * geo[5]) for cnt in contours: area_pixel = cv2.contourArea(cnt) if area_pixel * pixel_area < min_area_m2: continue ring = ogr.Geometry(ogr.wkbLinearRing) for pt in cnt[:, 0, :]: x = geo[0] + pt[0] * geo[1] + pt[1] * geo[2] y = geo[3] + pt[0] * geo[4] + pt[1] * geo[5] ring.AddPoint(x, y) ring.CloseRings() poly = ogr.Geometry(ogr.wkbPolygon) poly.AddGeometry(ring) feat = ogr.Feature(layer.GetLayerDefn()) feat.SetGeometry(poly) feat.SetField('DN', 255) layer.CreateFeature(feat) out_ds = None ds = None # ---------- 入口 ---------- if __name__ == "__main__": # 1. 读 XML（唯一需要改的地方） xml_file = r"test.xml" # 也可 sys.argv[1] 传入 tif_path, out_root = parse_xml(xml_file) # 2. 组装配置字典（其余逻辑零改动） base_name = Path(tif_path).stem out_dir = Path(out_root) / f"{base_name}_1" # F:/Results/admin/GF1/157/xxx_1 out_dir.mkdir(parents=True, exist_ok=True) MAIN_CFG = { "tif_path" : str(tif_path), "model_path" : r"seaweed_segmentation_improved_epoch500\best_checkpoint.pth", "config_path" : r"seaweed_segmentation_improved_epoch500\config.json", "tile_size" : 256, "overlap" : 64, "threshold" : 0.5, "sieve_pixels" : 10, "min_area" : 0.0, } # 3. 跑流程 run_one_shot(MAIN_CFG)