基于YOLOv8的船舶目标检测与分割(ONNX模型)

项目背景

  • 需求分析:在海洋监控、港口管理、海事安全等领域,自动化的船只检测与分割技术对于提高效率和安全性至关重要。
  • 技术选型:YOLOv8是YOLO系列的一个较新版本,以其速度快、准确率高而著称。使用ONNX(Open Neural Network Exchange)格式可以跨平台部署模型,并且通常能够获得更好的性能。

技术栈

  • Python:主要编程语言。
  • PyTorch:用于训练和加载YOLOv8模型。
  • ONNX:用于模型转换和部署。
  • OpenCV:用于图像处理和显示结果。
  • Pillow:用于读取和保存图像文件。

项目结构

  1. 数据准备

    • 收集带有标注的船只图像数据集。
    • 将数据集划分为训练集和测试集。
  2. 模型训练

    • 使用YOLOv8框架训练模型。
    • 调整超参数以优化检测和分割性能。
  3. 模型转换

    • 将训练好的PyTorch模型导出为ONNX格式。
    • 验证ONNX模型的正确性。
  4. 推理部署

    • 编写推理代码,支持从图像或视频流中检测并分割船只。
    • 使用ONNX Runtime进行高效推理。
  5. 结果展示

    • 可视化检测结果,包括边界框和分割掩码。
    • 计算并报告性能指标如准确率、召回率等。

示例代码

一个简化的示例代码片段,用于演示如何加载一个ONNX模型并在单张图片上进行船只检测和分割:

1import cv2
2import numpy as np
3import onnxruntime as ort
4
5# 加载ONNX模型
6ort_session = ort.InferenceSession("yolov8.onnx")
7
8# 加载图像
9img = cv2.imread('input.jpg')
10img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
11img = img.astype(np.float32) / 255.0
12img = np.expand_dims(img, axis=0)
13
14# 进行推理
15outputs = ort_session.run(None, {'images': img})
16
17# 解析输出
18boxes, scores, labels, masks = outputs
19
20# 可视化结果
21for box, score, label, mask in zip(boxes[0], scores[0], labels[0], masks[0]):
22    if score > 0.5:
23        x1, y1, x2, y2 = box
24        # 绘制边界框
25        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
26        # 绘制分割掩码
27        mask = (mask > 0.5).astype(np.uint8) * 255
28        img_masked = cv2.bitwise_and(img, img, mask=mask)
29        img = cv2.addWeighted(img, 1, img_masked, 0.5, 0)
30
31cv2.imshow('Detection and Segmentation', cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
32cv2.waitKey(0)
33cv2.destroyAllWindows()

首先,确保安装了必要的库:

1pip install onnxruntime opencv-python Pillow numpy

接下来是Python代码示例:

1import cv2
2import numpy as np
3from PIL import Image
4import onnxruntime as ort
5
6def letterbox_image(image, new_size):
7    """
8    Resize the image with unchanged aspect ratio using padding.
9    """
10    old_size = image.shape[:2]  # old_size is in (height, width) format
11    ratio = min(new_size[0] / old_size[0], new_size[1] / old_size[1])
12    new_size = tuple([int(x * ratio) for x in old_size])
13    image = cv2.resize(image, (new_size[1], new_size[0]))
14    
15    delta_w = new_size[1] - new_size[1]
16    delta_h = new_size[0] - new_size[0]
17    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
18    left, right = delta_w // 2, delta_w - (delta_w // 2)
19    
20    color = [0, 0, 0]
21    new_img = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
22    return new_img, ratio, (top, left)
23
24def preprocess_image(image, input_size=(640, 640)):
25    """
26    Preprocess the image to match the input size of the model.
27    """
28    img, _, _ = letterbox_image(image, input_size)
29    img = img[:, :, ::-1].transpose((2, 0, 1))  # BGR to RGB, HWC to CHW
30    img = np.ascontiguousarray(img, dtype=np.float32) / 255.0
31    return img
32
33def postprocess_output(output, confidence_threshold=0.5, iou_threshold=0.5):
34    """
35    Postprocess the output from the model.
36    """
37    boxes = output[0][0]
38    scores = output[0][1]
39    labels = output[0][2]
40    masks = output[0][3]
41
42    # Apply non-max suppression
43    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), confidence_threshold, iou_threshold)
44
45    filtered_boxes = []
46    filtered_scores = []
47    filtered_labels = []
48    filtered_masks = []
49
50    for i in indices:
51        idx = i[0]
52        filtered_boxes.append(boxes[idx])
53        filtered_scores.append(scores[idx])
54        filtered_labels.append(labels[idx])
55        filtered_masks.append(masks[idx])
56
57    return filtered_boxes, filtered_scores, filtered_labels, filtered_masks
58
59def visualize(image, boxes, scores, labels, masks, orig_image_shape, ratio, padding):
60    """
61    Visualize the detection results.
62    """
63    top, left = padding
64    for box, score, label, mask in zip(boxes, scores, labels, masks):
65        box = np.array(box).astype(int)
66        box /= ratio
67        box[[0, 2]] -= left
68        box[[1, 3]] -= top
69        box = box.clip(min=0)
70        
71        # Draw bounding box
72        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
73        
74        # Draw segmentation mask
75        mask = (mask > 0.5).astype(np.uint8) * 255
76        mask = cv2.resize(mask, (orig_image_shape[1], orig_image_shape[0]))
77        image_masked = cv2.bitwise_and(image, image, mask=mask)
78        image = cv2.addWeighted(image, 1, image_masked, 0.5, 0)
79
80    return image
81
82# Load the ONNX model
83ort_session = ort.InferenceSession("yolov8.onnx")
84
85# Load an example image
86image_path = 'input.jpg'
87image = cv2.imread(image_path)
88orig_image_shape = image.shape[:2]
89
90# Preprocess the image
91input_image = preprocess_image(image)
92input_image = np.expand_dims(input_image, axis=0)
93
94# Perform inference
95outputs = ort_session.run(None, {'images': input_image})
96
97# Postprocess the output
98filtered_boxes, filtered_scores, filtered_labels, filtered_masks = postprocess_output(outputs)
99
100# Visualize the results
101visualized_image = visualize(image, filtered_boxes, filtered_scores, filtered_labels, filtered_masks, orig_image_shape, 1.0, (0, 0))
102
103# Display the result
104cv2.imshow('Detection and Segmentation', visualized_image)
105cv2.waitKey(0)
106cv2.destroyAllWindows()

