util/utils.py

import torch
import numpy as np
from PIL import Image
import torch.nn.functional as F

def pad_to_multiple_of_16(latent, pad_value, patch_size=16):
    h, w = latent.size(2), latent.size(3)
    target_h = ((h - 1) // patch_size + 1) * patch_size
    target_w = ((w - 1) // patch_size + 1) * patch_size
    pad_h = (target_h - h) // 2
    pad_w = (target_w - w) // 2
    # 额外处理奇数padding的情况
    pad_h_extra = (target_h - h) % 2
    pad_w_extra = (target_w - w) % 2
    padded_latent = F.pad(latent, (pad_w, pad_w + pad_w_extra, pad_h, pad_h + pad_h_extra), mode='constant', value=pad_value)
    # print("After padding: ", padded_latent.shape)
    return padded_latent

def split_into_blocks(latent, patch_size=16):
    b, c, h, w = latent.size()
    blocks = latent.view(b, c, h // patch_size, patch_size, w // patch_size, patch_size)
    blocks = blocks.permute(0, 2, 4, 1, 3, 5).contiguous().view(-1, c, patch_size, patch_size)
    # print("After splitting into blocks: ", blocks.shape)
    return blocks

def merge_blocks(blocks, original_shape, patch_size=16):
    b, c, h, w = original_shape
    num_blocks_per_row = w // patch_size
    num_blocks_per_col = h // patch_size

    # 恢复到原始形状的顺序
    blocks = blocks.view(b, num_blocks_per_col, num_blocks_per_row, c, patch_size, patch_size)
    blocks = blocks.permute(0, 3, 1, 4, 2, 5).contiguous()
    blocks = blocks.view(b, c, h, w)
    # print("After merging blocks: ", blocks.shape)
    return blocks

def crop_to_original_shape(blocks, original_shape):
    _, _, padded_height, padded_width = blocks.shape
    original_height, original_width = original_shape[2], original_shape[3]
    start_h = (padded_height - original_height) // 2
    end_h = start_h + original_height
    start_w = (padded_width - original_width) // 2
    end_w = start_w + original_width
    cropped_blocks = blocks[:, :, start_h:end_h, start_w:end_w]
    # print("After cropping to original shape: ", cropped_blocks.shape)
    return cropped_blocks

def adaptively_split_and_pad(image_tensor, pad_value, target_patch_size=16):
    """
    return: 
        patches_tensor: (N * num_blocks_h * num_blocks_w, c, target_patch_size, target_patch_size) patched tensors after spilt
        patch_sizes: a list, ori size of each blocks
        num_blocks_h, num_blocks_w
    """
    c, h, w = image_tensor.size(1), image_tensor.size(2), image_tensor.size(3)
    # 计算每个方向上的块数量
    num_blocks_h = h // target_patch_size if h % target_patch_size == 0 else h // target_patch_size + 1
    num_blocks_w = w // target_patch_size if w % target_patch_size == 0 else w // target_patch_size + 1

    # 确定每个块的尺寸
    block_h = h // num_blocks_h
    block_w = w // num_blocks_w
    patches = []
    patch_sizes = []

    for i in range(num_blocks_h):
        for j in range(num_blocks_w):
            # 计算每个块的起始和结束索引
            start_h = i * block_h
            start_w = j * block_w
            end_h = start_h + block_h if i < num_blocks_h - 1 else h
            end_w = start_w + block_w if j < num_blocks_w - 1 else w
            # 切割块
            patch = image_tensor[:, :, start_h:end_h, start_w:end_w]
            # 打印每个block在padding前的分辨率
            # print(f"Block {i*num_blocks_w + j} size before padding: {end_h - start_h}x{end_w - start_w}")

            # 计算每个块的padding需求
            pad_top = (target_patch_size - (end_h - start_h)) // 2
            pad_bottom = target_patch_size - (end_h - start_h) - pad_top
            pad_left = (target_patch_size - (end_w - start_w)) // 2
            pad_right = target_patch_size - (end_w - start_w) - pad_left

            # 应用padding
            patch_padded = F.pad(patch, (pad_left, pad_right, pad_top, pad_bottom), mode='constant', value=pad_value)

            patches.append(patch_padded)
            patch_sizes.append((end_h - start_h, end_w - start_w))
    # 将所有patch合并成一个tensor
    patches_tensor = torch.cat(patches, dim=0)
    return patches_tensor, patch_sizes, num_blocks_h, num_blocks_w


def crop_and_reconstruct(patches, patch_sizes, num_blocks_h, num_blocks_w, target_patch_size=16):
    """
    inverse operation of adaptively_split_and_pad
    """
    index = 0
    reconstructed_rows = []

    for i in range(num_blocks_h):
        row_patches = []
        for j in range(num_blocks_w):
            patch = patches[index]
            patch_height, patch_width = patch_sizes[index]

            valid_h_start = (target_patch_size - patch_height) // 2
            valid_w_start = (target_patch_size - patch_width) // 2
            valid_h_end = valid_h_start + patch_height
            valid_w_end = valid_w_start + patch_width

            cropped_patch = patch[:, valid_h_start:valid_h_end, valid_w_start:valid_w_end]
            row_patches.append(cropped_patch)
            index += 1
        row_tensor = torch.cat(row_patches, dim=2)
        reconstructed_rows.append(row_tensor)

    reconstructed_image = torch.cat(reconstructed_rows, dim=1)
    return reconstructed_image

def save_image(tensor, file_path):
    # 将张量转换为PIL图像并保存
    image = tensor.to('cpu').clone().detach()
    image = image.squeeze(0)
    image = torch.clamp(image, 0, 1)
    image = Image.fromarray((image.permute(1, 2, 0).numpy() * 255).astype(np.uint8))
    image.save(file_path)
    print(f"Image saved to {file_path}")


if __name__ == "__main__":
    # 假设有一个随机初始化的图像张量
    N, C, H, W = 1, 3, 36, 33  # 非标准尺寸，测试目的
    image_tensor = torch.rand(N, C, H, W)

    # 使用adaptively_split_and_pad函数
    target_patch_size = 16
    pad_value = 0  # 通常用于图像是黑色填充
    patches_tensor, patch_sizes, num_blocks_h, num_blocks_w = adaptively_split_and_pad(image_tensor, pad_value, target_patch_size)

    # 可视化每个block的crop结果
    for i, patch in enumerate(patches_tensor):
        save_image(patch, f"patch_{i}.png")

    # 使用crop_and_reconstruct函数
    reconstructed_image = crop_and_reconstruct(patches_tensor, patch_sizes, num_blocks_h, num_blocks_w, target_patch_size)

    # 保存和显示重建的图像
    save_image(reconstructed_image, "reconstructed_image.png")