Upload gradio demo

.gitattributes

@@ -36,3 +36,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 image/sr_colorize.jpg filter=lfs diff=lfs merge=lfs -text
 model/colorize_stable.axmodel filter=lfs diff=lfs merge=lfs -text
 model/colorize_artistic.axmodel filter=lfs diff=lfs merge=lfs -text
+assert/gradio_demo.JPG filter=lfs diff=lfs merge=lfs -text

README.md

@@ -56,6 +56,25 @@ Input Data:
 |   `-- 1850Geography.jpg
 ```
 
+#### Inference with M.2 Accelerator card
+```
+$python3 gradio_demo.py
+[INFO] Available providers:  ['AXCLRTExecutionProvider']
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1 2295293f
+[INFO] Using provider: AXCLRTExecutionProvider
+[INFO] SOC Name: AX650N
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Compiler version: 5.0-patch1 2295293f
+* Running on local URL:  http://0.0.0.0:7860
+* To create a public link, set `share=True` in `launch()`.
+```
+Then use the M.2 Accelerator card IP instead of the 0.0.0.0, and use chrome open the URL: http://[your ip]:7860
+
+![gradio demo](./assert/gradio_demo.JPG)
+
 #### Inference with AX650 Host, such as M4N-Dock(爱芯派Pro)
 
 ```

python/gradio_demo.py

@@ -0,0 +1,178 @@
+import gradio as gr
+import os
+import tempfile
+from PIL import Image, ImageEnhance
+import numpy as np
+import axengine as axe
+import cv2
+
+# ==============================
+# 模拟上色函数（请替换为你的实际模型）
+# ==============================
+def init_DeOldifymodel(DeOldifyStable_path="../model/colorize_stable.axmodel", 
+                 DeOldifyArtistic_path="../model/colorize_artistic.axmodel"):
+    
+    DeOldifyStable_session = axe.InferenceSession(DeOldifyStable_path)
+    DeOldifyArtistic_session = axe.InferenceSession(DeOldifyArtistic_path)
+
+    return [DeOldifyStable_session, DeOldifyArtistic_session]
+
+DeOldify_sessions=init_DeOldifymodel()
+
+def from_numpy(x):
+    return x if isinstance(x, np.ndarray) else np.array(x)
+    
+def post_process(raw_color, orig):
+    color_np = np.asarray(raw_color)
+    orig_np = np.asarray(orig)
+    color_yuv = cv2.cvtColor(color_np, cv2.COLOR_RGB2YUV)
+    # do a black and white transform first to get better luminance values
+    orig_yuv = cv2.cvtColor(orig_np, cv2.COLOR_RGB2YUV)
+    hires = np.copy(orig_yuv)
+    hires[:, :, 1:3] = color_yuv[:, :, 1:3]
+    final = cv2.cvtColor(hires, cv2.COLOR_YUV2RGB)
+    return final
+
+def colorize_with_model(img_path, session):
+    output_names = [x.name for x in session.get_outputs()]
+    input_name = session.get_inputs()[0].name
+
+    ori_image = cv2.imread(img_path)
+    h, w = ori_image.shape[:2]
+    image = cv2.resize(ori_image, (512, 512))
+    image = (image[..., ::-1] /255.0).astype(np.float32)
+    
+    mean = [0.485, 0.456, 0.406]
+    std = [0.229, 0.224, 0.225]
+    image = ((image - mean) / std).astype(np.float32)
+
+    #image = (image /1.0).astype(np.float32)
+    image = np.transpose(np.expand_dims(np.ascontiguousarray(image), axis=0), (0,3,1,2))
+    
+    # Use the model to generate super-resolved images
+    sr = session.run(output_names, {input_name: image})
+    
+    if isinstance(sr, (list, tuple)):
+        sr = from_numpy(sr[0]) if len(sr) == 1 else [from_numpy(x) for x in sr]
+    else:
+        sr = from_numpy(sr)
+
+    #sr_y_image = imgproc.array_to_image(sr)
+    sr = np.transpose(sr.squeeze(0), (1,2,0))
+    sr = (sr*std + mean).astype(np.float32)
+    
+    # Save image
+    ndarr = np.clip((sr*255.0), 0, 255.0).astype(np.uint8)
+    ndarr = cv2.resize(ndarr[..., ::-1], (w, h))
+    out_image = post_process(ndarr, ori_image)
+
+    return out_image
+
+def colorize_image(input_img_path: str, model_name: str, progress=gr.Progress()):
+    if not input_img_path:
+        raise gr.Error("未上传图片")
+
+    # 加载图像
+    progress(0.3, desc="加载图像...")
+
+    # 根据模型选择调用不同函数
+    if model_name == "colorize_stable":
+        session = DeOldify_sessions[0]
+    else:
+        session = DeOldify_sessions[1]
+    out = colorize_with_model(input_img_path, session)
+
+    progress(0.9, desc="保存结果...")
+
+    # 保存到临时文件
+    output_path = os.path.join(tempfile.gettempdir(), "colorized_output.jpg")
+    cv2.imwrite(output_path, out)
+
+    progress(1.0, desc="完成！")
+    return output_path
+
+
+# ==============================
+# Gradio 界面
+# ==============================
+custom_css = """
+body, .gradio-container {
+    font-family: 'Microsoft YaHei', 'PingFang SC', 'Helvetica Neue', Arial, sans-serif;
+}
+.model-buttons .wrap {
+    display: flex;
+    gap: 10px;
+}
+.model-buttons .wrap label {
+    background-color: #f0f0f0;
+    padding: 10px 20px;
+    border-radius: 8px;
+    cursor: pointer;
+    text-align: center;
+    font-weight: 600;
+    border: 2px solid transparent;
+    flex: 1;
+}
+.model-buttons .wrap label:hover {
+    background-color: #e0e0e0;
+}
+.model-buttons .wrap input[type="radio"]:checked + label {
+    background-color: #4CAF50;
+    color: white;
+    border-color: #45a049;
+}
+"""
+
+with gr.Blocks(title="AI 图片上色工具") as demo:
+    gr.Markdown("## 🎨 AI 黑白图片自动上色演示")
+
+    with gr.Row(equal_height=True):
+        # 左侧：输入区
+        with gr.Column(scale=1, min_width=300):
+            gr.Markdown("### 📤 输入")
+            input_image = gr.Image(
+                type="filepath",
+                label="上传黑白/灰度图片",
+                sources=["upload"],
+                height=300
+            )
+            
+            gr.Markdown("### 🔧 选择上色模型")
+            model_choice = gr.Radio(
+                choices=["colorize_stable", "colorize_artistic"],
+                value="colorize_stable",
+                label=None,
+                elem_classes="model-buttons"
+            )
+            
+            run_btn = gr.Button("🚀 开始上色", variant="primary")
+
+        # 右侧：输出区
+        with gr.Column(scale=1, min_width=600):
+            gr.Markdown("### 🖼️ 上色结果")
+            output_image = gr.Image(
+                label="上色后图片",
+                interactive=False,
+                height=600
+            )
+            download_btn = gr.File(label="📥 下载上色图片")
+
+    # 绑定事件
+    def on_colorize(img_path, model, progress=gr.Progress()):
+        if img_path is None:
+            raise gr.Error("请先上传图片！")
+        try:
+            result_path = colorize_image(img_path, model, progress=progress)
+            return result_path, result_path
+        except Exception as e:
+            raise gr.Error(f"处理失败: {str(e)}")
+
+    run_btn.click(
+        fn=on_colorize,
+        inputs=[input_image, model_choice],
+        outputs=[output_image, download_btn]
+    )
+
+# 启动
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft(), css=custom_css)