Update AutoShape.forward() model.classes example (ultralytics#5324)

fhkiel-mlaip · Oct 25, 2021 · 424934d · 424934d
1 parent 20c74db
commit 424934d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/common.py b/models/common.py
@@ -277,7 +277,7 @@ class AutoShape(nn.Module):
     # YOLOv5 input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
     conf = 0.25  # NMS confidence threshold
     iou = 0.45  # NMS IoU threshold
-    classes = None  # (optional list) filter by class
+    classes = None  # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
     multi_label = False  # NMS multiple labels per box
     max_det = 1000  # maximum number of detections per image