Skip to content

Register now for

# Reference for `ultralytics/data/build.py`

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/build.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

## `` ultralytics.data.build.InfiniteDataLoader

    InfiniteDataLoader(*args: Any, **kwargs: Any)

Bases: `DataLoader`

Dataloader that reuses workers for infinite iteration.

This dataloader extends the PyTorch DataLoader to provide infinite recycling of workers, which improves efficiency for training loops that need to iterate through the dataset multiple times without recreating workers.

Attributes:

Name | Type | Description
--- | --- | ---
`batch_sampler` | `_RepeatSampler` | A sampler that repeats indefinitely.
`iterator` | `Iterator` | The iterator from the parent DataLoader.

Methods:

Name | Description
--- | ---
`__len__` | Return the length of the batch sampler's sampler.
`__iter__` | Create a sampler that repeats indefinitely.
`__del__` | Ensure workers are properly terminated.
`reset` | Reset the iterator, useful when modifying dataset settings during training.

Examples:

Create an infinite dataloader for training

    >>> dataset = YOLODataset(...)
    >>> dataloader = InfiniteDataLoader(dataset, batch_size=16, shuffle=True)
    >>> for batch in dataloader:  # Infinite iteration
    >>>     train_step(batch)

Source code in `ultralytics/data/build.py`

    60
    61
    62
    63
    64
    65
    66
    def __init__(self, *args: Any, **kwargs: Any):
        """Initialize the InfiniteDataLoader with the same arguments as DataLoader."""
        if not TORCH_2_0:
            kwargs.pop("prefetch_factor", None)  # not supported by earlier versions
        super().__init__(*args, **kwargs)
        object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))
        self.iterator = super().__iter__()

### `` __del__

    __del__()

Ensure that workers are properly terminated when the dataloader is deleted.

Source code in `ultralytics/data/build.py`

    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    def __del__(self):
        """Ensure that workers are properly terminated when the dataloader is deleted."""
        try:
            if not hasattr(self.iterator, "_workers"):
                return
            for w in self.iterator._workers:  # force terminate
                if w.is_alive():
                    w.terminate()
            self.iterator._shutdown_workers()  # cleanup
        except Exception:
            pass

### `` __iter__

    __iter__() -> Iterator

Create an iterator that yields indefinitely from the underlying iterator.

Source code in `ultralytics/data/build.py`

    72
    73
    74
    75
    def __iter__(self) -> Iterator:
        """Create an iterator that yields indefinitely from the underlying iterator."""
        for _ in range(len(self)):
            yield next(self.iterator)

### `` __len__

    __len__() -> int

Return the length of the batch sampler's sampler.

Source code in `ultralytics/data/build.py`

    68
    69
    70
    def __len__(self) -> int:
        """Return the length of the batch sampler's sampler."""
        return len(self.batch_sampler.sampler)

### `` reset

    reset()

Reset the iterator to allow modifications to the dataset during training.

Source code in `ultralytics/data/build.py`

    89
    90
    91
    def reset(self):
        """Reset the iterator to allow modifications to the dataset during training."""
        self.iterator = self._get_iterator()

## `` ultralytics.data.build._RepeatSampler

    _RepeatSampler(sampler: Any)

Sampler that repeats forever for infinite iteration.

This sampler wraps another sampler and yields its contents indefinitely, allowing for infinite iteration over a dataset without recreating the sampler.

Attributes:

Name | Type | Description
--- | --- | ---
`sampler` | `sampler` | The sampler to repeat.

Source code in `ultralytics/data/build.py`

    105
    106
    107
    def __init__(self, sampler: Any):
        """Initialize the _RepeatSampler with a sampler to repeat indefinitely."""
        self.sampler = sampler

### `` __iter__

    __iter__() -> Iterator

Iterate over the sampler indefinitely, yielding its contents.

Source code in `ultralytics/data/build.py`

    109
    110
    111
    112
    def __iter__(self) -> Iterator:
        """Iterate over the sampler indefinitely, yielding its contents."""
        while True:
            yield from iter(self.sampler)

