app.py

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)