predict_large_image_tiles_256.py

"""
大图浒苔分割预测脚本 - 256x256瓦片版本
读取大图遥感图像，切割为256x256小patch，对每个patch进行预测
预测结果格式与predict_seaweed_256x256_simple.py保持一致
"""
import os
import torch
import numpy as np
from pathlib import Path
import json
import time
from datetime import datetime
import cv2
from tqdm import tqdm
import rasterio
from rasterio.windows import Window
import warnings
warnings.filterwarnings('ignore')

# 导入自定义模块
from dinov3_deeplabv3plus import DinoV3DeepLabV3Plus

def load_model(model_path, config, device):
    """加载训练好的模型 - 与predict_seaweed_256x256_simple.py保持一致"""
    print(f"正在加载模型: {model_path}")
    
    # 创建模型
    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)
    
    # 加载权重
    checkpoint = torch.load(model_path, map_location=device)
    
    # 处理不同的checkpoint格式
    if 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
        print(f"加载模型权重 (epoch {checkpoint.get('epoch', 'unknown')})")
    elif 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['state_dict'])
        print("加载模型权重")
    else:
        # 直接加载state dict
        model.load_state_dict(checkpoint)
        print("加载模型权重")
    
    model.eval()
    print("模型加载完成，进入评估模式")
    return model

