File size: 4,641 Bytes
58bb53a dc276a0 58bb53a 096308a 58bb53a 096308a 58bb53a dc276a0 58bb53a dc276a0 58bb53a 474f8dc dc276a0 58bb53a dc276a0 474f8dc 734274e 58bb53a dc276a0 58bb53a dc276a0 58bb53a dc276a0 58bb53a dc276a0 |
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 |
import gradio as gr
import torch
from ultralyticsplus import YOLO, render_result
import numpy as np
from PIL import Image
import cv2
model_names = [
"cell-seg.pt",
"yolov8n-seg.pt",
"yolov8s-seg.pt",
"yolov8m-seg.pt",
"yolov8l-seg.pt",
"yolov8x-seg.pt",
]
current_model_name = "cell-seg.pt"
model = YOLO(current_model_name)
def sort_instance_masks_by_centroid(instances_mask, reverse=False):
# Calculate centroid of each instance mask
centroids = []
for mask in instances_mask:
# Find contours of the mask
mask_np = mask.astype(np.uint8)
#mask_np[mask_np !=0] = 255
contours, hierarchy = cv2.findContours(mask_np, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# Calculate moments of the contour
moments = cv2.moments(contours[0])
# Calculate centroid coordinates
c_x = int(moments["m10"] / moments["m00"])
c_y = int(moments["m01"] / moments["m00"])
centroids.append((c_x, c_y))
# Sort instance masks by centroid coordinates
sorted_instances_mask = [instance_mask for _, instance_mask in
sorted(zip(centroids, instances_mask), reverse=reverse)]
return sorted_instances_mask
def visualize_masks(masks):
masks = masks.detach().cpu().numpy()
height, width = masks.shape[1:]
# 计算有多少个 mask
num_masks = masks.shape[0]
masks = sort_instance_masks_by_centroid(masks)
# 创建一个空白图像，背景颜色为黑色
img = Image.new('RGB', (width, height),(0,0,0))
#img.putpalette([0, 0, 0] * 256)
img_array = np.array(img)
colors = []
# 将每个 mask 标记为不同的颜色
for i in range(num_masks):
color = np.random.randint(0, 256, size=3)
colors.append(tuple(color))
#colorimg.paste
#colorimg = Image.new('RGB', (width,height), color=tuple(np.random.randint(0, 256, size=3)))
#mask_img_tmp = Image.fromarray(masks[i]).convert('RGB')
#mask_array = Image.fromarray(masks[i])
img_array[masks[i] != 0,:] = color
#mask_img = mask_img.putpalette(color)
#img.paste(mask_img,(0,0),mask_img_tmp)
#img.putpalette(color + (0,) * 253)
# 将 mask 根据颜色映射显示为 RGB 图像
img_rgb = Image.fromarray(img_array)
return img_rgb,colors
def yolov8_inference(
image = None,
model_name = None,
dest_width = 512,
dest_height = 512,
conf_threshold = 0.25,
iou_threshold = 0.45,
):
"""
YOLOv8 inference function
Args:
image: Input image
model_name: Name of the model
image_size: Image size
conf_threshold: Confidence threshold
iou_threshold: IOU threshold
Returns:
Rendered image
"""
global model
global current_model_name
if model_name != current_model_name:
model = YOLO(model_name)
current_model_name = model_name
model.overrides["conf"] = conf_threshold
model.overrides["iou"] = iou_threshold
#model.overrides["classes"] = [0]
results = model.predict(image)
renders = []
colorarray = []
for image_results in model.predict(image):
#print("predict results: ",type(image_results.masks))
#render = render_result(
# model=model, image=image, result=image_results
#)
#render ,colors= visualize_masks(image_results.masks.data)
#render = render.resize((dest_width,dest_height))
#renders.append(render)
#colorarray.append(colors)
renders.append(image_results)
num_masks = image_results.masks.data.shape[0]
return renders[0].plot(), "Total num: "+str(num_masks)
inputs = [
gr.Image(type="filepath", label="Input Image"),
gr.Dropdown(
model_names,
value=current_model_name,
label="Model type",
),
gr.inputs.Slider(minimum=128, maximum=2048, step=64, default=512, label="Width"),
gr.inputs.Slider(minimum=128, maximum=2048, step=64, default=512, label="Height"),
gr.Slider(
minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"
),
gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
]
outputs = [gr.Image(type="filepath", label="Output Image"),gr.Textbox(label="Output Text")]
title = "Ultralytics YOLOv8 Segmentation For HumanBody Only Now"
demo_app = gr.Interface(
fn=yolov8_inference,
inputs=inputs,
outputs=outputs,
title=title,
examples=None,
cache_examples=False,
theme="default",
)
demo_app.launch(debug=True, enable_queue=True) |