How do make edits to detect only person

[uzumakinaruto19]

I want to detect only people class not any other .What changes should i make in the code

[zhiqwang]

Hi @vipinap123 ,

If you don't want to retrain a human-only model, then you need to do a filter on the current output. In the coco dataset, the human label is 0, so you need to add a filter here and only output it if label==0. Such as belows:

import torch
from yolort.models import yolov5s

score_thresh = 0.55
img_h, img_w = 640, 640

model = yolov5s(pretrained=True, score_thresh=score_thresh, size=(img_h, img_w))
model = model.eval()
# Perform inference on an image tensor
out = model.predict("bus.jpg")[0]

keep = torch.where(out["labels"] == 0)
out_person = {}
for k, v in out.items():
    out_person[k] = v[keep]

And then the out_person is all you need here.