From c8a98cb7cbbf5d05abb9b134ada0c75d0dc62a6c Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 19 Jul 2021 13:10:21 +0200
Subject: [PATCH] Missing `nc` and `names` handling in check_dataset() (#4066)

---
 utils/general.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/utils/general.py b/utils/general.py
index 846c1464..08a3ff65 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -231,6 +231,9 @@ def check_dataset(data, autodownload=True):
             if data.get(k):  # prepend path
                 data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
 
+    assert 'nc' in data, "Dataset 'nc' key missing."
+    if 'names' not in data:
+        data['names'] = [str(i) for i in range(data['nc'])]  # assign class names if missing
     train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')]
     if val:
         val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])]  # val path
-- 
GitLab