import numpy as np
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])

    snr = signal_mean / background_std if background_std != 0 else float('inf')

    return signal_mean, background_std, snr, signal, background

def process_image(image, roi_sketch):
    if image is None or roi_sketch is None:
        return None, None

    # Convert sketch to binary mask
    mask = np.array(roi_sketch)
    _, binary_mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)

    # Find contours in the binary mask
    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Convert contours to list of ROIs
    rois = [contour.squeeze().tolist() for contour in contours if len(contour) > 2]

    # 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']

    results = []
    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.Sketchpad(label="绘制ROI", type="numpy", shape=(256, 256))
    
    with gr.Row():
        result_text = gr.Textbox(label="结果")
        result_image = gr.Image(label="ROI 图像")

    inputs = [input_image, roi_sketch]
    outputs = [result_text, result_image]

    input_image.change(process_image, inputs=inputs, outputs=outputs)
    roi_sketch.change(process_image, inputs=inputs, outputs=outputs)

demo.launch()