diff --git a/train.py b/train.py index df0972a67c70..35e5fd6ccdf0 100644 --- a/train.py +++ b/train.py @@ -176,7 +176,14 @@ def train(hyp, opt, device, callbacks): data_dict = data_dict or check_dataset(data) # check if None train_path, val_path = data_dict["train"], data_dict["val"] nc = 1 if single_cls else int(data_dict["nc"]) # number of classes - names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class names + if single_cls and len(data_dict["names"]) != 1: + LOGGER.warning( + "WARNING ⚠️ Please check the dataset to ensure that when --single-cls is enabled, " + "the number of classes in the dataset is 1." + ) + names = {0: "item"} # combine all classes into a single class + else: + names = data_dict["names"] # class names is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset # Model