detect.py --classes IndexError

[VvvvvGH]

🐛 Bug

Traceback (most recent call last):
  File "detect.py", line 172, in <module>
    detect()
  File "detect.py", line 79, in detect
    fast=True, classes=opt.classes, agnostic=opt.agnostic_nms)
  File "D:\yolov5-master\utils\utils.py", line 511, in non_max_suppression
    x = x[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)]
IndexError: The shape of the mask [21] at index 0does not match the shape of the indexed tensor [20, 6] at index 0

To Reproduce

REQUIRED: Code to reproduce your issue below

 python detect.py --source 0 --weights yolov5x.pt --view-img --classes 0

Expected behavior

No error produce

Environment

If applicable, add screenshots to help explain your problem.

• OS: Windows
• GPU 1080

Additional context

yolov5/utils/utils.py

Lines 501 to 512 in 7c2832c

	# Detections matrix nx6 (xyxy, conf, cls)
	if multi_label:
	i, j = (x[:, 5:] > conf_thres).nonzero().t()
	x = torch.cat((box[i], x[i, j + 5].unsqueeze(1), j.float().unsqueeze(1)), 1)
	else: # best class only
	conf, j = x[:, 5:].max(1)
	x = torch.cat((box, conf.unsqueeze(1), j.float().unsqueeze(1)), 1)[conf > conf_thres]
	# Filter by class
	if classes:
	x = x[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)]

Possible fix

Change

 if classes: 
     x = x[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)] 

To

 if classes: 
     x = x[(x[:, 5:].view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)] 

[github-actions]

Hello @VvvvvGH, thank you for your interest in our work! Please visit our Custom Training Tutorial to get started, and see our Google Colab Notebook, Docker Image, and GCP Quickstart Guide for example environments.

If this is a bug report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom model or data training question, please note that Ultralytics does not provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as:

• Cloud-based AI surveillance systems operating on hundreds of HD video streams in realtime.
• Edge AI integrated into custom iOS and Android apps for realtime 30 FPS video inference.
• Custom data training, hyperparameter evolution, and model exportation to any destination.

For more information please visit https://www.ultralytics.com.

[glenn-jocher]

@VvvvvGH thank you for your bug report! We can successfully reproduce this following your example. I will look into this.

[glenn-jocher]

@VvvvvGH yes you are correct, x and j are becoming different shapes causing the error. Your solution works well. I will push a fix shortly.

[Tamme]

Same fix needs to be applied in yolov3 repo also

[glenn-jocher]

@Tamme submit a PR on the relevant repo please.