upsample_deep_features.py

import os
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageFilter
from scipy import ndimage

def upsample_feature_maps(input_dir, output_dir, target_size=(98, 98)):
    """
    将深层特征图上采样到指定尺寸
    
    Args:
        input_dir: 输入特征图目录（24×24的深层特征图）
        output_dir: 输出目录
        target_size: 目标尺寸，默认(98, 98)
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取所有PNG文件
    png_files = [f for f in os.listdir(input_dir) if f.endswith('.png')]
    
    for png_file in png_files:
        input_path = os.path.join(input_dir, png_file)
        output_path = os.path.join(output_dir, f"upsampled_{png_file}")
        
        # 读取图像 - 保持原始颜色
        img = cv2.imread(input_path, cv2.IMREAD_COLOR)  # 改为 IMREAD_COLOR
        
        if img is None:
            print(f"无法读取图像: {input_path}")
            continue
            
        print(f"原始尺寸: {img.shape}, 目标尺寸: {target_size}")
        
        # 方法1: 使用OpenCV双三次插值上采样
        upsampled_cv = cv2.resize(img, target_size, interpolation=cv2.INTER_CUBIC)
        
        # 方法2: 使用PyTorch双线性插值上采样
        img_tensor = torch.from_numpy(img).float().unsqueeze(0).unsqueeze(0)
        upsampled_torch = F.interpolate(img_tensor, size=target_size, mode='bilinear', align_corners=False)
        upsampled_torch = upsampled_torch.squeeze().numpy().astype(np.uint8)
        
        # 保存上采样结果
        cv2.imwrite(output_path.replace('.png', '_cubic.png'), upsampled_cv)
        cv2.imwrite(output_path.replace('.png', '_bilinear.png'), upsampled_torch)
        
        print(f"已保存: {output_path}")

def create_comparison_grid(shallow_dir, deep_dir, upsampled_dir, output_path, num_channels=5):
    """
    创建浅层、深层和上采样深层特征图的对比网格
    
    Args:
        shallow_dir: 浅层特征图目录（98×98）
        deep_dir: 深层特征图目录（24×24）
        upsampled_dir: 上采样深层特征图目录（98×98）
        output_path: 对比图保存路径
        num_channels: 要对比的通道数
    """
    fig, axes = plt.subplots(3, num_channels, figsize=(num_channels*4, 12))
    
    for i in range(num_channels):
        channel_file = f"channel_{i+1:03d}.png"
        
        # 读取浅层特征图
        shallow_path = os.path.join(shallow_dir, channel_file)
        if os.path.exists(shallow_path):
            shallow_img = cv2.imread(shallow_path, cv2.IMREAD_GRAYSCALE)
            axes[0, i].imshow(shallow_img, cmap='RdYlBu_r')
            axes[0, i].set_title(f'浅层 Ch{i+1} (98×98)', fontsize=12)
            axes[0, i].axis('off')
        
        # 读取深层特征图
        deep_path = os.path.join(deep_dir, channel_file)
        if os.path.exists(deep_path):
            deep_img = cv2.imread(deep_path, cv2.IMREAD_GRAYSCALE)
            axes[1, i].imshow(deep_img, cmap='RdYlBu_r')
            axes[1, i].set_title(f'深层 Ch{i+1} (24×24)', fontsize=12)
            axes[1, i].axis('off')
        
        # 读取上采样深层特征图
        upsampled_path = os.path.join(upsampled_dir, f"upsampled_channel_{i+1:03d}_bilinear.png")
        if os.path.exists(upsampled_path):
            upsampled_img = cv2.imread(upsampled_path, cv2.IMREAD_GRAYSCALE)
            axes[2, i].imshow(upsampled_img, cmap='RdYlBu_r')
            axes[2, i].set_title(f'上采样深层 Ch{i+1} (98×98)', fontsize=12)
            axes[2, i].axis('off')
    
    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"对比图已保存: {output_path}")

def batch_upsample_layers(base_dir, image_name, target_size=(98, 98)):
    """
    批量处理指定图像的所有深层特征图
    
    Args:
        base_dir: feature_maps_output基础目录
        image_name: 图像名称（如'ROIs2017_winter_108_p282.tif'）
        target_size: 目标尺寸
    """
    image_dir = os.path.join(base_dir, image_name)
    
    if not os.path.exists(image_dir):
        print(f"图像目录不存在: {image_dir}")
        return
    
    # 需要上采样的深层特征图层
    deep_layers = ['spatial_layer3', 'sar_arconv3', 'freq_sdtb3']
    
    for layer in deep_layers:
        layer_dir = os.path.join(image_dir, layer)
        if os.path.exists(layer_dir):
            output_dir = os.path.join(image_dir, f"{layer}_upsampled")
            print(f"正在处理层: {layer}")
            upsample_feature_maps(layer_dir, output_dir, target_size)
        else:
            print(f"层目录不存在: {layer_dir}")
    
    # 创建对比图
    shallow_dir = os.path.join(image_dir, 'spatial_layer1')  # 98×98
    deep_dir = os.path.join(image_dir, 'spatial_layer3')     # 24×24
    upsampled_dir = os.path.join(image_dir, 'spatial_layer3_upsampled')  # 98×98
    
    if all(os.path.exists(d) for d in [shallow_dir, deep_dir, upsampled_dir]):
        comparison_path = os.path.join(image_dir, 'layer_comparison.png')
        create_comparison_grid(shallow_dir, deep_dir, upsampled_dir, comparison_path)

def upsample_single_image(input_path, output_dir, target_size=(98, 98), method='lanczos'):
    """
    使用多种高质量插值方法上采样图像
    
    Args:
        input_path: 输入图像路径
        output_dir: 输出目录
        target_size: 目标尺寸
        method: 插值方法 ('lanczos', 'cubic', 'bilinear', 'nearest', 'area', 'bicubic_pil', 'antialias')
    """
    if not os.path.exists(input_path):
        print(f"输入文件不存在: {input_path}")
        return None
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 读取图像 - 保持原始颜色
    img = cv2.imread(input_path, cv2.IMREAD_COLOR)
    
    if img is None:
        print(f"无法读取图像: {input_path}")
        return None
    
    print(f"原始尺寸: {img.shape}, 目标尺寸: {target_size}")
    
    # 生成输出文件名
    input_filename = os.path.basename(input_path)
    base_name = os.path.splitext(input_filename)[0]
    output_filename = f"{base_name}_upsampled_{target_size[0]}x{target_size[1]}_{method}.png"
    output_path = os.path.join(output_dir, output_filename)
    
    # 根据方法选择上采样
    if method == 'lanczos':
        # Lanczos插值 - 高质量，保持锐度
        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        upsampled_pil = img_pil.resize(target_size, Image.LANCZOS)
        upsampled = cv2.cvtColor(np.array(upsampled_pil), cv2.COLOR_RGB2BGR)
        
    elif method == 'bicubic_pil':
        # PIL的双三次插值 - 比OpenCV更平滑
        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        upsampled_pil = img_pil.resize(target_size, Image.BICUBIC)
        upsampled = cv2.cvtColor(np.array(upsampled_pil), cv2.COLOR_RGB2BGR)
        
    elif method == 'antialias':
        # 抗锯齿插值 - 减少锯齿效应
        img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        upsampled_pil = img_pil.resize(target_size, Image.Resampling.LANCZOS)
        # 应用轻微的抗锯齿滤波
        upsampled_pil = upsampled_pil.filter(ImageFilter.SMOOTH_MORE)
        upsampled = cv2.cvtColor(np.array(upsampled_pil), cv2.COLOR_RGB2BGR)
        
    elif method == 'area':
        # 区域插值 - 适合缩小，但也可用于放大
        upsampled = cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
        
    elif method == 'cubic':
        # OpenCV双三次插值
        upsampled = cv2.resize(img, target_size, interpolation=cv2.INTER_CUBIC)
        
    elif method == 'nearest':
        # 最近邻插值 - 保持像素值不变
        upsampled = cv2.resize(img, target_size, interpolation=cv2.INTER_NEAREST)
        
    elif method == 'bilinear':
        # 双线性插值
        if len(img.shape) == 3:  # 彩色图像
            img_tensor = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0)
            upsampled_tensor = F.interpolate(img_tensor, size=target_size, mode='bilinear', align_corners=False)
            upsampled = upsampled_tensor.squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8)
        else:  # 灰度图像
            img_tensor = torch.from_numpy(img).float().unsqueeze(0).unsqueeze(0)
            upsampled_tensor = F.interpolate(img_tensor, size=target_size, mode='bilinear', align_corners=False)
            upsampled = upsampled_tensor.squeeze().numpy().astype(np.uint8)
            
    elif method == 'edsr_like':
        # 模拟EDSR风格的上采样（简化版）
        # 先用双三次插值，然后应用锐化
        upsampled = cv2.resize(img, target_size, interpolation=cv2.INTER_CUBIC)
        # 应用锐化核
        kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
        upsampled = cv2.filter2D(upsampled, -1, kernel)
        upsampled = np.clip(upsampled, 0, 255).astype(np.uint8)
        
    else:
        print(f"不支持的方法: {method}")
        return None
    
    # 保存结果
    cv2.imwrite(output_path, upsampled)
    print(f"上采样完成，已保存: {output_path}")
    
    return output_path

def batch_test_methods(input_path, output_dir, target_size=(98, 98)):
    """
    测试所有插值方法，生成对比图
    """
    methods = ['lanczos', 'bicubic_pil', 'antialias', 'cubic', 'bilinear', 'nearest', 'area', 'edsr_like']
    
    results = []
    for method in methods:
        result = upsample_single_image(input_path, output_dir, target_size, method)
        if result:
            results.append((method, result))
    
    print(f"\n完成所有方法测试，共生成 {len(results)} 个结果:")
    for method, path in results:
        print(f"- {method}: {path}")
    
    return results

if __name__ == "__main__":
    # 配置参数
    input_image_path = "feature_maps_output/ROIs2017_winter_108_p147.tif/align_fusion/channel_202.png"
    output_directory = "upsampled_results"
    target_size = (98, 98)
    
    # 方法1: 测试单个方法（推荐Lanczos获得最佳质量）
    method = 'lanczos'  # 推荐: 'lanczos', 'bicubic_pil', 'antialias'
    result = upsample_single_image(input_image_path, output_directory, target_size, method)
    
    # 方法2: 测试所有方法进行对比
    # batch_test_methods(input_image_path, output_directory, target_size)
    
    if result:
        print("处理完成！")
    else:
        print("处理失败！")