PyTorch

Question

我有点卡在某些代码上了。

我只想在检测到一个或多个行人时保存正面图像。当没有检测到任何东西时，什么也不做。

我从阅读开始：https://github.com/ultralytics/yolov5/issues/36

我写了这个：

import torch
import os

f = []
for dirpath, subdirs, files in os.walk('MyFolderWithPictures'):
    for x in files:
        if x.endswith(".jpg"):
            f.append(os.path.join(dirpath, x))


model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

model.conf = 0.25  # NMS confidence threshold
model.iou = 0.45  # NMS IoU threshold
model.classes = 0   # Only pedestrian
model.multi_label = False  # NMS multiple labels per box
model.max_det = 1000  # maximum number of detections per image

img = f  # list of pictures

results = model(img)

results.print()
results.save()

但是这会打印并保存所有图像（正片和负片）。

我只想保存有行人的图片。

你能帮帮我吗？提前致谢。

ps ：输出给：

image 1/13: 1080x1920 1 person
image 2/13: 1080x1920 (no detections)
image 3/13: 1080x1920 (no detections)
image 4/13: 1080x1920 (no detections)
image 5/13: 1080x1920 (no detections)
image 6/13: 1080x1920 (no detections)
image 7/13: 1080x1920 (no detections)
image 8/13: 1080x1920 (no detections)
image 9/13: 1080x1920 (no detections)
image 10/13: 1080x1920 (no detections)
image 11/13: 1080x1920 (no detections)
image 12/13: 1080x1920 1 person
image 13/13: 1080x1920 (no detections)
Speed: 18.6ms pre-process, 119.8ms inference, 1.9ms NMS per image at shape (13, 3, 384, 640)
Saved 13 images to runs\detect\exp

添加解决方案：

for item in f:
    # Images
    img = item  # or file, Path, PIL, OpenCV, numpy, list
    # Inference
    results = model(img)
    # Results
    results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
    if 0 in results.pandas().xyxy[0]['class']:
        results.save()

Answer 1

model(img) 总是会 return 某些类型的结果，即使没有检测到对象。您需要做的是检查结果，看看它是否包含您感兴趣的 class。结果可以很容易地转换为 Pandas Dataframe，以便您可以查询它们。

这是一个示例，用于检查结果是否包含 class 0 的实例，如果包含则保存结果。

if 0 in results.pandas().xyxy[0]['class']:
  results.save()

PyTorch - 仅保存检测到行人的正面图片

PyTorch - save only positive pictures where pedestrian are detected

python

yolov5