# Ultralytics YOLO 🚀, GPL-3.0 license from itertools import repeat from multiprocessing.pool import ThreadPool from pathlib import Path import torchvision from tqdm import tqdm from ..utils import NUM_THREADS, TQDM_BAR_FORMAT, is_dir_writeable from .augment import * from .base import BaseDataset from .utils import HELP_URL, LOCAL_RANK, get_hash, img2label_paths, verify_image_label class YOLODataset(BaseDataset): cache_version = '1.0.1' # dataset labels *.cache version, >= 1.0.0 for YOLOv8 rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4] """ Dataset class for loading images object detection and/or segmentation labels in YOLO format. Args: img_path (str): path to the folder containing images. imgsz (int): image size (default: 640). cache (bool): if True, a cache file of the labels is created to speed up future creation of dataset instances (default: False). augment (bool): if True, data augmentation is applied (default: True). hyp (dict): hyperparameters to apply data augmentation (default: None). prefix (str): prefix to print in log messages (default: ''). rect (bool): if True, rectangular training is used (default: False). batch_size (int): size of batches (default: None). stride (int): stride (default: 32). pad (float): padding (default: 0.0). single_cls (bool): if True, single class training is used (default: False). use_segments (bool): if True, segmentation masks are used as labels (default: False). use_keypoints (bool): if True, keypoints are used as labels (default: False). names (list): class names (default: None). Returns: A PyTorch dataset object that can be used for training an object detection or segmentation model. """ def __init__(self, img_path, imgsz=640, cache=False, augment=True, hyp=None, prefix='', rect=False, batch_size=None, stride=32, pad=0.0, single_cls=False, use_segments=False, use_keypoints=False, names=None): self.use_segments = use_segments self.use_keypoints = use_keypoints self.names = names assert not (self.use_segments and self.use_keypoints), 'Can not use both segments and keypoints.' super().__init__(img_path, imgsz, cache, augment, hyp, prefix, rect, batch_size, stride, pad, single_cls) def cache_labels(self, path=Path('./labels.cache')): """Cache dataset labels, check images and read shapes. Args: path (Path): path where to save the cache file (default: Path('./labels.cache')). Returns: (dict): labels. """ x = {'labels': []} nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages desc = f'{self.prefix}Scanning {path.parent / path.stem}...' total = len(self.im_files) with ThreadPool(NUM_THREADS) as pool: results = pool.imap(func=verify_image_label, iterable=zip(self.im_files, self.label_files, repeat(self.prefix), repeat(self.use_keypoints), repeat(len(self.names)))) pbar = tqdm(results, desc=desc, total=total, bar_format=TQDM_BAR_FORMAT) for im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg in pbar: nm += nm_f nf += nf_f ne += ne_f nc += nc_f if im_file: x['labels'].append( dict( im_file=im_file, shape=shape, cls=lb[:, 0:1], # n, 1 bboxes=lb[:, 1:], # n, 4 segments=segments, keypoints=keypoint, normalized=True, bbox_format='xywh')) if msg: msgs.append(msg) pbar.desc = f'{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt' pbar.close() if msgs: LOGGER.info('\n'.join(msgs)) if nf == 0: LOGGER.warning(f'{self.prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}') x['hash'] = get_hash(self.label_files + self.im_files) x['results'] = nf, nm, ne, nc, len(self.im_files) x['msgs'] = msgs # warnings x['version'] = self.cache_version # cache version if is_dir_writeable(path.parent): if path.exists(): path.unlink() # remove *.cache file if exists np.save(str(path), x) # save cache for next time path.with_suffix('.cache.npy').rename(path) # remove .npy suffix LOGGER.info(f'{self.prefix}New cache created: {path}') else: LOGGER.warning(f'{self.prefix}WARNING ⚠️ Cache directory {path.parent} is not writeable, cache not saved.') return x def get_labels(self): self.label_files = img2label_paths(self.im_files) cache_path = Path(self.label_files[0]).parent.with_suffix('.cache') try: cache, exists = np.load(str(cache_path), allow_pickle=True).item(), True # load dict assert cache['version'] == self.cache_version # matches current version assert cache['hash'] == get_hash(self.label_files + self.im_files) # identical hash except (FileNotFoundError, AssertionError, AttributeError): cache, exists = self.cache_labels(cache_path), False # run cache ops # Display cache nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupt, total if exists and LOCAL_RANK in {-1, 0}: d = f'Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt' tqdm(None, desc=self.prefix + d, total=n, initial=n, bar_format=TQDM_BAR_FORMAT) # display cache results if cache['msgs']: LOGGER.info('\n'.join(cache['msgs'])) # display warnings if nf == 0: # number of labels found raise FileNotFoundError(f'{self.prefix}No labels found in {cache_path}, can not start training. {HELP_URL}') # Read cache [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items labels = cache['labels'] self.im_files = [lb['im_file'] for lb in labels] # update im_files # Check if the dataset is all boxes or all segments lengths = ((len(lb['cls']), len(lb['bboxes']), len(lb['segments'])) for lb in labels) len_cls, len_boxes, len_segments = (sum(x) for x in zip(*lengths)) if len_segments and len_boxes != len_segments: LOGGER.warning( f'WARNING ⚠️ Box and segment counts should be equal, but got len(segments) = {len_segments}, ' f'len(boxes) = {len_boxes}. To resolve this only boxes will be used and all segments will be removed. ' 'To avoid this please supply either a detect or segment dataset, not a detect-segment mixed dataset.') for lb in labels: lb['segments'] = [] if len_cls == 0: raise ValueError(f'All labels empty in {cache_path}, can not start training without labels. {HELP_URL}') return labels # TODO: use hyp config to set all these augmentations def build_transforms(self, hyp=None): if self.augment: hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0 hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0 transforms = v8_transforms(self, self.imgsz, hyp) else: transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)]) transforms.append( Format(bbox_format='xywh', normalize=True, return_mask=self.use_segments, return_keypoint=self.use_keypoints, batch_idx=True, mask_ratio=hyp.mask_ratio, mask_overlap=hyp.overlap_mask)) return transforms def close_mosaic(self, hyp): hyp.mosaic = 0.0 # set mosaic ratio=0.0 hyp.copy_paste = 0.0 # keep the same behavior as previous v8 close-mosaic hyp.mixup = 0.0 # keep the same behavior as previous v8 close-mosaic self.transforms = self.build_transforms(hyp) def update_labels_info(self, label): """custom your label format here""" # NOTE: cls is not with bboxes now, classification and semantic segmentation need an independent cls label # we can make it also support classification and semantic segmentation by add or remove some dict keys there. bboxes = label.pop('bboxes') segments = label.pop('segments') keypoints = label.pop('keypoints', None) bbox_format = label.pop('bbox_format') normalized = label.pop('normalized') label['instances'] = Instances(bboxes, segments, keypoints, bbox_format=bbox_format, normalized=normalized) return label @staticmethod def collate_fn(batch): new_batch = {} keys = batch[0].keys() values = list(zip(*[list(b.values()) for b in batch])) for i, k in enumerate(keys): value = values[i] if k == 'img': value = torch.stack(value, 0) if k in ['masks', 'keypoints', 'bboxes', 'cls']: value = torch.cat(value, 0) new_batch[k] = value new_batch['batch_idx'] = list(new_batch['batch_idx']) for i in range(len(new_batch['batch_idx'])): new_batch['batch_idx'][i] += i # add target image index for build_targets() new_batch['batch_idx'] = torch.cat(new_batch['batch_idx'], 0) return new_batch # Classification dataloaders ------------------------------------------------------------------------------------------- class ClassificationDataset(torchvision.datasets.ImageFolder): """ YOLOv5 Classification Dataset. Arguments root: Dataset path transform: torchvision transforms, used by default album_transform: Albumentations transforms, used if installed """ def __init__(self, root, augment, imgsz, cache=False): super().__init__(root=root) self.torch_transforms = classify_transforms(imgsz) self.album_transforms = classify_albumentations(augment, imgsz) if augment else None self.cache_ram = cache is True or cache == 'ram' self.cache_disk = cache == 'disk' self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples] # file, index, npy, im def __getitem__(self, i): f, j, fn, im = self.samples[i] # filename, index, filename.with_suffix('.npy'), image if self.cache_ram and im is None: im = self.samples[i][3] = cv2.imread(f) elif self.cache_disk: if not fn.exists(): # load npy np.save(fn.as_posix(), cv2.imread(f)) im = np.load(fn) else: # read image im = cv2.imread(f) # BGR if self.album_transforms: sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))['image'] else: sample = self.torch_transforms(im) return {'img': sample, 'cls': j} def __len__(self) -> int: return len(self.samples) # TODO: support semantic segmentation class SemanticDataset(BaseDataset): def __init__(self): pass