yolo分割模型转mask代码

输入图片，yolov8分割模型，输出图片
def run_object_detection(image,model_input):
    # Load ONNX model
    model = model_input

    # Read image using OpenCV
    img = image
    img_height, img_width, _ = img.shape
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (640, 640))
    input_image = np.array(img)
    input_image = input_image.transpose(2, 0, 1)
    input_image = input_image.reshape(1, 3, 640, 640).astype('float32')
    input_image = input_image / 255.0

    # Run the model

    outputs = model.run(None, {"images": input_image})

    # Process model output
    output0 = outputs[0][0].transpose()
    output1 = outputs[1][0]
    boxes = output0[:, 0:6]
    masks = output0[:, 6:]
    output1 = output1.reshape(32, 160 * 160)
    masks = masks @ output1
    boxes = np.hstack((boxes, masks))

    yolo_classes = [
        'chejiao', '0'
    ]

    def intersection(box1, box2):
        box1_x1, box1_y1, box1_x2, box1_y2 = box1[:4]
        box2_x1, box2_y1, box2_x2, box2_y2 = box2[:4]
        x1 = max(box1_x1, box2_x1)
        y1 = max(box1_y1, box2_y1)
        x2 = min(box1_x2, box2_x2)
        y2 = min(box1_y2, box2_y2)
        return (x2 - x1) * (y2 - y1)

    def union(box1, box2):
        box1_x1, box1_y1, box1_x2, box1_y2 = box1[:4]
        box2_x1, box2_y1, box2_x2, box2_y2 = box2[:4]
        box1_area = (box1_x2 - box1_x1) * (box1_y2 - box1_y1)
        box2_area = (box2_x2 - box2_x1) * (box2_y2 - box2_y1)
        return box1_area + box2_area - intersection(box1, box2)

    def iou(box1, box2):
        return intersection(box1, box2) / union(box1, box2)

    def get_mask(row, box):
        mask = row.reshape(160, 160)
        return mask

    objects = []
    for row in boxes:
        prob = row[4:6].max()
        if prob < 0.1:
            continue
        xc, yc, w, h = row[:4]
        class_id = row[4:6].argmax()
        x1 = (xc - w / 2) / 640 * img_width
        y1 = (yc - h / 2) / 640 * img_height
        x2 = (xc + w / 2) / 640 * img_width
        y2 = (yc + h / 2) / 640 * img_height
        label = yolo_classes[class_id]
        mask = get_mask(row[6:25684], (x1, y1, x2, y2))
        objects.append([x1, y1, x2, y2, label, prob, mask])

    objects.sort(key=lambda x: x[5], reverse=True)
    result = []
    while len(objects) > 0:
        result.append(objects[0])
        objects = [object for object in objects if iou(object, objects[0]) < 0.7]

    mask = result[0][6]
    mask = 1 / (1 + np.exp(-mask))
    mask = (mask > 0.9).astype('uint8') * 255
    img_mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)

    # Resize the mask to match the original image size
    img_mask = cv2.resize(img_mask, (img_width, img_height))

    return img_mask