概述
本文旨在基于检测到的物品位置信息，分析、计数并提取相关目标。通过对检测结果的坐标数据进行分析，将确定货架的数量以及货架上的物品数量。为此，这里将使用 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 YOLOmodel = 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]) / 2midpoints = 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 += 1objects = 0else:objects += 1group_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 argparseparser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, required=True, help='图像路径')
args = parser.parse_args()image_path = args.image_pathclass ShelfDetector:def __init__(self, model_path: str, confidence: float = 0.45):self.model = YOLO(model_path)self.confidence = confidencedef detect_shelves(self, image_path: str) -> List[int]:result = self.model.predict(source=image_path, conf=self.confidence, save=False)arrxy = result[0].boxes.xyxycoordinates = np.array(arrxy)x_coords = (coordinates[:, 0] + coordinates[:, 2]) / 2y_coords = (coordinates[:, 1] + coordinates[:, 3]) / 2midpoints = 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 = 0for 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 = 0else:objects += 1group_sizes.append(objects + 1)return group_sizesdetector = ShelfDetector('best.pt')
result = detector.detect_shelves(image_path)
for i, size in enumerate(result):print(f"第 {i + 1} 层货架上有 {size} 件商品")
在我们的 CMD 提示符中，输出将如下所示。通过这种方式，我们可以提供不同图像的路径来运行代码。
