yolov5-master 实例分割推理 精简代码

以下是使用yolov5-master库进行实例分割推理的简化代码示例：

1. 导入所需库和模块：

import torch
from PIL import Image

from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression, scale_coords
from yolov5.utils.datasets import letterbox
from yolov5.utils.plots import plot_one_box

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

2. 定义推理函数：

def instance_segmentation(image_path):
    # 加载模型
    model_path = 'path/to/your/model.pt'
    model = attempt_load(model_path, map_location=device)  # 加载模型
    model.to(device).eval()

    # 读取图像
    image = Image.open(image_path)
    original_image = image.copy()
    image = letterbox(image, new_shape=model.imgsz)[0]

    # 转换为Tensor
    image = torch.from_numpy(image.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0).to(device)

    # 进行推理
    pred = model(image)[0]
    pred = non_max_suppression(pred, conf_thres=0.3, iou_thres=0.45)

    # 处理预测结果
    for det in pred[0]:
        if det is not None:
            det[:, :4] = scale_coords(image.shape[2:], det[:, :4], original_image.size).round()

            # 绘制边框和标签
            for *xyxy, conf, cls in reversed(det):
                plot_one_box(xyxy, original_image, label=f'{conf:.2f}', color=(0, 255, 0), line_thickness=3)

    # 显示结果
    original_image.show()

3. 调用推理函数：

if __name__ == '__main__':
    image_path = 'path/to/your/image.jpg'
    instance_segmentation(image_path)

请确保已将模型文件和待推理的图像文件路径正确设置，并对代码进行适当的调整以满足自己的需求。这是一个简化的示例代码，你还可以根据具体情况对其进行修改和优化。