## `` ultralytics.data.build.seed_worker

    seed_worker(worker_id: int)

Set dataloader worker seed for reproducibility across worker processes.

Source code in `ultralytics/data/build.py`

    115
    116
    117
    118
    119
    def seed_worker(worker_id: int):  # noqa
        """Set dataloader worker seed for reproducibility across worker processes."""
        worker_seed = torch.initial_seed() % 2**32
        np.random.seed(worker_seed)
        random.seed(worker_seed)

## `` ultralytics.data.build.build_yolo_dataset

    build_yolo_dataset(
        cfg: IterableSimpleNamespace,
        img_path: str,
        batch: int,
        data: dict[str, Any],
        mode: str = "train",
        rect: bool = False,
        stride: int = 32,
        multi_modal: bool = False,
    )

Build and return a YOLO dataset based on configuration parameters.

Source code in `ultralytics/data/build.py`

    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    def build_yolo_dataset(
        cfg: IterableSimpleNamespace,
        img_path: str,
        batch: int,
        data: dict[str, Any],
        mode: str = "train",
        rect: bool = False,
        stride: int = 32,
        multi_modal: bool = False,
    ):
        """Build and return a YOLO dataset based on configuration parameters."""
        dataset = YOLOMultiModalDataset if multi_modal else YOLODataset
        return dataset(
            img_path=img_path,
            imgsz=cfg.imgsz,
            batch_size=batch,
            augment=mode == "train",  # augmentation
            hyp=cfg,  # TODO: probably add a get_hyps_from_cfg function
            rect=cfg.rect or rect,  # rectangular batches
            cache=cfg.cache or None,
            single_cls=cfg.single_cls or False,
            stride=stride,
            pad=0.0 if mode == "train" else 0.5,
            prefix=colorstr(f"{mode}: "),
            task=cfg.task,
            classes=cfg.classes,
            data=data,
            fraction=cfg.fraction if mode == "train" else 1.0,
        )

## `` ultralytics.data.build.build_grounding

    build_grounding(
        cfg: IterableSimpleNamespace,
        img_path: str,
        json_file: str,
        batch: int,
        mode: str = "train",
        rect: bool = False,
        stride: int = 32,
        max_samples: int = 80,
    )

Build and return a GroundingDataset based on configuration parameters.

Source code in `ultralytics/data/build.py`

    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    def build_grounding(
        cfg: IterableSimpleNamespace,
        img_path: str,
        json_file: str,
        batch: int,
        mode: str = "train",
        rect: bool = False,
        stride: int = 32,
        max_samples: int = 80,
    ):
        """Build and return a GroundingDataset based on configuration parameters."""
        return GroundingDataset(
            img_path=img_path,
            json_file=json_file,
            max_samples=max_samples,
            imgsz=cfg.imgsz,
            batch_size=batch,
            augment=mode == "train",  # augmentation
            hyp=cfg,  # TODO: probably add a get_hyps_from_cfg function
            rect=cfg.rect or rect,  # rectangular batches
            cache=cfg.cache or None,
            single_cls=cfg.single_cls or False,
            stride=stride,
            pad=0.0 if mode == "train" else 0.5,
            prefix=colorstr(f"{mode}: "),
            task=cfg.task,
            classes=cfg.classes,
            fraction=cfg.fraction if mode == "train" else 1.0,
        )

## `` ultralytics.data.build.build_dataloader

    build_dataloader(
        dataset,
        batch: int,
        workers: int,
        shuffle: bool = True,
        rank: int = -1,
        drop_last: bool = False,
    )

Create and return an InfiniteDataLoader or DataLoader for training or validation.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`dataset` | `Dataset` | Dataset to load data from. | required
`batch` | `int` | Batch size for the dataloader. | required
`workers` | `int` | Number of worker threads for loading data. | required
`shuffle` | `bool` | Whether to shuffle the dataset. | `True`
`rank` | `int` | Process rank in distributed training. -1 for single-GPU training. | `-1`
`drop_last` | `bool` | Whether to drop the last incomplete batch. | `False`

