基于YOLOv8与SKU110K数据集实现超市货架物品目标检测与计算

概述

本文旨在基于检测到的物品位置信息，分析、计数并提取相关目标。通过对检测结果的坐标数据进行分析，将确定货架的数量以及货架上的物品数量。为此，这里将使用 SKU110K 数据集来训练目标检测模型。该数据集包含商店货架上物品的边界框标注，仅包含一个名为“物品”的类别。

SKU110K数据集

在 SKU-110K 是专注于密集场景下的目标检测问题。此类场景中的图像包含大量外观相似甚至相同的物体，且物体位置紧密相邻。这些场景通常是人造的，例如零售货架展示、交通和城市景观图像。

数据集 SKU-110K 中的一张典型图像，展示了密集排列的物体

算法实现

创建环境与导入相关库，使用预训练模型进行预测，并将结果赋值给一个变量。

conda create -n yolov8 python=3.8
activate ylolv8
conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.7 -c pytorch -c nvidia
pip install ultralytics

测试环境：

import numpy as np
from ultralytics import YOLO

model = YOLO('best.pt')
result = model.predict(
   source='test_88.jpg',
   conf=0.45,
   save=True
)

模型训练

关于模型训练，可以看之前关于YOLOv8如果训练自定义数据

提取边界框坐标并排序

提取边界框的 xyxy 格式坐标，并将其转换为 NumPy 数组。这些数据分别代表边界框的 xmin、ymin、xmax 和 ymax 坐标。以下是前 25 个输出：

arrxy = result[0].boxes.xyxy
coordinates = np.array(arrxy)
coordinates[:25]

array([[       2082,        1426,        2318,        1635],
       [       2356,        1106,        2678,        1321],
       [       1927,        2442,        2284,        2799],
       [        647,         961,         865,        1149],
       [       2101,        1644,        2323,        1841],
       [       1565,        2472,        1913,        2822],
       [       2334,        1420,        2567,        1640],
       [       1094,         957,        1301,        1138],
       [        967,        3186,        1243,        3369],
       [        873,         956,        1087,        1148],
       [        739,        3466,         993,        3662],
       [       1318,         968,        1512,        1138],
       [       1528,        2948,        1782,        3159],
       [       2109,        1914,        2497,        2128],
       [       1329,        2482,        1559,        2662],
       [       1264,        2947,        1522,        3156],
       [        466,        3655,         738,        3847],
       [       2139,        2135,        2527,        2330],
       [       1526,        1129,        1764,        1330],
       [        468,        3469,         731,        3646],
       [        691,        2976,         963,        3181],
       [       2233,         384,        2458,         582],
       [       1256,        3161,        1509,        3370],
       [        426,           0,         913,         154],
       [        975,        2964,        1252,        3179]], dtype=float32)

由于图像的 (0,0) 坐标代表图像的左上角，因此需要相应地对它们进行排序。为此，计算 x 和 y 坐标的中点，并根据 y 坐标进行排序。

arrxy = result[0].boxes.xyxy
coordinates = np.array(arrxy)

x_coords = (coordinates[:, 0] + coordinates[:, 2]) / 2
y_coords = (coordinates[:, 1] + coordinates[:, 3]) / 2

midpoints = np.column_stack((x_coords, y_coords))

rounded_n_sorted_arr = np.round(midpoints[midpoints[:, 1].argsort()]).astype(int)

print(rounded_n_sorted_arr[:25])

[[2762   63]
 [2463   66]
 [1998   68]
 [ 670   77]
 [ 241   80]
 [1547  370]
 [ 978  378]
 [1370  399]
 [2088  416]
 [2102  476]
 [ 916  478]
 [2346  483]
 [ 363  504]
 [2774  514]
 [1842  527]
 [1392  542]
 [2559  544]
 [1178  552]
 [1599  554]
 [ 652  579]
 [ 916  662]
 [3009  677]
 [2778  678]
 [2122  684]
 [2364  688]]

使用 OpenCV 可视化数据

使用 OpenCV 分析物品之间的坐标关系，并通过将物品中心的坐标用深蓝色圆圈标出，直观地理解这些关系。从排序后的坐标可以看出，y 轴的显著增加表明物品已经跳到了货架上。

例如，y 轴从 80 增加到 370，变化非常大，因此可以断定这是一个货架。

计数物品和货架

从分析中可以看出，如果 y 轴的增加超过某个值，就可以认为是一个货架。在编写代码时，有机会在不同图像上进行测试，可以说如果图像捕捉到了所有货架并且是平行拍摄的，那么 y 轴的跳跃值超过 130 就可以认为是一个货架。

count = 1
objects = 0
group_sizes = []

for i in range(1, len(rounded_n_sorted_arr)):
    if rounded_n_sorted_arr[i][1] - rounded_n_sorted_arr[i - 1][1] > 130:
        group_sizes.append(objects + 1)
        count += 1
        objects = 0
    else:
        objects += 1

group_sizes.append(objects + 1)

for i, size in enumerate(group_sizes):
    print(f"第 {i + 1} 层货架上有 {size} 件商品")

转换为参数化程序

将所有代码整合在一起。使用 argparse 库，可以通过以下代码提供输入文件路径。作为输出，将打印出每个货架上有多少件商品的数据。

from typing import List
import numpy as np
from ultralytics import YOLO
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, required=True, help='图像路径')
args = parser.parse_args()

image_path = args.image_path

class ShelfDetector:
    def __init__(self, model_path: str, confidence: float = 0.45):
        self.model = YOLO(model_path)
        self.confidence = confidence

    def detect_shelves(self, image_path: str) -> List[int]:
        result = self.model.predict(source=image_path, conf=self.confidence, save=False)
        arrxy = result[0].boxes.xyxy
        coordinates = np.array(arrxy)

        x_coords = (coordinates[:, 0] + coordinates[:, 2]) / 2
        y_coords = (coordinates[:, 1] + coordinates[:, 3]) / 2
        midpoints = np.column_stack((x_coords, y_coords))

        sorted_midpoints = midpoints[midpoints[:, 1].argsort()]
        rounded_n_sorted_arr = np.round(sorted_midpoints).astype(int)

        group_sizes = []
        objects = 0
        for i in range(1, len(rounded_n_sorted_arr)):
            if rounded_n_sorted_arr[i][1] - rounded_n_sorted_arr[i - 1][1] > 130:
                group_sizes.append(objects + 1)
                objects = 0
            else:
                objects += 1

        group_sizes.append(objects + 1)
        return group_sizes

detector = ShelfDetector('best.pt')
result = detector.detect_shelves(image_path)
for i, size in enumerate(result):
    print(f"第 {i + 1} 层货架上有 {size} 件商品")

在我们的 CMD 提示符中，输出将如下所示。通过这种方式，我们可以提供不同图像的路径来运行代码。