From 6dd17516c8610fca8a75b8c003866c1bbe921daa Mon Sep 17 00:00:00 2001 From: Eljas Hyyrynen Date: Thu, 23 Mar 2023 21:29:11 +0200 Subject: [PATCH] dataloaders: fix class filtering for segmentation (#11171) * dataloaders: fix class filtering for segmentation self.segments[i] and segment[j] are lists so they cannot be indexed with booleans self.segments is a tuple so it has to be converted into a list first * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- utils/dataloaders.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/utils/dataloaders.py b/utils/dataloaders.py index 7687a2ba2..28d5b7974 100644 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -531,13 +531,14 @@ class LoadImagesAndLabels(Dataset): # Update labels include_class = [] # filter labels to include only these classes (optional) + self.segments = list(self.segments) include_class_array = np.array(include_class).reshape(1, -1) for i, (label, segment) in enumerate(zip(self.labels, self.segments)): if include_class: j = (label[:, 0:1] == include_class_array).any(1) self.labels[i] = label[j] if segment: - self.segments[i] = segment[j] + self.segments[i] = [segment[idx] for idx, elem in enumerate(j) if elem] if single_cls: # single-class training, merge all classes into 0 self.labels[i][:, 0] = 0