# Ultralytics YOLO 🚀, GPL-3.0 license import os import random from pathlib import Path import numpy as np import torch from PIL import Image from torch.utils.data import DataLoader, dataloader, distributed from ultralytics.yolo.data.dataloaders.stream_loaders import (LOADERS, LoadImages, LoadPilAndNumpy, LoadScreenshots, LoadStreams, SourceTypes, autocast_list) from ultralytics.yolo.data.utils import IMG_FORMATS, VID_FORMATS from ultralytics.yolo.utils.checks import check_file from ..utils import LOGGER, colorstr from ..utils.torch_utils import torch_distributed_zero_first from .dataset import ClassificationDataset, YOLODataset from .utils import PIN_MEMORY, RANK class InfiniteDataLoader(dataloader.DataLoader): """Dataloader that reuses workers Uses same syntax as vanilla DataLoader """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler)) self.iterator = super().__iter__() def __len__(self): return len(self.batch_sampler.sampler) def __iter__(self): for _ in range(len(self)): yield next(self.iterator) class _RepeatSampler: """Sampler that repeats forever Args: sampler (Sampler) """ def __init__(self, sampler): self.sampler = sampler def __iter__(self): while True: yield from iter(self.sampler) def seed_worker(worker_id): # Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader worker_seed = torch.initial_seed() % 2 ** 32 np.random.seed(worker_seed) random.seed(worker_seed) def build_dataloader(cfg, batch, img_path, stride=32, rect=False, names=None, rank=-1, mode='train'): assert mode in ['train', 'val'] shuffle = mode == 'train' if cfg.rect and shuffle: LOGGER.warning("WARNING ⚠️ 'rect=True' is incompatible with DataLoader shuffle, setting shuffle=False") shuffle = False with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP dataset = YOLODataset( 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=int(stride), pad=0.0 if mode == 'train' else 0.5, prefix=colorstr(f'{mode}: '), use_segments=cfg.task == 'segment', use_keypoints=cfg.task == 'keypoint', names=names) batch = min(batch, len(dataset)) nd = torch.cuda.device_count() # number of CUDA devices workers = cfg.workers if mode == 'train' else cfg.workers * 2 nw = min([os.cpu_count() // max(nd, 1), batch if batch > 1 else 0, workers]) # number of workers sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) loader = DataLoader if cfg.image_weights or cfg.close_mosaic else InfiniteDataLoader # allow attribute updates generator = torch.Generator() generator.manual_seed(6148914691236517205 + RANK) return loader(dataset=dataset, batch_size=batch, shuffle=shuffle and sampler is None, num_workers=nw, sampler=sampler, pin_memory=PIN_MEMORY, collate_fn=getattr(dataset, 'collate_fn', None), worker_init_fn=seed_worker, generator=generator), dataset # build classification # TODO: using cfg like `build_dataloader` def build_classification_dataloader(path, imgsz=224, batch_size=16, augment=True, cache=False, rank=-1, workers=8, shuffle=True): # Returns Dataloader object to be used with YOLOv5 Classifier with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache) batch_size = min(batch_size, len(dataset)) nd = torch.cuda.device_count() nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) generator = torch.Generator() generator.manual_seed(6148914691236517205 + RANK) return InfiniteDataLoader(dataset, batch_size=batch_size, shuffle=shuffle and sampler is None, num_workers=nw, sampler=sampler, pin_memory=PIN_MEMORY, worker_init_fn=seed_worker, generator=generator) # or DataLoader(persistent_workers=True) def check_source(source): webcam, screenshot, from_img, in_memory = False, False, False, False if isinstance(source, (str, int, Path)): # int for local usb carame source = str(source) is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) is_url = source.lower().startswith(('https://', 'http://', 'rtsp://', 'rtmp://')) webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file) screenshot = source.lower().startswith('screen') if is_url and is_file: source = check_file(source) # download elif isinstance(source, tuple(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 else: raise Exception( 'Unsupported type encountered! See docs for supported types https://docs.ultralytics.com/predict') return source, webcam, screenshot, from_img, in_memory def load_inference_source(source=None, transforms=None, imgsz=640, vid_stride=1, stride=32, auto=True): """ TODO: docs """ # source source, webcam, screenshot, from_img, in_memory = check_source(source) source_type = source.source_type if in_memory else SourceTypes(webcam, screenshot, from_img) # Dataloader if in_memory: dataset = source elif webcam: dataset = LoadStreams(source, imgsz=imgsz, stride=stride, auto=auto, transforms=transforms, vid_stride=vid_stride) elif screenshot: dataset = LoadScreenshots(source, imgsz=imgsz, stride=stride, auto=auto, transforms=transforms) elif from_img: dataset = LoadPilAndNumpy(source, imgsz=imgsz, stride=stride, auto=auto, transforms=transforms) else: dataset = LoadImages(source, imgsz=imgsz, stride=stride, auto=auto, transforms=transforms, vid_stride=vid_stride) setattr(dataset, 'source_type', source_type) # attach source types return dataset