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

news2024/11/23 12:55:07

项目背景

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

技术栈

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

项目结构

数据准备：
- 收集带有标注的船只图像数据集。
- 将数据集划分为训练集和测试集。
模型训练：
- 使用YOLOv8框架训练模型。
- 调整超参数以优化检测和分割性能。
模型转换：
- 将训练好的PyTorch模型导出为ONNX格式。
- 验证ONNX模型的正确性。
推理部署：
- 编写推理代码，支持从图像或视频流中检测并分割船只。
- 使用ONNX Runtime进行高效推理。
结果展示：
- 可视化检测结果，包括边界框和分割掩码。
- 计算并报告性能指标如准确率、召回率等。

示例代码

一个简化的示例代码片段，用于演示如何加载一个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()