Fix `__len__` and `single_cls` (#249)

pull/254/head
Laughing 2 years ago committed by GitHub
parent 88c9418087
commit 53fa0f4550
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 4
      ultralytics/yolo/data/base.py

@ -110,7 +110,7 @@ class BaseDataset(Dataset):
if segments:
self.labels[i]["segments"] = segments[j]
if self.single_cls:
self.labels[i]["cls"] = 0
self.labels[i]["cls"][:, 0] = 0
def load_image(self, i):
# Loads 1 image from dataset index 'i', returns (im, resized hw)
@ -191,7 +191,7 @@ class BaseDataset(Dataset):
return label
def __len__(self):
return len(self.im_files)
return len(self.labels)
def update_labels_info(self, label):
"""custom your label format here"""

Loading…
Cancel
Save