example_usage.py

#!/usr/bin/env python3
"""
Grid Sampler CUDA 使用示例
"""

import numpy as np
import grid_sampler_cuda
import time

def example_basic_usage():
    """基本使用示例"""
    print("=== 基本使用示例 ===")

    # 创建一个简单的图像 [1, 3, 4, 4] (批次=1, 通道=3, 高=4, 宽=4)
    input_image = np.array([
        # 红色通道
        [[[1, 2, 3, 4],
          [5, 6, 7, 8],
          [9, 10, 11, 12],
          [13, 14, 15, 16]]],
        # 绿色通道
        [[[17, 18, 19, 20],
          [21, 22, 23, 24],
          [25, 26, 27, 28],
          [29, 30, 31, 32]]],
        # 蓝色通道
        [[[33, 34, 35, 36],
          [37, 38, 39, 40],
          [41, 42, 43, 44],
          [45, 46, 47, 48]]]
    ], dtype=np.float32).transpose(1, 0, 2, 3)  # 转换为 [N, C, H, W]

    print(f"输入图像形状: {input_image.shape}")
    print("输入图像 (红色通道):")
    print(input_image[0, 0])

    # 创建网格 - 将4x4图像采样到2x2
    # 网格坐标范围 [-1, 1]
    grid = np.array([
        [[[-1, -1], [1, -1]],   # 左上角到右上角
         [[-1, 1], [1, 1]]]     # 左下角到右下角
    ], dtype=np.float32)

    print(f"网格形状: {grid.shape}")
    print("网格坐标:")
    print(grid[0])

    # 执行grid sampling
    output = grid_sampler_cuda.grid_sampler(
        input_image, grid,
        mode="bilinear",
        padding_mode="zeros",
        align_corners=False
    )

    print(f"输出形状: {output.shape}")
    print("输出图像 (红色通道):")
    print(output[0, 0])
    print()

def example_image_rotation():
    """图像旋转示例"""
    print("=== 图像旋转示例 ===")

    # 创建一个简单的图像
    height, width = 5, 5
    input_image = np.zeros((1, 1, height, width), dtype=np.float32)

    # 在中心画一个十字
    center_h, center_w = height // 2, width // 2
    input_image[0, 0, center_h, :] = 1.0  # 水平线
    input_image[0, 0, :, center_w] = 1.0  # 垂直线

    print("原始图像:")
    print(input_image[0, 0])

    # 创建旋转45度的网格
    angle = np.pi / 4  # 45度
    cos_a, sin_a = np.cos(angle), np.sin(angle)

    # 创建输出网格
    output_h, output_w = 5, 5
    grid = np.zeros((1, output_h, output_w, 2), dtype=np.float32)

    for h in range(output_h):
        for w in range(output_w):
            # 将输出坐标转换为输入坐标
            x = (w - output_w // 2) / (output_w // 2)
            y = (h - output_h // 2) / (output_h // 2)

            # 应用旋转变换
            x_rot = x * cos_a - y * sin_a
            y_rot = x * sin_a + y * cos_a

            grid[0, h, w, 0] = x_rot
            grid[0, h, w, 1] = y_rot

    # 执行旋转
    rotated = grid_sampler_cuda.grid_sampler(
        input_image, grid,
        mode="bilinear",
        padding_mode="zeros"
    )

    print("旋转45度后的图像:")
    print(rotated[0, 0])
    print()

def example_performance_comparison():
    """性能对比示例"""
    print("=== 性能对比示例 ===")

    # 创建较大的测试数据
    batch_size = 2
    channels = 64
    height = 256
    width = 256
    output_height = 128
    output_width = 128

    print(f"测试数据: [{batch_size}, {channels}, {height}, {width}] -> [{batch_size}, {channels}, {output_height}, {output_width}]")

    # 创建输入数据
    input_data = np.random.randn(batch_size, channels, height, width).astype(np.float32)

    # 创建网格数据
    grid_data = np.random.randn(batch_size, output_height, output_width, 2).astype(np.float32)
    grid_data = np.clip(grid_data, -1, 1)

    # 测试不同模式
    modes = ["bilinear", "nearest"]
    padding_modes = ["zeros", "border", "reflection"]

    for mode in modes:
        for padding_mode in padding_modes:
            # 预热
            _ = grid_sampler_cuda.grid_sampler(input_data, grid_data, mode=mode, padding_mode=padding_mode)

            # 性能测试
            num_iterations = 5
            start_time = time.time()

            for _ in range(num_iterations):
                output = grid_sampler_cuda.grid_sampler(input_data, grid_data, mode=mode, padding_mode=padding_mode)

            end_time = time.time()
            avg_time = (end_time - start_time) / num_iterations

            print(f"{mode:8} + {padding_mode:10}: {avg_time:.4f} 秒")

    print()

def example_batch_processing():
    """批处理示例"""
    print("=== 批处理示例 ===")

    # 创建多个图像
    batch_size = 3
    channels = 3
    height, width = 4, 4

    # 为每个批次创建不同的图像
    input_batch = np.zeros((batch_size, channels, height, width), dtype=np.float32)
    for b in range(batch_size):
        for c in range(channels):
            input_batch[b, c, :, :] = b * 10 + c

    print(f"输入批次形状: {input_batch.shape}")
    print("批次0, 通道0:")
    print(input_batch[0, 0])
    print("批次1, 通道0:")
    print(input_batch[1, 0])

    # 为每个批次创建不同的网格
    output_height, output_width = 2, 2
    grid_batch = np.zeros((batch_size, output_height, output_width, 2), dtype=np.float32)

    for b in range(batch_size):
        # 每个批次使用不同的采样模式
        scale = 0.5 + b * 0.2  # 不同的缩放因子
        for h in range(output_height):
            for w in range(output_width):
                grid_batch[b, h, w, 0] = (w * 2.0 / (output_width - 1) - 1) * scale
                grid_batch[b, h, w, 1] = (h * 2.0 / (output_height - 1) - 1) * scale

    # 批处理
    output_batch = grid_sampler_cuda.grid_sampler(
        input_batch, grid_batch,
        mode="bilinear",
        padding_mode="border"
    )

    print(f"输出批次形状: {output_batch.shape}")
    print("批次0, 通道0:")
    print(output_batch[0, 0])
    print("批次1, 通道0:")
    print(output_batch[1, 0])
    print()

if __name__ == "__main__":
    print("Grid Sampler CUDA 使用示例\n")

    try:
        example_basic_usage()
        example_image_rotation()
        example_performance_comparison()
        example_batch_processing()

        print("🎉 所有示例运行成功！")

    except Exception as e:
        print(f"❌ 示例运行失败: {e}")
        import traceback
        traceback.print_exc()