Update app.py

app.py

@@ -1,45 +1,91 @@
-import gradio as gr
 import numpy as np
-from PIL import Image, ImageDraw
 import cv2
+import matplotlib.pyplot as plt
+import gradio as gr
+import tempfile
+import os
+from PIL import Image, ImageDraw
+
+def calculate_snr(channel, rois):
+    signal_mask = np.zeros(channel.shape, np.uint8)
+    for roi in rois:
+        cv2.fillPoly(signal_mask, [np.array(roi)], 255)
+    signal = cv2.bitwise_and(channel, channel, mask=signal_mask)
+
+    background_mask = cv2.bitwise_not(signal_mask)
+    background = cv2.bitwise_and(channel, channel, mask=background_mask)
+
+    signal_mean = np.mean(signal[signal_mask == 255])
+    background_std = np.std(background[background_mask == 255])
 
-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)
+    snr = signal_mean / background_std if background_std != 0 else float('inf')
 
-    # Create a mask for the ROI
-    mask = np.zeros_like(image)
-    cv2.fillPoly(mask, [roi], 255)
+    return signal_mean, background_std, snr, signal, background
 
-    # Calculate signal and noise
-    signal = np.mean(image[mask > 0])
-    noise = np.std(image[mask == 0])
+def process_image(image, roi_sketch):
+    if image is None or roi_sketch is None:
+        return None, None
 
-    # Calculate SNR
-    snr = signal / noise
+    # Convert sketch to binary mask
+    mask = np.array(roi_sketch)
+    _, binary_mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
 
-    return signal, noise, snr
+    # Find contours in the binary mask
+    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
-def process_image(input_image, roi_points):
-    # Load the input image
-    image = Image.open(input_image)
-    image = np.array(image)
+    # Convert contours to list of ROIs
+    rois = [contour.squeeze().tolist() for contour in contours if len(contour) > 2]
 
-    # 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)
+    # Convert image to numpy array
+    img_array = 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']
 
-    # Calculate SNR for each channel
     results = []
-    for i in range(image.shape[2]):
-        signal, noise, snr = calculate_snr(image[:, :, i], roi_image)
-        results.append((signal, noise, snr))
+    plt.figure(figsize=(5*len(channels), 10))
+
+    for i, (channel, name) in enumerate(zip(channels, channel_names)):
+        signal_mean, background_std, snr, signal, background = calculate_snr(channel, rois)
+        
+        results.append(f"{name} channel:\n"
+                       f"信号平均强度(Signal): {signal_mean:.2f}\n"
+                       f"背景标准差(Noise): {background_std:.2f}\n"
+                       f"信噪比(SNR): {snr:.2f}\n")
+
+        plt.subplot(2, len(channels), i+1)
+        plt.imshow(signal, cmap='gray')
+        plt.title(f'{name} Signal ROI')
+        plt.subplot(2, len(channels), i+1+len(channels))
+        plt.imshow(background, cmap='gray')
+        plt.title(f'{name} Background ROI')
+
+    result_filename = tempfile.mktemp(suffix='.png')
+    plt.savefig(result_filename)
+    plt.close()
+
+    return ("\n".join(results), result_filename)
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 图像SNR计算器")
+    gr.Markdown("上传一张图像，在第二个框中绘制多个感兴趣区域，结果会实时更新。支持单通道和多通道图像。")
+    
+    with gr.Row():
+        input_image = gr.Image(label="上传图像", type="numpy")
+        roi_sketch = gr.Image(label="绘制ROI", source="upload", tool="sketch", type="numpy", interactive=True)
+    
+    with gr.Row():
+        result_text = gr.Textbox(label="结果")
+        result_image = gr.Image(label="ROI 图像")
+
+    inputs = [input_image, roi_sketch]
+    outputs = [result_text, result_image]
 
-    return results, roi_image
+    input_image.change(process_image, inputs=inputs, outputs=outputs)
+    roi_sketch.change(process_image, inputs=inputs, outputs=outputs)
 
-iface = gr.Interface(fn=process_image, inputs=[gr.inputs.Image(), gr.inputs.Image()], outputs=[gr.outputs.Textbox(), gr.outputs.Image()], title="图像SNR计算器")
-iface.launch()
+demo.launch()