AXERA-TECH/CodeFormer

1import gradio as gr2import os3import tempfile4import numpy as np5import axengine as axe6import cv27from utils.restoration_helper import RestoreHelper8import socket9 10restore_helper = RestoreHelper(11        upscale_factor=1,12        face_size=512,13        crop_ratio=(1, 1),14        det_model="../model/yolov5l-face.axmodel",15        res_model="../model/codeformer.axmodel",16        bg_model="../model/realesrgan-x2.axmodel",17        save_ext='png',18        use_parse=True19        )20 21def face(img_path, session):22 23    output_names = [x.name for x in session.get_outputs()]24    input_name = session.get_inputs()[0].name25 26    ori_image = cv2.imread(img_path)27    h, w = ori_image.shape[:2]28    image = cv2.resize(ori_image, (512, 512))29    image = (image[..., ::-1] /255.0).astype(np.float32)30    31    mean = [0.5, 0.5, 0.5]32    std = [0.5, 0.5, 0.5]33    image = ((image - mean) / std).astype(np.float32)34 35    #image = (image /1.0).astype(np.float32)36    img = np.transpose(np.expand_dims(np.ascontiguousarray(image), axis=0), (0,3,1,2))37    38    # Use the model to generate super-resolved images39    sr = session.run(output_names, {input_name: img})40 41    #sr_y_image = imgproc.array_to_image(sr)42    sr = np.transpose(sr[0].squeeze(0), (1,2,0))43    sr = (sr*std + mean).astype(np.float32)44    45    # Save image46    ndarr = np.clip((sr*255.0), 0, 255.0).astype(np.uint8)47    out_image = cv2.resize(ndarr[..., ::-1], (w, h))48 49    return out_image50 51def full_image(img_path, restore_helper=restore_helper):52 53    restore_helper.clean_all()54    img = cv2.imread(img_path, cv2.IMREAD_COLOR)55 56    restore_helper.read_image(img)57    # get face landmarks for each face58    num_det_faces = restore_helper.get_face_landmarks_5(59        only_center_face=False, resize=640, eye_dist_threshold=5)60    # align and warp each face61    restore_helper.align_warp_face()62    # face restoration for each cropped face63    for idx, cropped_face in enumerate(restore_helper.cropped_faces):64        # prepare data65        cropped_face_t = (cropped_face.astype(np.float32) / 255.0) * 2.0 - 1.066        cropped_face_t = np.transpose(67            np.expand_dims(np.ascontiguousarray(cropped_face_t[...,::-1]), axis=0), 68            (0,3,1,2)69        )70        #print('cropped_face_t', cropped_face_t.shape)71 72        try:73            ort_outs = restore_helper.rs_sessison.run(74                restore_helper.rs_output, 75                {restore_helper.rs_input: cropped_face_t}76                )77            restored_face = ort_outs[0]78            restored_face = (restored_face.squeeze().transpose(1, 2, 0) * 0.5 + 0.5) * 25579            restored_face = np.clip(restored_face[...,::-1], 0, 255).astype(np.uint8)80        except Exception as error:81            print(f'\tFailed inference for CodeFormer: {error}')82            restored_face = (cropped_face_t.squeeze().transpose(1, 2, 0) * 0.5 + 0.5) * 25583            restored_face = np.clip(restored_face, 0, 255).astype(np.uint8)84 85        restored_face = restored_face.astype('uint8')86        restore_helper.add_restored_face(restored_face, cropped_face)87 88    # upsample the background89    # Now only support RealESRGAN for upsampling background90    bg_img = restore_helper.background_upsampling(img)91    restore_helper.get_inverse_affine(None)92    # paste each restored face to the input image93    restored_img = restore_helper.paste_faces_to_input_image(upsample_img=bg_img, draw_box=False)94 95    return restored_img96    97 98def colorize_image(input_img_path: str, model_name: str, progress=gr.Progress()):99    if not input_img_path:100        raise gr.Error("未上传图片")101 102    # 加载图像103    progress(0.3, desc="加载图像...")104 105    # 根据模型选择调用不同函数106    if model_name == "Face":107        out = face(input_img_path, session=restore_helper.rs_sessison)108    else:109        out = full_image(input_img_path, restore_helper=restore_helper)110 111    progress(0.9, desc="保存结果...")112 113    # 保存到临时文件114    output_path = os.path.join(tempfile.gettempdir(), "restore_output.jpg")115    cv2.imwrite(output_path, out)116 117    progress(1.0, desc="完成！")118    return output_path119 120 121# ==============================122# Gradio 界面123# ==============================124custom_css = """125body, .gradio-container {126    font-family: 'Microsoft YaHei', 'PingFang SC', 'Helvetica Neue', Arial, sans-serif;127}128.model-buttons .wrap {129    display: flex;130    gap: 10px;131}132.model-buttons .wrap label {133    background-color: #f0f0f0;134    padding: 10px 20px;135    border-radius: 8px;136    cursor: pointer;137    text-align: center;138    font-weight: 600;139    border: 2px solid transparent;140    flex: 1;141}142.model-buttons .wrap label:hover {143    background-color: #e0e0e0;144}145.model-buttons .wrap input[type="radio"]:checked + label {146    background-color: #4CAF50;147    color: white;148    border-color: #45a049;149}150"""151 152with gr.Blocks(title="人脸修复工具") as demo:153    gr.Markdown("## 🎨 人脸修复演示DEMO")154 155    with gr.Row(equal_height=True):156        # 左侧：输入区157        with gr.Column(scale=1, min_width=300):158            gr.Markdown("### 📤 输入")159            input_image = gr.Image(160                type="filepath",161                label="上传图片",162                sources=["upload"],163                height=300164            )165            166            gr.Markdown("### 🔧 选择修复模式")167            model_choice = gr.Radio(168                choices=["Face", "Full image"],169                value="Face",170                label=None,171                elem_classes="model-buttons"172            )173            174            run_btn = gr.Button("🚀 开始修复", variant="primary")175 176        # 右侧：输出区177        with gr.Column(scale=1, min_width=600):178            gr.Markdown("### 🖼️ 修复结果")179            output_image = gr.Image(180                label="修复后图片",181                interactive=False,182                height=600183            )184            download_btn = gr.File(label="📥 下载修复图片")185 186    # 绑定事件187    def on_colorize(img_path, model, progress=gr.Progress()):188        if img_path is None:189            raise gr.Error("请先上传图片！")190        try:191            result_path = colorize_image(img_path, model, progress=progress)192            return result_path, result_path193        except Exception as e:194            raise gr.Error(f"处理失败: {str(e)}")195 196    run_btn.click(197        fn=on_colorize,198        inputs=[input_image, model_choice],199        outputs=[output_image, download_btn]200    )201 202# 启动203def get_local_ip():204    """获取本机局域网IP地址"""205    try:206        # 创建一个UDP连接（不会真正发送数据）207        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:208            s.connect(("8.8.8.8", 80))  # 连接到公共DNS（Google）209            ip = s.getsockname()[0]210        return ip211    except Exception:212        # 回退到 localhost213        return "127.0.0.1"214 215 216if __name__ == "__main__":217    # demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())218 219    server_port = 7860220    server_name = "0.0.0.0"221    222    # 获取本机IP223    local_ip = get_local_ip()224    225    # 打印可点击的URL（大多数终端支持点击）226    print("\n" + "="*50)227    print("🌐 人脸修复工具 Web UI 已启动！")228    print(f"🔗 本地访问: http://127.0.0.1:{server_port}")229    if local_ip != "127.0.0.1":230        print(f"🔗 局域网访问: http://{local_ip}:{server_port}")231    print("="*50 + "\n")232 233    # 启动Gradio应用234    demo.launch(235        server_name=server_name,236        server_port=server_port,237        theme=gr.themes.Soft(),238        css=custom_css239    )