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