From 025694f22cadaf9010dbd2c816b095ea07b5766e Mon Sep 17 00:00:00 2001
From: Le0v1n <zjkljd@163.com>
Date: Tue, 21 May 2024 10:45:07 +0800
Subject: [PATCH] Added warning message when --single-cls is used incorrectly

---
 train.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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