Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import gradio as gr | |
| import cv2 | |
| import torch | |
| # import queue | |
| # import threading | |
| # 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) | |
| # # 预测队列 | |
| # prediction_queue = queue.Queue(maxsize=5) | |
| # # 线程锁 | |
| # lock = threading.Lock() | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| def predict(input, input_size=512, high_visual_quality=False): | |
| input_size = int(input_size) # 确保 imgsz 是整数 | |
| # # 获取线程锁 | |
| # with lock: | |
| # print('5') | |
| # # 将任务添加到队列 | |
| # prediction_queue.put((input, input_size, high_visual_quality)) | |
| # # 等待结果 | |
| # print('6') | |
| # fig = prediction_queue.get()[0] | |
| # print(fig) | |
| # return fig | |
| 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 worker(): | |
| # while True: | |
| # # 从队列获取任务 | |
| # print('1') | |
| # input, input_size, high_visual_quality = prediction_queue.get() | |
| # # 执行模型预测 | |
| # print('2') | |
| # results = model(input, device=device, retina_masks=True, iou=0.7, conf=0.25, imgsz=input_size) | |
| # print('3') | |
| # fig = fast_process(annotations=results[0].masks.data, | |
| # image=input, high_quality=high_visual_quality, device=device) | |
| # print('4') | |
| # # 将结果放回队列 | |
| # prediction_queue.put(fig) | |
| # # 在一个新的线程中启动工作函数 | |
| # threading.Thread(target=worker).start() | |
| # # 将耗时的函数包装在另一个函数中，用于控制队列和线程同步 | |
| # def process_request(): | |
| # while True: | |
| # if not request_queue.empty(): | |
| # # 如果请求队列不为空，则处理该请求 | |
| # try: | |
| # lock.put(1) # 加锁，防止同时处理多个请求 | |
| # input, input_size, high_visual_quality = request_queue.get() | |
| # fig = predict(input, input_size, high_visual_quality) | |
| # 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, label='input_size'), | |
| gr.components.Checkbox(value=False, label='high_visual_quality')], | |
| 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=True, | |
| title="Fast Segment Anything (Everything mode)" | |
| ) | |
| # # 定义一个请求处理函数，将请求添加到队列中 | |
| # def handle_request(value): | |
| # try: | |
| # request_queue.put_nowait(value) # 添加请求到队列 | |
| # except: | |
| # return "当前队列已满，请稍后再试！" | |
| # return None | |
| # # 添加请求处理函数到应用程序界面 | |
| # app_interface.call_function() | |
| app_interface.queue(concurrency_count=1, max_size=20) | |
| app_interface.launch() |