feifeilib/feifeisharpened.py

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

def feifeisharpened(image,num_strength):
    # 将PIL图像转换为NumPy数组
    image_np = np.array(image)

    # 将NumPy数组转换为PyTorch张量
    image_tensor = (torch.tensor(image_np).permute(
        2, 0, 1).unsqueeze(0).float().to("cuda"))

    # 定义锐化滤镜，并调整中心值
    strength = num_strength
    sharpen_kernel = (torch.tensor(
        [
            [0, -1 * strength, 0],
            [-1 * strength, 1 + 4 * strength, -1 * strength],
            [0, -1 * strength, 0],
        ],
        dtype=torch.float32,
    ).unsqueeze(0).unsqueeze(0).to("cuda"))

    # 分别对每个通道应用卷积核
    sharpened_channels = []
    for i in range(3):
        channel_tensor = image_tensor[:, i:i + 1, :, :]
        sharpened_channel = F.conv2d(channel_tensor,
                                     sharpen_kernel,
                                     padding=1)
        sharpened_channels.append(sharpened_channel)

    # 合并通道
    sharpened_image_tensor = torch.cat(sharpened_channels, dim=1)

    # 将增强后的图像转换回PIL格式
    sharpened_image_np = (sharpened_image_tensor.squeeze(0).permute(
        1, 2, 0).cpu().numpy())
    sharpened_image_np = np.clip(sharpened_image_np, 0,
                                 255).astype(np.uint8)
    image = Image.fromarray(sharpened_image_np)
    return image