File size: 9,299 Bytes
c987532 d852f6a c987532 dc120b2 c987532 dc120b2 c987532 dc120b2 c987532 dc120b2 c987532 dc120b2 c987532 dc120b2 c987532 dc120b2 c987532 dc120b2 c987532 d852f6a c987532 ec68a7d dc120b2 c987532 dc120b2 d852f6a ec68a7d d852f6a c987532 dc120b2 c987532 dc120b2 d852f6a dc120b2 ec68a7d c987532 dc120b2 e0a1444 dc120b2 d852f6a c987532 d852f6a |
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 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 |
from ultralytics import YOLO
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
import cv2
import torch
# import queue
# import time
# from PIL import Image
model = YOLO('checkpoints/FastSAM.pt') # load a custom model
def fast_process(annotations, image, high_quality, device):
if isinstance(annotations[0],dict):
annotations = [annotation['segmentation'] for annotation in annotations]
original_h = image.height
original_w = image.width
fig = plt.figure(figsize=(10, 10))
plt.imshow(image)
if high_quality == True:
if isinstance(annotations[0],torch.Tensor):
annotations = np.array(annotations.cpu())
for i, mask in enumerate(annotations):
mask = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8))
annotations[i] = cv2.morphologyEx(mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8))
if device == 'cpu':
annotations = np.array(annotations)
fast_show_mask(annotations,
plt.gca(),
bbox=None,
points=None,
pointlabel=None,
retinamask=True,
target_height=original_h,
target_width=original_w)
else:
if isinstance(annotations[0],np.ndarray):
annotations = torch.from_numpy(annotations)
fast_show_mask_gpu(annotations,
plt.gca(),
bbox=None,
points=None,
pointlabel=None)
if isinstance(annotations, torch.Tensor):
annotations = annotations.cpu().numpy()
if high_quality == True:
contour_all = []
temp = np.zeros((original_h, original_w,1))
for i, mask in enumerate(annotations):
if type(mask) == dict:
mask = mask['segmentation']
annotation = mask.astype(np.uint8)
contours, _ = cv2.findContours(annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
for contour in contours:
contour_all.append(contour)
cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
color = np.array([0 / 255, 0 / 255, 255 / 255, 0.8])
contour_mask = temp / 225 * color.reshape(1, 1, -1)
plt.imshow(contour_mask)
plt.axis('off')
plt.tight_layout()
return fig
# CPU post process
def fast_show_mask(annotation, ax, bbox=None,
points=None, pointlabel=None,
retinamask=True, target_height=960,
target_width=960):
msak_sum = annotation.shape[0]
height = annotation.shape[1]
weight = annotation.shape[2]
# 将annotation 按照面积 排序
areas = np.sum(annotation, axis=(1, 2))
sorted_indices = np.argsort(areas)[::1]
annotation = annotation[sorted_indices]
index = (annotation != 0).argmax(axis=0)
color = np.random.random((msak_sum,1,1,3))
transparency = np.ones((msak_sum,1,1,1)) * 0.6
visual = np.concatenate([color,transparency],axis=-1)
mask_image = np.expand_dims(annotation,-1) * visual
show = np.zeros((height,weight,4))
h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij')
indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
# 使用向量化索引更新show的值
show[h_indices, w_indices, :] = mask_image[indices]
if bbox is not None:
x1, y1, x2, y2 = bbox
ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
# draw point
if points is not None:
plt.scatter([point[0] for i, point in enumerate(points) if pointlabel[i]==1], [point[1] for i, point in enumerate(points) if pointlabel[i]==1], s=20, c='y')
plt.scatter([point[0] for i, point in enumerate(points) if pointlabel[i]==0], [point[1] for i, point in enumerate(points) if pointlabel[i]==0], s=20, c='m')
if retinamask==False:
show = cv2.resize(show,(target_width,target_height),interpolation=cv2.INTER_NEAREST)
ax.imshow(show)
def fast_show_mask_gpu(annotation, ax,
bbox=None, points=None,
pointlabel=None):
msak_sum = annotation.shape[0]
height = annotation.shape[1]
weight = annotation.shape[2]
areas = torch.sum(annotation, dim=(1, 2))
sorted_indices = torch.argsort(areas, descending=False)
annotation = annotation[sorted_indices]
# 找每个位置第一个非零值下标
index = (annotation != 0).to(torch.long).argmax(dim=0)
color = torch.rand((msak_sum,1,1,3)).to(annotation.device)
transparency = torch.ones((msak_sum,1,1,1)).to(annotation.device) * 0.6
visual = torch.cat([color,transparency],dim=-1)
mask_image = torch.unsqueeze(annotation,-1) * visual
# 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
show = torch.zeros((height,weight,4)).to(annotation.device)
h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
# 使用向量化索引更新show的值
show[h_indices, w_indices, :] = mask_image[indices]
show_cpu = show.cpu().numpy()
if bbox is not None:
x1, y1, x2, y2 = bbox
ax.add_patch(plt.Rectangle((x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor='b', linewidth=1))
# draw point
if points is not None:
plt.scatter([point[0] for i, point in enumerate(points) if pointlabel[i]==1], [point[1] for i, point in enumerate(points) if pointlabel[i]==1], s=20, c='y')
plt.scatter([point[0] for i, point in enumerate(points) if pointlabel[i]==0], [point[1] for i, point in enumerate(points) if pointlabel[i]==0], s=20, c='m')
ax.imshow(show_cpu)
# # 建立请求队列和线程同步锁
# request_queue = queue.Queue(maxsize=10)
# lock = queue.Queue()
def predict(input, input_size=512, high_visual_quality=True):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
input_size = int(input_size) # 确保 imgsz 是整数
results = model(input, device=device, retina_masks=True, iou=0.7, conf=0.25, imgsz=input_size)
fig = fast_process(annotations=results[0].masks.data,
image=input, high_quality=high_visual_quality, device=device)
return fig
# # 将耗时的函数包装在另一个函数中，用于控制队列和线程同步
# def process_request():
# while True:
# if not request_queue.empty():
# # 如果请求队列不为空，则处理该请求
# try:
# lock.put(1) # 加锁，防止同时处理多个请求
# input_package = request_queue.get()
# fig = predict(input_package)
# request_queue.task_done() # 请求处理结束，移除请求
# lock.get() # 解锁
# yield fig # 返回预测结果
# except:
# lock.get() # 出错时也需要解锁
# else:
# # 如果请求队列为空，则等待新的请求到达
# time.sleep(1)
# input_size=1024
# high_quality_visual=True
# inp = 'assets/sa_192.jpg'
# input = Image.open(inp)
# device = 'cuda' if torch.cuda.is_available() else 'cpu'
# input_size = int(input_size) # 确保 imgsz 是整数
# results = model(input, device=device, retina_masks=True, iou=0.7, conf=0.25, imgsz=input_size)
# pil_image = fast_process(annotations=results[0].masks.data,
# image=input, high_quality=high_quality_visual, device=device)
app_interface = gr.Interface(fn=predict,
inputs=[gr.components.Image(type='pil'),
gr.components.Slider(minimum=512, maximum=1024, value=1024, step=64),
gr.components.Checkbox(value=True)],
outputs=['plot'],
# examples=[["assets/sa_8776.jpg", 1024, True]],
# ["assets/sa_1309.jpg", 1024]],
examples=[["assets/sa_192.jpg"], ["assets/sa_414.jpg"],
["assets/sa_561.jpg"], ["assets/sa_862.jpg"],
["assets/sa_1309.jpg"], ["assets/sa_8776.jpg"],
["assets/sa_10039.jpg"], ["assets/sa_11025.jpg"],],
cache_examples=False,
title="Fast Segment Anthing (Everything mode)"
)
# # 定义一个请求处理函数，将请求添加到队列中
# def handle_request(value):
# try:
# request_queue.put_nowait(value) # 添加请求到队列
# except:
# return "当前队列已满，请稍后再试！"
# return None
# # 添加请求处理函数到应用程序界面
# app_interface.add_transition("submit", handle_request)
app_interface.queue(concurrency_count=2)
app_interface.launch() |