def process_image_for_model(image_path, target_size=256, use_4channel=True):
    """处理图像用于模型预测 - 完全复用predict_seaweed_256x256_simple.py的逻辑"""
    # 使用OpenCV读取16位TIF图像
    img_16bit = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    
    if img_16bit is None:
        raise ValueError(f"无法读取图像文件: {image_path}")
    
    # 使用rasterio方式读取多波段数据
    try:
        import rasterio
        with rasterio.open(image_path) as src:
            # 读取所有波段
            image = src.read()
            # 转换为HWC格式
            image = np.transpose(image, (1, 2, 0))
    except Exception as e:
        print(f"rasterio读取失败，使用OpenCV: {e}")
        # 回退到OpenCV
        image = img_16bit
    
    # 处理通道数
    if use_4channel and image.shape[2] >= 4:
        # 使用4通道
        processed_image = image[:, :, :4]
    else:
        # 使用3通道假彩色（432波段）
        if image.shape[2] >= 4:
            # 选择第4,3,2波段（索引3,2,1）
            processed_image = image[:, :, [3, 2, 1]]
        else:
            # 如果通道数不足，使用可用通道
            processed_image = image[:, :, :3]
            if image.shape[2] < 3:
                # 如果还是不足3通道，重复填充
                while processed_image.shape[2] < 3:
                    processed_image = np.concatenate([processed_image, processed_image[:, :, -1:]], axis=2)
    
    # 调整尺寸
    if processed_image.shape[0] != target_size or processed_image.shape[1] != target_size:
        processed_image = cv2.resize(processed_image, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
    
    # 转换为float32并标准化到0-1范围
    processed_image = processed_image.astype(np.float32)
    if processed_image.max() > 1.0:
        processed_image = processed_image / 65535.0  # 16位最大值
    
    # 标准化
    if use_4channel:
        # 4通道标准化
        mean = np.array([0.430, 0.411, 0.296, 0.350])
        std = np.array([0.213, 0.156, 0.143, 0.180])
    else:
        # 3通道标准化
        mean = np.array([0.430, 0.411, 0.296])
        std = np.array([0.213, 0.156, 0.143])
    
    # 应用标准化
    for i in range(processed_image.shape[2]):
        processed_image[:, :, i] = (processed_image[:, :, i] - mean[i]) / std[i]
    
    # 转换为tensor并添加batch维度
    img_tensor = torch.from_numpy(processed_image).permute(2, 0, 1).float()
    
    return img_tensor.unsqueeze(0)  # 添加batch维度

def predict_image(model, image_tensor, device, config, threshold=0.5):
    """
    预测单张图像 - 与predict_seaweed_256x256_simple.py保持一致
    
    Args:
        model: 训练好的模型
        image_tensor: 图像tensor
        device: 设备
        config: 配置字典
        threshold: 浒苔预测阈值，只有当浒苔概率超过此阈值时才预测为浒苔（默认0.5）
    """
    # 预测
    with torch.no_grad():
        start_time = time.time()
        output = model(image_tensor)
        inference_time = time.time() - start_time
        
        # 处理输出格式
        if isinstance(output, dict):
            output = output['out']
        
        # 获取预测结果
        pred = torch.softmax(output, dim=1)
        
        # 计算浒苔概率
        seaweed_prob = pred[0, 1].squeeze(0).cpu().numpy()  # 类别1是浒苔
        
        # 使用阈值进行预测：只有当浒苔概率超过阈值时才预测为浒苔
        # 这样可以避免在无浒苔图片中，由于两个类别概率接近而误判
        pred_class = (seaweed_prob > threshold).astype(np.uint8)
        
        return {
            'prediction': pred_class,
            'seaweed_probability': seaweed_prob,
            'inference_time': inference_time
        }

def save_prediction_results(prediction, seaweed_prob, output_path, base_name):
    """保存预测结果到文件 - 与predict_seaweed_256x256_simple.py保持一致"""
    # 保存预测掩码 (0-1)
    pred_mask = prediction.astype(np.uint8)
    np.save(str(output_path / f"{base_name}_prediction.npy"), pred_mask)
    
    # 保存概率图 (0-1)
    np.save(str(output_path / f"{base_name}_probability.npy"), seaweed_prob)
    
    # 保存为PNG图像
    pred_png = (prediction * 255).astype(np.uint8)
    prob_png = (seaweed_prob * 255).astype(np.uint8)
    
    cv2.imwrite(str(output_path / f"{base_name}_prediction.png"), pred_png)
    cv2.imwrite(str(output_path / f"{base_name}_probability.png"), prob_png)

def calculate_statistics(prediction, seaweed_prob):
    """计算统计信息 - 与predict_seaweed_256x256_simple.py保持一致"""
    total_pixels = prediction.size
    seaweed_pixels = np.sum(prediction == 1)
    seaweed_ratio = seaweed_pixels / total_pixels
    
    # 计算平均概率
    avg_probability = np.mean(seaweed_prob)
    max_probability = np.max(seaweed_prob)
    
    return {
        'total_pixels': int(total_pixels),
        'seaweed_pixels': int(seaweed_pixels),
        'seaweed_ratio': float(seaweed_ratio),
        'avg_probability': float(avg_probability),
        'max_probability': float(max_probability)
    }

def get_tile_info_from_large_image(image_path, tile_size=256, overlap=0):
    """
    获取大图的瓦片信息，但不实际读取数据（内存友好）
    
    Args:
        image_path: 大图路径
        tile_size: 瓦片大小
        overlap: 重叠大小
    
    Returns:
        list: 瓦片信息列表，每个元素包含(x, y, window, tile_id)
    """
    tile_info_list = []
    
    with rasterio.open(image_path) as src:
        height, width = src.height, src.width
        bands = src.count
        
        print(f"大图信息: {width}x{height}, {bands}波段")
        
        # 计算步长
        stride = tile_size - overlap
        
        # 计算瓦片数量
        tiles_y = (height - tile_size) // stride + 1
        tiles_x = (width - tile_size) // stride + 1
        
        # 处理边界情况
        if (height - tile_size) % stride != 0:
            tiles_y += 1
        if (width - tile_size) % stride != 0:
            tiles_x += 1
        
        total_tiles = tiles_x * tiles_y
        print(f"将提取 {total_tiles} 个瓦片 ({tiles_x} x {tiles_y})")
        
        # 生成瓦片信息
        tile_count = 0
        for y in range(0, height - tile_size + 1, stride):
            for x in range(0, width - tile_size + 1, stride):
                
                # 处理边界情况
                actual_x = min(x, width - tile_size)
                actual_y = min(y, height - tile_size)
                
                # 创建窗口信息
                window = Window(actual_x, actual_y, tile_size, tile_size)
                
                tile_info_list.append({
                    'x': actual_x,
                    'y': actual_y,
                    'window': window,
                    'tile_id': tile_count
                })
                
                tile_count += 1
    
    return tile_info_list

def read_tile_data(image_path, window, bands):
    """读取单个瓦片的数据"""
    with rasterio.open(image_path) as src:
        # 读取多波段数据
        if bands >= 4:
            # 使用4通道
            tile_data = src.read([1, 2, 3, 4], window=window)
        else:
            # 使用可用通道
            tile_data = src.read(list(range(1, min(bands, 4) + 1)), window=window)
        
        # 转换为HWC格式
        tile_data = np.transpose(tile_data, (1, 2, 0))
        
        return tile_data

def save_tile_temporarily(tile_data, temp_dir, tile_id):
    """临时保存瓦片为文件，用于process_image_for_model函数"""
    temp_path = os.path.join(temp_dir, f"temp_tile_{tile_id}.tif")
    
    # 确保临时目录存在
    os.makedirs(temp_dir, exist_ok=True)
    
    # 保存为临时TIFF文件
    with rasterio.open(
        temp_path, 'w',
        driver='GTiff',
        height=tile_data.shape[0],
        width=tile_data.shape[1],
        count=tile_data.shape[2],
        dtype=tile_data.dtype,
        compress='lzw'
    ) as dst:
        for i in range(tile_data.shape[2]):
            dst.write(tile_data[:, :, i], i + 1)
    
    return temp_path

def predict_tiles(tiles, model, device, config, output_dir, threshold=0.5):
    """
    对所有瓦片进行预测
    
    Args:
        tiles: 瓦片列表
        model: 模型
        device: 设备
        config: 配置
        output_dir: 输出目录
        threshold: 预测阈值
    """
    results = []
    total_inference_time = 0
    
    # 创建临时目录
    temp_dir = os.path.join(output_dir, "temp_tiles")
    
    print(f"\n开始预测 {len(tiles)} 个瓦片...")
    
    # 使用tqdm显示进度条
    for tile_info in tqdm(tiles, desc="预测瓦片"):
        try:
            tile_id = tile_info['tile_id']
            tile_data = tile_info['data']
            
            # 临时保存瓦片为文件
            temp_tile_path = save_tile_temporarily(tile_data, temp_dir, tile_id)
            
            # 处理图像用于模型
            image_tensor = process_image_for_model(
                temp_tile_path,
                target_size=config['image_size'],
                use_4channel=config['use_4channel']
            ).to(device)
            
            # 预测（使用阈值）
            result = predict_image(model, image_tensor, device, config, threshold=threshold)
            
            # 计算统计信息
            stats = calculate_statistics(result['prediction'], result['seaweed_probability'])
            
            # 保存预测结果
            base_name = f"patch{tile_id:05d}"
            save_prediction_results(
                result['prediction'], 
                result['seaweed_probability'], 
                Path(output_dir), 
                base_name
            )
            
            # 保存结果信息
            result_info = {
                'tile_id': tile_id,
                'x': tile_info['x'],
                'y': tile_info['y'],
                'inference_time': float(result['inference_time']),
                'statistics': stats
            }
            results.append(result_info)
            
            total_inference_time += result['inference_time']
            
            # 清理临时文件
            if os.path.exists(temp_tile_path):
                os.remove(temp_tile_path)
                
        except Exception as e:
            print(f"  - 瓦片 {tile_id} 预测失败: {str(e)}")
            continue
    
    # 清理临时目录
    if os.path.exists(temp_dir):
        try:
            os.rmdir(temp_dir)
        except:
            pass
    
    return results, total_inference_time

def predict_tiles_memory_friendly(tile_info_list, image_path, model, device, config, output_dir, threshold=0.5):
    """
    内存友好的瓦片预测 - 逐块处理
    
    Args:
        tile_info_list: 瓦片信息列表
        image_path: 图像路径
        model: 模型
        device: 设备
        config: 配置
        output_dir: 输出目录
        threshold: 预测阈值
    """
    results = []
    total_inference_time = 0
    
    # 创建临时目录
    temp_dir = os.path.join(output_dir, "temp_tiles")
    os.makedirs(temp_dir, exist_ok=True)
    
    # 获取图像信息
    with rasterio.open(image_path) as src:
        bands = src.count
    
    print(f"\n开始预测 {len(tile_info_list)} 个瓦片...")
    
    # 使用tqdm显示进度条
    for tile_info in tqdm(tile_info_list, desc="预测瓦片"):
        try:
            tile_id = tile_info['tile_id']
            
            # 读取瓦片数据
            tile_data = read_tile_data(image_path, tile_info['window'], bands)
            
            # 临时保存瓦片为文件
            temp_tile_path = save_tile_temporarily(tile_data, temp_dir, tile_id)
            
            # 处理图像用于模型
            image_tensor = process_image_for_model(
                temp_tile_path,
                target_size=config['image_size'],
                use_4channel=config['use_4channel']
            ).to(device)
            
            # 预测（使用阈值）
            result = predict_image(model, image_tensor, device, config, threshold=threshold)
            
            # 计算统计信息
            stats = calculate_statistics(result['prediction'], result['seaweed_probability'])
            
            # 保存预测结果
            base_name = f"patch{tile_id:05d}"
            save_prediction_results(
                result['prediction'], 
                result['seaweed_probability'], 
                Path(output_dir), 
                base_name
            )
            
            # 保存结果信息
            result_info = {
                'tile_id': tile_id,
                'x': tile_info['x'],
                'y': tile_info['y'],
                'inference_time': float(result['inference_time']),
                'statistics': stats
            }
            results.append(result_info)
            
            total_inference_time += result['inference_time']
            
            # 清理临时文件
            if os.path.exists(temp_tile_path):
                os.remove(temp_tile_path)
                
        except Exception as e:
            print(f"  - 瓦片 {tile_id} 预测失败: {str(e)}")
            continue
    
    # 清理临时目录
    try:
        os.rmdir(temp_dir)
    except:
        pass
    
    return results, total_inference_time

def main():
    """主函数"""
    print("=" * 60)
    print("大图瓦片浒苔分割预测系统 - 256x256版本")
    print("=" * 60)
    
    # 配置
    model_path = "seaweed_segmentation_improved_epoch500/best_checkpoint.pth"
    config_path = "seaweed_segmentation_improved_epoch500/config.json"
    large_image_dir = "E:\GF6\山科影像\浒苔\GF6_PMS_E121.1_N33.6_20250604_L1A1420584616-MUX_fuse.tif"
    output_dir = "outputs/test_large"
    
    # 瓦片配置
    tile_size = 256
    overlap = 0  # 无重叠
    
    # 预测阈值：只有当浒苔概率超过此阈值时才预测为浒苔
    # 建议值：0.5-0.7，可以根据验证集调整
    prediction_threshold = 0.5
    
    # 创建设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 加载配置
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    
    print(f"模型配置:")
    print(f"  - 输入通道: {4 if config['use_4channel'] else 3}")
    print(f"  - 图像尺寸: {config['image_size']}x{config['image_size']}")
    print(f"  - 类别数: {config['num_classes']}")
    print(f"  - 瓦片大小: {tile_size}x{tile_size}")
    print(f"  - 重叠: {overlap}")
    print(f"  - 预测阈值: {prediction_threshold}")
    
    # 加载模型
    model = load_model(model_path, config, device)
    
    # 创建输出目录
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 获取大图文件
    image_extensions = ['.tif', '.tiff']
    large_images = []
    for ext in image_extensions:
        large_images.extend(Path(large_image_dir).glob(f"*{ext}"))
    
    print(f"\n找到 {len(large_images)} 个大图文件")
    
    # 处理每个大图
    all_results = []
    
    for image_file in large_images:
        print(f"\n处理大图: {image_file.name}")
        
        try:
            # 获取瓦片信息（不读取实际数据）
            tile_info_list = get_tile_info_from_large_image(
                str(image_file), 
                tile_size=tile_size, 
                overlap=overlap
            )
            
            if not tile_info_list:
                print(f"  - 未能生成瓦片信息，跳过")
                continue
            
            # 内存友好的瓦片预测
            results, total_inference_time = predict_tiles_memory_friendly(
                tile_info_list, str(image_file), model, device, config, output_path, threshold=prediction_threshold
            )
            
            if results:
                # 计算平均推理时间
                avg_inference_time = total_inference_time / len(results)
                
                # 计算平均统计
                all_ratios = [r['statistics']['seaweed_ratio'] for r in results]
                avg_seaweed_ratio = np.mean(all_ratios)
                
                print(f"\n瓦片预测完成:")
                print(f"  - 总瓦片数: {len(results)}")
                print(f"  - 平均推理时间: {avg_inference_time:.3f}s")
                print(f"  - 平均浒苔比例: {avg_seaweed_ratio:.2%}")
                
                # 保存大图处理结果
                image_results = {
                    'image_file': str(image_file),
                    'total_tiles': len(results),
                    'avg_inference_time': avg_inference_time,
                    'avg_seaweed_ratio': float(avg_seaweed_ratio),
                    'tile_results': results
                }
                all_results.append(image_results)
                
                # 保存结果文件
                result_file = output_path / f"{image_file.stem}_tile_results.json"
                with open(result_file, 'w', encoding='utf-8') as f:
                    json.dump(image_results, f, indent=2, ensure_ascii=False, default=str)
                
                print(f"  - 结果保存: {result_file}")
            
        except Exception as e:
            print(f"  - 处理失败: {str(e)}")
            import traceback
            traceback.print_exc()
            continue
    
    # 保存总体结果
    if all_results:
        summary = {
            'timestamp': datetime.now().isoformat(),
            'total_images': len(all_results),
            'total_tiles': sum([r['total_tiles'] for r in all_results]),
            'results': all_results
        }
        
        summary_file = output_path / "large_image_tile_summary.json"
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False, default=str)
        
        print(f"\n" + "=" * 60)
        print("大图瓦片预测完成总结")
        print("=" * 60)
        print(f"总图像数: {len(all_results)}")
        print(f"总瓦片数: {summary['total_tiles']}")
        print(f"结果保存目录: {output_path}")
        print(f"总结文件: {summary_file}")
        
        # 显示浒苔分布统计
        all_ratios = []
        for result in all_results:
            all_ratios.extend([r['statistics']['seaweed_ratio'] for r in result['tile_results']])
        
        if all_ratios:
            print(f"\n浒苔分布统计:")
            print(f"  - 最小浒苔比例: {min(all_ratios):.2%}")
            print(f"  - 最大浒苔比例: {max(all_ratios):.2%}")
            print(f"  - 浒苔比例标准差: {np.std(all_ratios):.2%}")
    
    print("\n🎉 大图瓦片浒苔预测完成！")
    print(f"预测结果保存在: {output_dir}")

if __name__ == "__main__":
    main()