import gradio as gr
import cv2
import numpy as np
from sam_segment import segment_image_with_prompt

# 预定义分割颜色组
SEGMENT_COLORS = [
    ((255, 99, 71), (255, 99, 71)),    # 红橙色
    ((65, 105, 225), (65, 105, 225)),  # 皇家蓝
    ((50, 205, 50), (50, 205, 50)),    # 酸橙绿
    ((255, 215, 0), (255, 215, 0)),    # 金色
    ((238, 130, 238), (238, 130, 238)), # 紫罗兰
    ((0, 191, 255), (0, 191, 255)),    # 深天蓝
    ((255, 165, 0), (255, 165, 0)),    # 橙色
    ((106, 90, 205), (106, 90, 205)),  # 石板蓝
]

def segment_image(input_image, model_size, conf_threshold, iou_threshold):
    """
    使用FastSAM模型对输入图片进行分割
    """
    try:
        # 进行预测
        results = segment_image_with_prompt(
            image=input_image,
            model_size=model_size,
            conf=conf_threshold,
            iou=iou_threshold,
        )
        
        # 创建输出图像的副本
        output_image = input_image.copy()
        
        # 获取图像尺寸
        h, w = output_image.shape[:2]
        
        # 创建一个总的遮罩层和一个累积掩码
        final_mask = np.zeros_like(output_image)
        accumulated_mask = np.zeros((h, w), dtype=np.uint8)
        
        # 为每个分割结果创建掩码
        for idx, points in enumerate(results["segments"]):
            # 将点列表转换为轮廓格式
            contour_points = np.array(points).reshape(-1, 2).astype(np.int32)
            
            # 创建空白掩码
            mask = np.zeros((h, w), dtype=np.uint8)
            
            # 填充轮廓
            cv2.fillPoly(mask, [contour_points], 1)
            
            # 更新累积掩码（避免重叠区域重复计算）
            mask = cv2.bitwise_and(mask, cv2.bitwise_not(accumulated_mask))
            accumulated_mask = cv2.bitwise_or(accumulated_mask, mask)
            
            # 使用预定义的颜色（循环使用）
            color_idx = idx % len(SEGMENT_COLORS)
            fill_color, stroke_color = SEGMENT_COLORS[color_idx]
            
            # 创建填充区域（半透明）
            fill_mask = np.zeros_like(output_image)
            fill_mask[mask > 0] = fill_color
            final_mask = cv2.addWeighted(final_mask, 1.0, fill_mask, 0.3, 0)
            
            # 绘制轮廓线
            cv2.drawContours(final_mask, [contour_points], -1, stroke_color, 2)
        
        # 混合原图和掩码
        output_image = cv2.addWeighted(output_image, 1.0, final_mask, 0.5, 0)
        
        return output_image
    
    except Exception as e:
        print(f"分割过程中出错: {str(e)}")
        return input_image

# 创建Gradio界面
demo = gr.Interface(
    fn=segment_image,
    inputs=[
        gr.Image(label="输入图片"),
        gr.Radio(
            choices=["small", "large"],
            value="large",
            label="模型大小",
            info="small: 更快但精度较低, large: 更慢但精度更高"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.4,
            step=0.1,
            label="置信度阈值",
            info="值越高，检测越严格"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.3,  # 降低默认值，使其能显示更多区域
            step=0.1,
            label="IoU阈值",
            info="值越低则保留更多重叠区域，值越高则保留更少重叠区域"
        )
    ],
    outputs=gr.Image(label="分割结果"),
    title="FastSAM图像分割演示",
    description="上传一张图片，调整参数，模型将对图片中的对象进行分割。",
    examples=[
        [
            "./images/test_1.png",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ],
        [
            "./images/test_2.jpg",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ],
        [
            "./images/test_3.jpg",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ],
        [
            "./images/test_4.jpg",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ],
        [
            "./images/test_5.jpg",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ],
        [
            "./images/test_6.jpg",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ],
        [
            "./images/test_7.jpg",  # 图片路径
            "large",                   # 模型大小
            0.3,                       # 置信度阈值
            0.3                        # IoU阈值，降低默认值
        ]
    ]
)

# 启动应用
if __name__ == "__main__":
    demo.launch(share=True)