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

single_channel
Laughing 2 years ago committed by GitHub
parent 88c9418087
commit 53fa0f4550
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

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

Loading…
Cancel
Save