sam_segment.py

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=[
[
"https://3vj-render.3vjia.com//UpFile_Render/C00006070/PMC/DesignSchemeRenderFile/20240831/592351213526335564/43f8d835b3a54869a34167ed7f2a27aa.jpg?x-oss-process=image/resize,m_fill,h_730,w_1220", # 图片路径
"large", # 模型大小
0.4, # 置信度阈值
0.3 # IoU阈值，降低默认值
]
]
)

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

app.py

@@ -0,0 +1,120 @@
+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=[
+        [
+            "https://3vj-render.3vjia.com//UpFile_Render/C00006070/PMC/DesignSchemeRenderFile/20240831/592351213526335564/43f8d835b3a54869a34167ed7f2a27aa.jpg?x-oss-process=image/resize,m_fill,h_730,w_1220",  # 图片路径
+            "large",                   # 模型大小
+            0.4,                       # 置信度阈值
+            0.3                        # IoU阈值，降低默认值
+        ]
+    ]
+)
+
+# 启动应用
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")