update

.gitignore

@@ -1,2 +1,3 @@
 /__pycache__
 /output*
+*.mnn

app.py

@@ -1,3 +1,4 @@
+from cgi import test
 import gradio as gr
 import requests
 import os
@@ -11,6 +12,12 @@ import gdown
 import sys
 from typing import Optional
 import datetime
+from mnnsr import modelTest_for_gradio
+from PIL import Image
+import cv2
+
+
+
 
 log_to_terminal = True
 task_counter = 0
@@ -52,9 +59,7 @@ def download_gdrive_file(
         Optional[str]: 如果下载成功，返回文件的完整路径；否则返回 None。
     """
     print(f"--- 开始处理链接: {url} ---")
-
-
-    # 3. 准备下载路径并创建文件夹
+    # 准备下载路径并创建文件夹
     try:
         os.makedirs(folder, exist_ok=True)
     except OSError as e:
@@ -179,7 +184,7 @@ def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min:
 async def _process_model(model_input: Union[str, gr.File], tilesize: int, output_dir: str,task_id:int,fp16:bool):
     log = ('初始化日志记录...\n')
     print_log(task_id, '初始化日志记录', '开始')
-    yield [],[], log
+    yield [],log
 
     if isinstance(model_input, str):
         input_path = model_input
@@ -188,12 +193,12 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
         input_path = model_input.name
         log += f'已上传文件: {input_path}\n'
     print_log(task_id, log.split('\n')[-1], '开始')
-    yield [], [], log
+    yield [], log
 
     if not input_path:
         log += ( f'未获得正确的模型文件\n')
         print_log(task_id, f'未获得正确的模型文件', '错误')
-        yield  [],[], log
+        yield  [],log
         return
 
 
@@ -201,7 +206,7 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
         onnx_path = input_path
         log += ( '输入已经是 ONNX 文件\n')
         print_log(task_id, '输入已经是 ONNX 文件', '跳过')
-        yield [],[], log
+        yield [],log
     else:
         print_log(task_id, f'转换 PTH 模型为 ONNX, folder={output_dir}', '开始')
         onnx_path = convert_pth_to_onnx(input_path, tilesize=tilesize, output_folder=output_dir,use_fp16=fp16)
@@ -211,7 +216,7 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
         else:
             log += ( '生成ONNX模型失败\n')
             print_log(task_id, '生成ONNX模型', '错误')
-            yield [], [], log
+            yield [],log
             return
 
 
@@ -222,11 +227,11 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
         log += ( '正在将 ONNX 模型转换为 MNN 格式...\n')
         print_log(task_id, '正在将 ONNX 模型转换为 MNN 格式', '开始')
         convertmnn(onnx_path, mnn_path,fp16)
-        yield onnx_path,[], log
+        yield onnx_path,log
     except Exception as e:
         log += ( f'转换 MNN 模型时出错: {str(e)}\n')
         print_log(task_id, f'转换 MNN 模型时出错: {str(e)}', '错误')
-        yield onnx_path,[], log
+        yield onnx_path,log
 
     print_log(task_id, '模型转换任务完成', '完成')
     
@@ -240,7 +245,7 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
     yield onnx_path, mnn_path, log
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 模型转换工具")
+    gr.Markdown("# MNN模型转换工具")
     model_type_opt = ['从链接下载', '直接上传文件']
     with gr.Row():
         with gr.Column():
@@ -259,13 +264,20 @@ with gr.Blocks() as demo:
             tilesize = gr.Number(label="Tilesize", value=0, precision=0)
             # 添加fp16和try_run复选框
             fp16 = gr.Checkbox(label="FP16", value=False)
-            try_run = gr.Checkbox(label="测试mnn推理", value=False)
+            try_run = gr.Checkbox(label="pymnnsr测试", value=False)
             convert_btn = gr.Button("开始转换")
         with gr.Column():
         # with gr.Row():
             log_box = gr.Textbox(label="转换日志", lines=10, interactive=False)
-            onnx_output = gr.File(label="ONNX 模型输出")
-            mnn_output = gr.File(label="MNN 模型输出")
+            onnx_output = gr.File(label="ONNX 模型输出",file_types=["filepath"])
+            mnn_output = gr.File(label="MNN 模型输出",file_types=["filepath"])
+            img_output = gr.Image(type="pil", label="测试输出(边缘正常说明模型转换成功，色彩有bug)" ,visible=False)
+            def show_try_run(try_run):
+                if try_run:
+                    return gr.update(visible=True)
+                else:
+                    return gr.update(visible=False)
+            try_run.change(show_try_run, inputs=try_run, outputs=img_output)
 
     async def process_model(input_type, url_input, file_input, tilesize, fp16, try_run):
         global task_counter
@@ -279,7 +291,7 @@ with gr.Blocks() as demo:
         if input_type == model_type_opt[0] and url_input:
             log = f'正在下载模型文件: {url_input}\n'
             print_log(task_counter, f'正在下载模型文件: {url_input}', '开始')
-            yield None, None, log
+            yield None, None, log, None
 
             if url_input.startswith("https://drive.google.com/"):
                 model_input = download_gdrive_file(
@@ -296,22 +308,21 @@ with gr.Blocks() as demo:
                     filesize_min=1024     # 1KB
                 )
 
-
             if not model_input:
                 log += f'\n模型文件下载失败\n'
                 print_log(task_counter, f'模型文件载', '失败')
-                yield None, None, log
+                yield None, None, log, None
                 return
             
             log += f'\n模型文件已下载到: {model_input}\n'
             print_log(task_counter, f'模型文件已下载到: {model_input}', '完成')
-            yield None, None, log
+            yield None, None, log, None
         elif input_type == model_type_opt[1] and file_input:
             model_input = file_input
         else:
             # 改为通过yield返回错误日志
             log = '\n请选择输入类型并提供有效的输入！'
-            yield None, None, log
+            yield None, None, log,None
             return
         
         onnx_path = None
@@ -320,17 +331,23 @@ with gr.Blocks() as demo:
         async for result in _process_model(model_input, int(tilesize), output_dir, task_counter, fp16):
             if isinstance(result, tuple) and len(result) == 3:
                 onnx_path, mnn_path, process_log = result
-                yield onnx_path, mnn_path, log+process_log
+                yield onnx_path, mnn_path, log+process_log,None
             elif isinstance(result, tuple) and len(result) == 2:
                 # 处理纯日志yield
                 _, process_log = result
-                yield None, None, log+process_log
+                yield None, None, log+process_log,None
             # yield onnx_path, mnn_path, log+process_log
 
+            if mnn_path and try_run:
+                processed_image_np = modelTest_for_gradio(mnn_path, "./sample.jpg")
+                processed_image_pil = Image.fromarray(cv2.cvtColor(processed_image_np, cv2.COLOR_BGR2RGB))
+                # processed_image_pil = Image.fromarray(processed_image_np)
+                yield onnx_path, mnn_path, log+process_log,processed_image_pil
+
     convert_btn.click(
         process_model,
         inputs=[input_type, url_input, file_input, tilesize, fp16, try_run],
-        outputs=[onnx_output, mnn_output, log_box],
+        outputs=[onnx_output, mnn_output, log_box, img_output],
         api_name="convert_model"
     )
 

app_mnnsr.py

@@ -0,0 +1,28 @@
+import gradio as gr
+import os
+import MNN
+import numpy as np
+import cv2
+from PIL import Image
+from mnnsr import modelTest_for_gradio
+
+def gradio_interface(modelPath, input_image):
+    processed_image_np = modelTest_for_gradio(modelPath, input_image)
+    
+    processed_image_pil = Image.fromarray(cv2.cvtColor(processed_image_np, cv2.COLOR_BGR2RGB))
+    
+    return processed_image_pil
+
+# 创建Gradio界面
+iface = gr.Interface(
+    fn=gradio_interface,
+    # inputs=gr.Image(type="pil", label="上传图像"),
+    inputs = [gr.File(label="上传MNN模型"), gr.Image(type="filepath", label="上传图像",value="./sample.jpg")],
+    outputs=gr.Image(type="pil", label="处理后的图像"),
+    title="MNN图像超分辨率处理",
+    description="上传图像，使用MNN模型进行超分辨率处理"
+)
+
+if __name__ == "__main__":
+    # 启动Gradio界面
+    iface.launch()

mnnsr.py

@@ -0,0 +1,150 @@
+import gradio as gr
+import os
+import MNN
+import numpy as np
+import cv2
+from PIL import Image
+
+
+# 复制原始modelTest函数中的必要函数
+def process_image(data, H, W, C, color='BGR'):
+    """
+    处理图像数据（灰度或彩色）并转换为指定色彩空间
+    
+    参数:
+        data: 输入数据指针 (numpy数组形式传入)
+        H: 高度
+        W: 宽度
+        C: 通道数 (1 或 3)
+        color: 目标色彩空间 ('BGR', 'RGB', 'YCbCr', 'YUV')
+    
+    返回:
+        numpy数组 处理后的图像
+    """
+    if C == 1:
+        # 灰度图像处理
+        gray = np.array(data, dtype=np.float32).reshape(H, W)
+        gray = (gray * 255).astype(np.uint8)
+        result = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
+        return result
+    else:
+        # 彩色图像处理
+        # 将数据拆分为各个通道
+        channels = [
+            np.array(data[i*H*W : (i+1)*H*W], dtype=np.float32).reshape(H, W)
+            for i in range(C)
+        ]
+        if color == 'RGB':
+            # RGB -> BGR (OpenCV默认顺序)
+            channels[0], channels[2] = channels[2], channels[0]  # 交换R和B通道
+            result = cv2.merge(channels)
+            result = (result * 255).astype(np.uint8)
+        else:
+            # 先合并为BGR格式
+            rgb = cv2.merge(channels)
+            rgb = (rgb * 255).astype(np.uint8)
+            # 转换为目标色彩空间
+            if color == 'YCbCr':
+                result = cv2.cvtColor(rgb, cv2.COLOR_BGR2YCrCb)
+            elif color == 'YUV':
+                result = cv2.cvtColor(rgb, cv2.COLOR_BGR2YUV)
+            else:
+                result = rgb  # 默认为BGR
+            
+        return result
+
+def createTensor(tensor):
+    shape = tensor.getShape()
+    data = np.ones(shape, dtype=np.float32)
+    return MNN.Tensor(shape, tensor.getDataType(), data, tensor.getDimensionType())
+
+def modelTest_for_gradio(modelPath, image_path, tilesize = 128):
+    model_name = os.path.basename(modelPath)
+    if "-Grayscale" in model_name:
+        model_channel = 1
+    elif "-4ch" in model_name:
+        model_channel = 4
+    else:
+        model_channel = 3
+
+    net = MNN.Interpreter(modelPath)
+    # set 9 for Session_Backend_Auto, Let BackGround Tuning
+    net.setSessionMode(9)
+    # set 0 for tune_num
+    # net.setSessionHint(0, 20)
+    config = {}
+    # "CPU"或0(默认), "OPENCL"或3,"OPENGL"或6, "VULKAN"或7, "METAL"或1, "TRT"或9, "CUDA"或2, "HIAI"或8
+    config['backend'] = 3
+    #config['precision'] = "low"
+    session = net.createSession(config)
+    
+    print("Run on backendtype: %d \n" % net.getSessionInfo(session, 2))
+
+    # 读取图像
+    image = cv2.imread(image_path)
+    if image.ndim == 2:
+        # 为了方便处理，先将其扩展为3维数组 (height, width, 1)
+        print("extend dims")
+        image = np.expand_dims(image, axis=-1)
+    image_channel = image.shape[2]
+    
+    image = cv2.resize(image, (tilesize, tilesize))
+
+    # 处理通道数不匹配的情况
+    if image_channel == 3:
+        if model_channel == 1:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        elif model_channel == 3:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        elif model_channel == 4:    
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
+        else:
+            print(f"unexpect input: model_channel {model_channel}, image_channel {image_channel}")
+    elif image_channel == 4:
+        if model_channel == 1:
+            image = cv2.cvtColor(image, cv2.COLOR_BGRA2GRAY)
+        elif model_channel == 3:
+            image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
+        elif model_channel == 4:
+            image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
+        else:
+            print(f"unexpect input: model_channel {model_channel}, image_channel {image_channel}")
+    else:
+        print(f"unexpect input: model_channel {model_channel}, image_channel {image_channel}")
+
+    # 显示图像(在Gradio中不需要)
+    # display(Image(data=cv2.imencode('.jpg', image)[1].tobytes()))
+    
+    image = image/255.0
+    #preprocess it
+    image = image.transpose((2, 0, 1))
+    #change numpy data type as np.float32 to match tensor's format
+    image = image.astype(np.float32)
+    #cv2 read shape is NHWC, Tensor's need is NCHW,transpose it
+    tmp_input = MNN.Tensor((1, model_channel, tilesize, tilesize), MNN.Halide_Type_Float, image, MNN.Tensor_DimensionType_Caffe)
+    
+    # input
+    inputTensor = net.getSessionInput(session)
+    net.resizeTensor(inputTensor, (1, model_channel, tilesize, tilesize))
+    net.resizeSession(session)
+    inputTensor.copyFrom(tmp_input)
+    # infer
+    net.runSession(session)
+    outputTensor = net.getSessionOutput(session)
+    # output
+    outputShape = outputTensor.getShape()
+    print("outputShape",outputShape)
+    outputHost = createTensor(outputTensor)
+    outputTensor.copyToHostTensor(outputHost)
+    
+    outimage = process_image(outputHost.getData(), outputShape[2], outputShape[3], outputShape[1], color='RGB')
+    
+    # 返回处理后的图像(numpy数组)
+    return outimage
+
+# def gradio_interface(modelPath, input_image):
+#     processed_image_np = modelTest_for_gradio(modelPath, input_image)
+    
+#     processed_image_pil = Image.fromarray(cv2.cvtColor(processed_image_np, cv2.COLOR_BGR2RGB))
+    
+#     return processed_image_pil

requirements.txt

@@ -3,6 +3,9 @@ torch
 pnnx
 onnx
 onnxsim
-mnn
 gradio
 gdown
+MNN
+numpy
+opencv-python
+Pillow