Update app.py

app.py

@@ -1,85 +1,45 @@
-import numpy as np
-import cv2
-import matplotlib.pyplot as plt
 import gradio as gr
-import tempfile
-import os
+import numpy as np
 from PIL import Image, ImageDraw
+import cv2
 
-def calculate_snr(channel, roi_points):
-    # 创建信号和背景掩码
-    mask = Image.new('L', (channel.shape[1], channel.shape[0]), 0)
-    ImageDraw.Draw(mask).polygon(roi_points, outline=1, fill=1)
-    signal_mask = np.array(mask)
-    
-    # 计算信号和背景区域
-    signal = np.where(signal_mask, channel, 0)
-    background = np.where(~signal_mask, channel, 0)
+def calculate_snr(image, roi):
+    # Convert image to grayscale if it's not already
+    if len(image.shape) == 3:
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
 
-    # 计算信号区域的平均强度
-    signal_mean = np.mean(signal[signal_mask == 1])
+    # Create a mask for the ROI
+    mask = np.zeros_like(image)
+    cv2.fillPoly(mask, [roi], 255)
 
-    # 计算背景区域的标准差
-    background_std = np.std(background[signal_mask == 0])
+    # Calculate signal and noise
+    signal = np.mean(image[mask > 0])
+    noise = np.std(image[mask == 0])
 
-    # 计算SNR
-    snr = signal_mean / background_std if background_std != 0 else float('inf')
+    # Calculate SNR
+    snr = signal / noise
 
-    return signal_mean, background_std, snr, signal, background
+    return signal, noise, snr
 
-def process_image(image, roi):
-    # 将图像转换为numpy数组
-    img_array = np.array(image)
+def process_image(input_image, roi_points):
+    # Load the input image
+    image = Image.open(input_image)
+    image = np.array(image)
 
-    # 判断图像是单通道还是多通道
-    if len(img_array.shape) == 2:
-        channels = [img_array]
-        channel_names = ['灰度']
-    else:
-        channels = cv2.split(img_array)
-        channel_names = ['Red', 'Green', 'Blue']
+    # Load the ROI image
+    roi_image = Image.new('L', (image.shape[1], image.shape[0]))
+    draw = ImageDraw.Draw(roi_image)
+    draw.polygon(roi_points, fill=255)
+    roi_image = np.array(roi_image)
 
+    # Calculate SNR for each channel
     results = []
-    result_images = []
-
-    for i, (channel, name) in enumerate(zip(channels, channel_names)):
-        signal_mean, background_std, snr, signal, background = calculate_snr(channel, roi)
-        
-        results.append(f"{name} channel:\n"
-                       f"信号平均强度(Signal): {signal_mean:.2f}\n"
-                       f"背景标准差(Noise): {background_std:.2f}\n"
-                       f"信噪比(SNR): {snr:.2f}\n")
-
-        # 创建结果图像
-        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
-        ax1.imshow(signal, cmap='gray')
-        ax1.set_title(f'{name} Signal ROI')
-        ax2.imshow(background, cmap='gray')
-        ax2.set_title(f'{name} Background ROI')
-        
-        # 保存图像结果到临时文件
-        result_filename = tempfile.mktemp(suffix='.png')
-        plt.savefig(result_filename)
-        plt.close(fig)
-        
-        result_images.append(result_filename)
-
-    return "\n".join(results), result_images
+    for i in range(image.shape[2]):
+        signal, noise, snr = calculate_snr(image[:, :, i], roi_image)
+        results.append((signal, noise, snr))
 
-# 创建Gradio接口
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[
-        gr.Image(label="上传图像", type="pil"),
-        gr.Image(label="绘制ROI", type="pil", tool="sketch", interactive=True)
-    ],
-    outputs=[
-        gr.Textbox(label="结果"),
-        gr.Gallery(label="ROI 图像")
-    ],
-    title="图像SNR计算器",
-    description="上传一张图像，在第二个框中绘制感兴趣区域，然后点击提交按钮计算SNR。支持单通道和多通道图像。"
-)
+    return results, roi_image
 
-# 运行Gradio应用
-iface.launch()
+iface = gr.Interface(fn=process_image, inputs=[gr.inputs.Image(), gr.inputs.Image()], outputs=[gr.outputs.Textbox(), gr.outputs.Image()], title="图像SNR计算器")
+iface.launch()