PyTorch - save only positive pictures where pedestrian are detected

Question

I'm a bit stuck on some code.

I want to save ONLY POSITIVE IMAGE when one or more pedestrian are detected. When nothing is detected, do nothing.

I started by reading : https://github.com/ultralytics/yolov5/issues/36

I wrote this :

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()

But this print and save ALL images (positive and negative).

I want to save only images with pedestrian.

Can you help me ? Thanks in advance.

ps : the output give :

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

ADD SOLUTION :

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) will always return some kinds of results, even if there are no objects detected. What you need to do is inspect the results and see if it includes the class that you are interested in. The results can easily be converted to a Pandas Dataframe so that you can query them.

Here is an example to check if the results contain an instance of class 0 and then save the results if it does.

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