Returns:

Type | Description
--- | ---
`InfiniteDataLoader` | A dataloader that can be used for training or validation.

Examples:

Create a dataloader for training

    >>> dataset = YOLODataset(...)
    >>> dataloader = build_dataloader(dataset, batch=16, workers=4, shuffle=True)

Source code in `ultralytics/data/build.py`

    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    def build_dataloader(dataset, batch: int, workers: int, shuffle: bool = True, rank: int = -1, drop_last: bool = False):
        """
        Create and return an InfiniteDataLoader or DataLoader for training or validation.

        Args:
            dataset (Dataset): Dataset to load data from.
            batch (int): Batch size for the dataloader.
            workers (int): Number of worker threads for loading data.
            shuffle (bool, optional): Whether to shuffle the dataset.
            rank (int, optional): Process rank in distributed training. -1 for single-GPU training.
            drop_last (bool, optional): Whether to drop the last incomplete batch.

        Returns:
            (InfiniteDataLoader): A dataloader that can be used for training or validation.

        Examples:
            Create a dataloader for training
            >>> dataset = YOLODataset(...)
            >>> dataloader = build_dataloader(dataset, batch=16, workers=4, shuffle=True)
        """
        batch = min(batch, len(dataset))
        nd = torch.cuda.device_count()  # number of CUDA devices
        nw = min(os.cpu_count() // max(nd, 1), workers)  # number of workers
        sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
        generator = torch.Generator()
        generator.manual_seed(6148914691236517205 + RANK)
        return InfiniteDataLoader(
            dataset=dataset,
            batch_size=batch,
            shuffle=shuffle and sampler is None,
            num_workers=nw,
            sampler=sampler,
            prefetch_factor=4 if nw > 0 else None,  # increase over default 2
            pin_memory=nd > 0,
            collate_fn=getattr(dataset, "collate_fn", None),
            worker_init_fn=seed_worker,
            generator=generator,
            drop_last=drop_last and len(dataset) % batch != 0,
        )

## `` ultralytics.data.build.check_source

    check_source(source)

Check the type of input source and return corresponding flag values.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`source` | `str | int | Path | list | tuple | ndarray | Image | Tensor` | The input source to check. | required

Returns:

Name | Type | Description
--- | --- | ---
`source` | `str | int | Path | list | tuple | ndarray | Image | Tensor` | The processed source.
`webcam` | `bool` | Whether the source is a webcam.
`screenshot` | `bool` | Whether the source is a screenshot.
`from_img` | `bool` | Whether the source is an image or list of images.
`in_memory` | `bool` | Whether the source is an in-memory object.
`tensor` | `bool` | Whether the source is a torch.Tensor.

Examples:

Check a file path source

    >>> source, webcam, screenshot, from_img, in_memory, tensor = check_source("image.jpg")

Check a webcam source

    >>> source, webcam, screenshot, from_img, in_memory, tensor = check_source(0)

Source code in `ultralytics/data/build.py`

    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    def check_source(source):
        """
        Check the type of input source and return corresponding flag values.

        Args:
            source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The input source to check.

        Returns:
            source (str | int | Path | list | tuple | np.ndarray | PIL.Image | torch.Tensor): The processed source.
            webcam (bool): Whether the source is a webcam.
            screenshot (bool): Whether the source is a screenshot.
            from_img (bool): Whether the source is an image or list of images.
            in_memory (bool): Whether the source is an in-memory object.
            tensor (bool): Whether the source is a torch.Tensor.

        Examples:
            Check a file path source
            >>> source, webcam, screenshot, from_img, in_memory, tensor = check_source("image.jpg")

            Check a webcam source
            >>> source, webcam, screenshot, from_img, in_memory, tensor = check_source(0)
        """
        webcam, screenshot, from_img, in_memory, tensor = False, False, False, False, False
        if isinstance(source, (str, int, Path)):  # int for local usb camera
            source = str(source)
            source_lower = source.lower()
            is_url = source_lower.startswith(("https://", "http://", "rtsp://", "rtmp://", "tcp://"))
            is_file = (urlsplit(source_lower).path if is_url else source_lower).rpartition(".")[-1] in (
                IMG_FORMATS | VID_FORMATS
            )
            webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
            screenshot = source_lower == "screen"
            if is_url and is_file:
                source = check_file(source)  # download
        elif isinstance(source, LOADERS):
            in_memory = True
        elif isinstance(source, (list, tuple)):
            source = autocast_list(source)  # convert all list elements to PIL or np arrays
            from_img = True
        elif isinstance(source, (Image.Image, np.ndarray)):
            from_img = True
        elif isinstance(source, torch.Tensor):
            tensor = True
        else:
            raise TypeError("Unsupported image type. For supported types see https://docs.ultralytics.com/modes/predict")

        return source, webcam, screenshot, from_img, in_memory, tensor

## `` ultralytics.data.build.load_inference_source

    load_inference_source(
        source=None,
        batch: int = 1,
        vid_stride: int = 1,
        buffer: bool = False,
        channels: int = 3,
    )

Load an inference source for object detection and apply necessary transformations.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`source` | `str | Path | Tensor | Image | ndarray` | The input source for inference. | `None`
`batch` | `int` | Batch size for dataloaders. | `1`
`vid_stride` | `int` | The frame interval for video sources. | `1`
`buffer` | `bool` | Whether stream frames will be buffered. | `False`
`channels` | `int` | The number of input channels for the model. | `3`

Returns:

Type | Description
--- | ---
`Dataset` | A dataset object for the specified input source with attached source_type attribute.

Examples:

Load an image source for inference

    >>> dataset = load_inference_source("image.jpg", batch=1)

Load a video stream source

    >>> dataset = load_inference_source("rtsp://example.com/stream", vid_stride=2)

Source code in `ultralytics/data/build.py`

    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    300
    301
    302
    303
    304
    305
    306
    307
    308
    309
    310
    311
    312
    313
    314
    315
    def load_inference_source(source=None, batch: int = 1, vid_stride: int = 1, buffer: bool = False, channels: int = 3):
        """
        Load an inference source for object detection and apply necessary transformations.

        Args:
            source (str | Path | torch.Tensor | PIL.Image | np.ndarray, optional): The input source for inference.
            batch (int, optional): Batch size for dataloaders.
            vid_stride (int, optional): The frame interval for video sources.
            buffer (bool, optional): Whether stream frames will be buffered.
            channels (int, optional): The number of input channels for the model.

        Returns:
            (Dataset): A dataset object for the specified input source with attached source_type attribute.

        Examples:
            Load an image source for inference
            >>> dataset = load_inference_source("image.jpg", batch=1)

            Load a video stream source
            >>> dataset = load_inference_source("rtsp://example.com/stream", vid_stride=2)
        """
        source, stream, screenshot, from_img, in_memory, tensor = check_source(source)
        source_type = source.source_type if in_memory else SourceTypes(stream, screenshot, from_img, tensor)

        # Dataloader
        if tensor:
            dataset = LoadTensor(source)
        elif in_memory:
            dataset = source
        elif stream:
            dataset = LoadStreams(source, vid_stride=vid_stride, buffer=buffer, channels=channels)
        elif screenshot:
            dataset = LoadScreenshots(source, channels=channels)
        elif from_img:
            dataset = LoadPilAndNumpy(source, channels=channels)
        else:
            dataset = LoadImagesAndVideos(source, batch=batch, vid_stride=vid_stride, channels=channels)

        # Attach source types to the dataset
        setattr(dataset, "source_type", source_type)

        return dataset

📅 Created 1 year ago ✏️ Updated 1 year ago