代码说明

  1. letterbox_image: 保持原始图像的长宽比不变,通过填充的方式调整图像大小。
  2. preprocess_image: 图像预处理函数,将图像调整到模型所需的尺寸,并将其转换为合适的格式。
  3. postprocess_output: 后处理函数,对模型输出进行非极大值抑制 (NMS),过滤掉低置信度和重叠的预测。
  4. visualize: 结果可视化函数,用于绘制边界框和分割掩码。

注意事项

  • 在运行这段代码之前,请确保已经训练了一个YOLOv8模型,并将其导出为ONNX格式。你可以从YOLOv8的官方仓库获取相应的代码或者使用预训练的模型。
  • 本示例假设模型输出包含四个维度:边界框坐标、置信度分数、类别标签以及分割掩码。
  • 对于实际应用,你可能还需要考虑更多的因素,例如模型的输入输出布局、后处理的具体细节等。

代码仅为示例,实际应用中可能需要根据具体需求调整细节。此外,确保安装了所有必要的库,并正确配置了环境。

相关推荐
江_小_白24 分钟前
自动驾驶之激光雷达
人工智能·机器学习·自动驾驶
yusaisai大鱼2 小时前
TensorFlow如何调用GPU?
人工智能·tensorflow
LNTON羚通2 小时前
摄像机视频分析软件下载LiteAIServer视频智能分析平台玩手机打电话检测算法技术的实现
算法·目标检测·音视频·监控·视频监控
珠海新立电子科技有限公司4 小时前
FPC柔性线路板与智能生活的融合
人工智能·生活·制造
IT古董4 小时前
【机器学习】机器学习中用到的高等数学知识-8. 图论 (Graph Theory)
人工智能·机器学习·图论
曼城周杰伦5 小时前
自然语言处理:第六十三章 阿里Qwen2 & 2.5系列
人工智能·阿里云·语言模型·自然语言处理·chatgpt·nlp·gpt-3
吃肉不能购5 小时前
Label-studio-ml-backend 和YOLOV8 YOLO11自动化标注,目标检测,实例分割,图像分类,关键点估计,视频跟踪
运维·yolo·自动化
余炜yw5 小时前
【LSTM实战】跨越千年,赋诗成文:用LSTM重现唐诗的韵律与情感
人工智能·rnn·深度学习
莫叫石榴姐6 小时前
数据科学与SQL:组距分组分析 | 区间分布问题
大数据·人工智能·sql·深度学习·算法·机器学习·数据挖掘
如若1236 小时前
利用 `OpenCV` 和 `Matplotlib` 库进行图像读取、颜色空间转换、掩膜创建、颜色替换
人工智能·opencv·matplotlib