vector.py

# -*- 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)