Spaces:
Paused
Paused
File size: 1,330 Bytes
8b7211f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
# Ultralytics YOLO 🚀, GPL-3.0 license
import collections
from copy import deepcopy
from .augment import LetterBox
class MixAndRectDataset:
"""A wrapper of multiple images mixed dataset.
Args:
dataset (:obj:`BaseDataset`): The dataset to be mixed.
transforms (Sequence[dict]): config dict to be composed.
"""
def __init__(self, dataset):
self.dataset = dataset
self.imgsz = dataset.imgsz
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
labels = deepcopy(self.dataset[index])
for transform in self.dataset.transforms.tolist():
# mosaic and mixup
if hasattr(transform, "get_indexes"):
indexes = transform.get_indexes(self.dataset)
if not isinstance(indexes, collections.abc.Sequence):
indexes = [indexes]
mix_labels = [deepcopy(self.dataset[index]) for index in indexes]
labels["mix_labels"] = mix_labels
if self.dataset.rect and isinstance(transform, LetterBox):
transform.new_shape = self.dataset.batch_shapes[self.dataset.batch[index]]
labels = transform(labels)
if "mix_labels" in labels:
labels.pop("mix_labels")
return labels
|