# 导入必要的库和模块
from deoldify.device_id import DeviceId
from deoldify import device
from deoldify.visualize import get_image_colorizer
import gradio as gr
import torch
import warnings
from PIL import Image
import tempfile
import numpy as np

# 检查是否有可用的GPU，如果有，设置设备为GPU；否则使用CPU
if torch.cuda.is_available():
    device.set(device=DeviceId.GPU0)  
else:
    print("GPU 不可用，将使用 CPU。")  

# 优化并忽略一些警告信息
torch.backends.cudnn.benchmark = True  
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")  # 忽略特定警告

# 创建一个colorizer对象，用于图像上色
colorizer = get_image_colorizer(artistic=True)  

# 定义colorize_image函数，用于上色黑白照片
def colorize_image(image, render_factor):
    """
    颜色化黑白照片的函数。
    参数:
    - image: 用户上传的黑白照片。
    - render_factor: 渲染分辨率的因子，影响颜色化质量和处理时间。
    返回:
    - 上色化后的照片。
    """
    # 将numpy数组转换回PIL图像格式
    image = Image.fromarray(image)

    # 创建一个临时文件，用于保存PIL图像并作为colorizer的输入
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
        image_path = f.name  # 获取临时文件的路径
    image.save(image_path, "JPEG")  # 保存PIL图像为JPEG格式

    # 使用colorizer对图像进行上色，image_path是上色后的图像路径
    colorized_image_path = colorizer.plot_transformed_image(image_path, render_factor=render_factor, compare=False)

    # 打开上色后的图像文件，以PIL图像格式读取
    colorized_image = Image.open(colorized_image_path)
    
    # 将PIL Image转换为numpy数组，以便于在Gradio界面中显示
    result_image = np.array(colorized_image)

    # 返回numpy数组形式的上色图像
    return result_image

# 创建Gradio界面，使用colorize_image函数作为后端逻辑
iface = gr.Interface(
    fn=colorize_image,  
    inputs=[
        gr.inputs.Image(label="上传黑白照片"),  
        gr.inputs.Slider(  
            minimum=10, maximum=40, step=1, default=10, label="render_factor",
        )
    ],
    outputs=gr.outputs.Image(label="上色后的照片"),  
    title="图片上色器",  
    description="上传一张黑白图片并通过滑动条选择渲染分辨率来进行上色。",  
    live=False  
)

# 启动Gradio网页应用
iface.launch()