File size: 5,600 Bytes
840f4e2 676e2ff 840f4e2 2ef17d8 840f4e2 ef1e412 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
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) |