k20hcmus/FishEye8K
3
1import contextlib2import glob3import hashlib4import json5import math6import os7import random8import shutil9import time10from itertools import repeat11from multiprocessing.pool import Pool, ThreadPool12from pathlib import Path13from threading import Thread14from urllib.parse import urlparse15 16import numpy as np17import psutil18import torch19import torch.nn.functional as F20import torchvision21import yaml22from PIL import ExifTags, Image, ImageOps23from torch.utils.data import DataLoader, Dataset, dataloader, distributed24from tqdm import tqdm25 26from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste,27 letterbox, mixup, random_perspective)28from utils.general import (DATASETS_DIR, LOGGER, NUM_THREADS, TQDM_BAR_FORMAT, check_dataset, check_requirements,29 check_yaml, clean_str, cv2, is_colab, is_kaggle, segments2boxes, unzip_file, xyn2xy,30 xywh2xyxy, xywhn2xyxy, xyxy2xywhn)31from utils.torch_utils import torch_distributed_zero_first32 33# Parameters34HELP_URL = 'See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'35IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm' # include image suffixes36VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' # include video suffixes37LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html38RANK = int(os.getenv('RANK', -1))39PIN_MEMORY = str(os.getenv('PIN_MEMORY', True)).lower() == 'true' # global pin_memory for dataloaders40 41# Get orientation exif tag42for orientation in ExifTags.TAGS.keys():43 if ExifTags.TAGS[orientation] == 'Orientation':44 break45 46 47def get_hash(paths):48 # Returns a single hash value of a list of paths (files or dirs)49 size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes50 h = hashlib.md5(str(size).encode()) # hash sizes51 h.update(''.join(paths).encode()) # hash paths52 return h.hexdigest() # return hash53 54 55def exif_size(img):56 # Returns exif-corrected PIL size57 s = img.size # (width, height)58 with contextlib.suppress(Exception):59 rotation = dict(img._getexif().items())[orientation]60 if rotation in [6, 8]: # rotation 270 or 9061 s = (s[1], s[0])62 return s63 64 65def exif_transpose(image):66 """67 Transpose a PIL image accordingly if it has an EXIF Orientation tag.68 Inplace version of https://github.com/python-pillow/Pillow/blob/master/src/PIL/ImageOps.py exif_transpose()69 70 :param image: The image to transpose.71 :return: An image.72 """73 exif = image.getexif()74 orientation = exif.get(0x0112, 1) # default 175 if orientation > 1:76 method = {77 2: Image.FLIP_LEFT_RIGHT,78 3: Image.ROTATE_180,79 4: Image.FLIP_TOP_BOTTOM,80 5: Image.TRANSPOSE,81 6: Image.ROTATE_270,82 7: Image.TRANSVERSE,83 8: Image.ROTATE_90}.get(orientation)84 if method is not None:85 image = image.transpose(method)86 del exif[0x0112]87 image.info["exif"] = exif.tobytes()88 return image89 90 91def seed_worker(worker_id):92 # Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader93 worker_seed = torch.initial_seed() % 2 ** 3294 np.random.seed(worker_seed)95 random.seed(worker_seed)96 97 98def create_dataloader(path,99 imgsz,100 batch_size,101 stride,102 single_cls=False,103 hyp=None,104 augment=False,105 cache=False,106 pad=0.0,107 rect=False,108 rank=-1,109 workers=8,110 image_weights=False,111 close_mosaic=False,112 quad=False,113 min_items=0,114 prefix='',115 shuffle=False):116 if rect and shuffle:117 LOGGER.warning('WARNING ⚠️ --rect is incompatible with DataLoader shuffle, setting shuffle=False')118 shuffle = False119 with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP120 dataset = LoadImagesAndLabels(121 path,122 imgsz,123 batch_size,124 augment=augment, # augmentation125 hyp=hyp, # hyperparameters126 rect=rect, # rectangular batches127 cache_images=cache,128 single_cls=single_cls,129 stride=int(stride),130 pad=pad,131 image_weights=image_weights,132 min_items=min_items,133 prefix=prefix)134 135 batch_size = min(batch_size, len(dataset))136 nd = torch.cuda.device_count() # number of CUDA devices137 nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers138 sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)139 #loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates140 loader = DataLoader if image_weights or close_mosaic else InfiniteDataLoader141 generator = torch.Generator()142 generator.manual_seed(6148914691236517205 + RANK)143 return loader(dataset,144 batch_size=batch_size,145 shuffle=shuffle and sampler is None,146 num_workers=nw,147 sampler=sampler,148 pin_memory=PIN_MEMORY,149 collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,150 worker_init_fn=seed_worker,151 generator=generator), dataset152 153 154class InfiniteDataLoader(dataloader.DataLoader):155 """ Dataloader that reuses workers156 157 Uses same syntax as vanilla DataLoader158 """159 160 def __init__(self, *args, **kwargs):161 super().__init__(*args, **kwargs)162 object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))163 self.iterator = super().__iter__()164 165 def __len__(self):166 return len(self.batch_sampler.sampler)167 168 def __iter__(self):169 for _ in range(len(self)):170 yield next(self.iterator)171 172 173class _RepeatSampler:174 """ Sampler that repeats forever175 176 Args:177 sampler (Sampler)178 """179 180 def __init__(self, sampler):181 self.sampler = sampler182 183 def __iter__(self):184 while True:185 yield from iter(self.sampler)186 187 188class LoadScreenshots:189 # YOLOv5 screenshot dataloader, i.e. `python detect.py --source "screen 0 100 100 512 256"`190 def __init__(self, source, img_size=640, stride=32, auto=True, transforms=None):191 # source = [screen_number left top width height] (pixels)192 check_requirements('mss')193 import mss194 195 source, *params = source.split()196 self.screen, left, top, width, height = 0, None, None, None, None # default to full screen 0197 if len(params) == 1:198 self.screen = int(params[0])199 elif len(params) == 4:200 left, top, width, height = (int(x) for x in params)201 elif len(params) == 5:202 self.screen, left, top, width, height = (int(x) for x in params)203 self.img_size = img_size204 self.stride = stride205 self.transforms = transforms206 self.auto = auto207 self.mode = 'stream'208 self.frame = 0209 self.sct = mss.mss()210 211 # Parse monitor shape212 monitor = self.sct.monitors[self.screen]213 self.top = monitor["top"] if top is None else (monitor["top"] + top)214 self.left = monitor["left"] if left is None else (monitor["left"] + left)215 self.width = width or monitor["width"]216 self.height = height or monitor["height"]217 self.monitor = {"left": self.left, "top": self.top, "width": self.width, "height": self.height}218 219 def __iter__(self):220 return self221 222 def __next__(self):223 # mss screen capture: get raw pixels from the screen as np array224 im0 = np.array(self.sct.grab(self.monitor))[:, :, :3] # [:, :, :3] BGRA to BGR225 s = f"screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: "226 227 if self.transforms:228 im = self.transforms(im0) # transforms229 else:230 im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0] # padded resize231 im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB232 im = np.ascontiguousarray(im) # contiguous233 self.frame += 1234 return str(self.screen), im, im0, None, s # screen, img, original img, im0s, s235 236 237class LoadImages:238 # YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`239 def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vid_stride=1):240 files = []241 for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:242 p = str(Path(p).resolve())243 if '*' in p:244 files.extend(sorted(glob.glob(p, recursive=True))) # glob245 elif os.path.isdir(p):246 files.extend(sorted(glob.glob(os.path.join(p, '*.*')))) # dir247 elif os.path.isfile(p):248 files.append(p) # files249 else:250 raise FileNotFoundError(f'{p} does not exist')251 252 images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]253 videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]254 ni, nv = len(images), len(videos)255 256 self.img_size = img_size257 self.stride = stride258 self.files = images + videos259 self.nf = ni + nv # number of files260 self.video_flag = [False] * ni + [True] * nv261 self.mode = 'image'262 self.auto = auto263 self.transforms = transforms # optional264 self.vid_stride = vid_stride # video frame-rate stride265 if any(videos):266 self._new_video(videos[0]) # new video267 else:268 self.cap = None269 assert self.nf > 0, f'No images or videos found in {p}. ' \270 f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}'271 272 def __iter__(self):273 self.count = 0274 return self275 276 def __next__(self):277 if self.count == self.nf:278 raise StopIteration279 path = self.files[self.count]280 281 if self.video_flag[self.count]:282 # Read video283 self.mode = 'video'284 for _ in range(self.vid_stride):285 self.cap.grab()286 ret_val, im0 = self.cap.retrieve()287 while not ret_val:288 self.count += 1289 self.cap.release()290 if self.count == self.nf: # last video291 raise StopIteration292 path = self.files[self.count]293 self._new_video(path)294 ret_val, im0 = self.cap.read()295 296 self.frame += 1297 # im0 = self._cv2_rotate(im0) # for use if cv2 autorotation is False298 s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '299 300 else:301 # Read image302 self.count += 1303 im0 = cv2.imread(path) # BGR304 assert im0 is not None, f'Image Not Found {path}'305 s = f'image {self.count}/{self.nf} {path}: '306 307 if self.transforms:308 im = self.transforms(im0) # transforms309 else:310 im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0] # padded resize311 im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB312 im = np.ascontiguousarray(im) # contiguous313 314 return path, im, im0, self.cap, s315 316 def _new_video(self, path):317 # Create a new video capture object318 self.frame = 0319 self.cap = cv2.VideoCapture(path)320 self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)321 self.orientation = int(self.cap.get(cv2.CAP_PROP_ORIENTATION_META)) # rotation degrees322 # self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0) # disable https://github.com/ultralytics/yolov5/issues/8493323 324 def _cv2_rotate(self, im):325 # Rotate a cv2 video manually326 if self.orientation == 0:327 return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE)328 elif self.orientation == 180:329 return cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE)330 elif self.orientation == 90:331 return cv2.rotate(im, cv2.ROTATE_180)332 return im333 334 def __len__(self):335 return self.nf # number of files336 337 338class LoadStreams:339 # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams`340 def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True, transforms=None, vid_stride=1):341 torch.backends.cudnn.benchmark = True # faster for fixed-size inference342 self.mode = 'stream'343 self.img_size = img_size344 self.stride = stride345 self.vid_stride = vid_stride # video frame-rate stride346 sources = Path(sources).read_text().rsplit() if os.path.isfile(sources) else [sources]347 n = len(sources)348 self.sources = [clean_str(x) for x in sources] # clean source names for later349 self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n350 for i, s in enumerate(sources): # index, source351 # Start thread to read frames from video stream352 st = f'{i + 1}/{n}: {s}... '353 if urlparse(s).hostname in ('www.youtube.com', 'youtube.com', 'youtu.be'): # if source is YouTube video354 # YouTube format i.e. 'https://www.youtube.com/watch?v=Zgi9g1ksQHc' or 'https://youtu.be/Zgi9g1ksQHc'355 check_requirements(('pafy', 'youtube_dl==2020.12.2'))356 import pafy357 s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL358 s = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam359 if s == 0:360 assert not is_colab(), '--source 0 webcam unsupported on Colab. Rerun command in a local environment.'361 assert not is_kaggle(), '--source 0 webcam unsupported on Kaggle. Rerun command in a local environment.'362 cap = cv2.VideoCapture(s)363 assert cap.isOpened(), f'{st}Failed to open {s}'364 w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))365 h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))366 fps = cap.get(cv2.CAP_PROP_FPS) # warning: may return 0 or nan367 self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback368 self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30 # 30 FPS fallback369 370 _, self.imgs[i] = cap.read() # guarantee first frame371 self.threads[i] = Thread(target=self.update, args=([i, cap, s]), daemon=True)372 LOGGER.info(f"{st} Success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")373 self.threads[i].start()374 LOGGER.info('') # newline375 376 # check for common shapes377 s = np.stack([letterbox(x, img_size, stride=stride, auto=auto)[0].shape for x in self.imgs])378 self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal379 self.auto = auto and self.rect380 self.transforms = transforms # optional381 if not self.rect:382 LOGGER.warning('WARNING ⚠️ Stream shapes differ. For optimal performance supply similarly-shaped streams.')383 384 def update(self, i, cap, stream):385 # Read stream `i` frames in daemon thread386 n, f = 0, self.frames[i] # frame number, frame array387 while cap.isOpened() and n < f:388 n += 1389 cap.grab() # .read() = .grab() followed by .retrieve()390 if n % self.vid_stride == 0:391 success, im = cap.retrieve()392 if success:393 self.imgs[i] = im394 else:395 LOGGER.warning('WARNING ⚠️ Video stream unresponsive, please check your IP camera connection.')396 self.imgs[i] = np.zeros_like(self.imgs[i])397 cap.open(stream) # re-open stream if signal was lost398 time.sleep(0.0) # wait time399 400 def __iter__(self):401 self.count = -1402 return self403 404 def __next__(self):405 self.count += 1406 if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit407 cv2.destroyAllWindows()408 raise StopIteration409 410 im0 = self.imgs.copy()411 if self.transforms:412 im = np.stack([self.transforms(x) for x in im0]) # transforms413 else:414 im = np.stack([letterbox(x, self.img_size, stride=self.stride, auto=self.auto)[0] for x in im0]) # resize415 im = im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW416 im = np.ascontiguousarray(im) # contiguous417 418 return self.sources, im, im0, None, ''419 420 def __len__(self):421 return len(self.sources) # 1E12 frames = 32 streams at 30 FPS for 30 years422 423 424def img2label_paths(img_paths):425 # Define label paths as a function of image paths426 sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}' # /images/, /labels/ substrings427 return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]428 429 430class LoadImagesAndLabels(Dataset):431 # YOLOv5 train_loader/val_loader, loads images and labels for training and validation432 cache_version = 0.6 # dataset labels *.cache version433 rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4]434 435 def __init__(self,436 path,437 img_size=640,438 batch_size=16,439 augment=False,440 hyp=None,441 rect=False,442 image_weights=False,443 cache_images=False,444 single_cls=False,445 stride=32,446 pad=0.0,447 min_items=0,448 prefix=''):449 self.img_size = img_size450 self.augment = augment451 self.hyp = hyp452 self.image_weights = image_weights453 self.rect = False if image_weights else rect454 self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)455 self.mosaic_border = [-img_size // 2, -img_size // 2]456 self.stride = stride457 self.path = path458 self.albumentations = Albumentations(size=img_size) if augment else None459 460 try:461 f = [] # image files462 for p in path if isinstance(path, list) else [path]:463 p = Path(p) # os-agnostic464 if p.is_dir(): # dir465 f += glob.glob(str(p / '**' / '*.*'), recursive=True)466 # f = list(p.rglob('*.*')) # pathlib467 elif p.is_file(): # file468 with open(p) as t:469 t = t.read().strip().splitlines()470 parent = str(p.parent) + os.sep471 f += [x.replace('./', parent, 1) if x.startswith('./') else x for x in t] # to global path472 # f += [p.parent / x.lstrip(os.sep) for x in t] # to global path (pathlib)473 else:474 raise FileNotFoundError(f'{prefix}{p} does not exist')475 self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)476 # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib477 assert self.im_files, f'{prefix}No images found'478 except Exception as e:479 raise Exception(f'{prefix}Error loading data from {path}: {e}\n{HELP_URL}') from e480 481 # Check cache482 self.label_files = img2label_paths(self.im_files) # labels483 cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')484 try:485 cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict486 assert cache['version'] == self.cache_version # matches current version487 assert cache['hash'] == get_hash(self.label_files + self.im_files) # identical hash488 except Exception:489 cache, exists = self.cache_labels(cache_path, prefix), False # run cache ops490 491 # Display cache492 nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupt, total493 if exists and LOCAL_RANK in {-1, 0}:494 d = f"Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt"495 tqdm(None, desc=prefix + d, total=n, initial=n, bar_format=TQDM_BAR_FORMAT) # display cache results496 if cache['msgs']:497 LOGGER.info('\n'.join(cache['msgs'])) # display warnings498 assert nf > 0 or not augment, f'{prefix}No labels found in {cache_path}, can not start training. {HELP_URL}'499 500 # Read cache501 [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items502 labels, shapes, self.segments = zip(*cache.values())503 nl = len(np.concatenate(labels, 0)) # number of labels504 assert nl > 0 or not augment, f'{prefix}All labels empty in {cache_path}, can not start training. {HELP_URL}'505 self.labels = list(labels)506 self.shapes = np.array(shapes)507 self.im_files = list(cache.keys()) # update508 self.label_files = img2label_paths(cache.keys()) # update509 510 # Filter images511 if min_items:512 include = np.array([len(x) >= min_items for x in self.labels]).nonzero()[0].astype(int)513 LOGGER.info(f'{prefix}{n - len(include)}/{n} images filtered from dataset')514 self.im_files = [self.im_files[i] for i in include]515 self.label_files = [self.label_files[i] for i in include]516 self.labels = [self.labels[i] for i in include]517 self.segments = [self.segments[i] for i in include]518 self.shapes = self.shapes[include] # wh519 520 # Create indices521 n = len(self.shapes) # number of images522 bi = np.floor(np.arange(n) / batch_size).astype(int) # batch index523 nb = bi[-1] + 1 # number of batches524 self.batch = bi # batch index of image525 self.n = n526 self.indices = range(n)527 528 # Update labels529 include_class = [] # filter labels to include only these classes (optional)530 include_class_array = np.array(include_class).reshape(1, -1)531 for i, (label, segment) in enumerate(zip(self.labels, self.segments)):532 if include_class:533 j = (label[:, 0:1] == include_class_array).any(1)534 self.labels[i] = label[j]535 if segment:536 self.segments[i] = segment[j]537 if single_cls: # single-class training, merge all classes into 0538 self.labels[i][:, 0] = 0539 540 # Rectangular Training541 if self.rect:542 # Sort by aspect ratio543 s = self.shapes # wh544 ar = s[:, 1] / s[:, 0] # aspect ratio545 irect = ar.argsort()546 self.im_files = [self.im_files[i] for i in irect]547 self.label_files = [self.label_files[i] for i in irect]548 self.labels = [self.labels[i] for i in irect]549 self.segments = [self.segments[i] for i in irect]550 self.shapes = s[irect] # wh551 ar = ar[irect]552 553 # Set training image shapes554 shapes = [[1, 1]] * nb555 for i in range(nb):556 ari = ar[bi == i]557 mini, maxi = ari.min(), ari.max()558 if maxi < 1:559 shapes[i] = [maxi, 1]560 elif mini > 1:561 shapes[i] = [1, 1 / mini]562 563 self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(int) * stride564 565 # Cache images into RAM/disk for faster training566 if cache_images == 'ram' and not self.check_cache_ram(prefix=prefix):567 cache_images = False568 self.ims = [None] * n569 self.npy_files = [Path(f).with_suffix('.npy') for f in self.im_files]570 if cache_images:571 b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes572 self.im_hw0, self.im_hw = [None] * n, [None] * n573 fcn = self.cache_images_to_disk if cache_images == 'disk' else self.load_image574 results = ThreadPool(NUM_THREADS).imap(fcn, range(n))575 pbar = tqdm(enumerate(results), total=n, bar_format=TQDM_BAR_FORMAT, disable=LOCAL_RANK > 0)576 for i, x in pbar:577 if cache_images == 'disk':578 b += self.npy_files[i].stat().st_size579 else: # 'ram'580 self.ims[i], self.im_hw0[i], self.im_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)581 b += self.ims[i].nbytes582 pbar.desc = f'{prefix}Caching images ({b / gb:.1f}GB {cache_images})'583 pbar.close()584 585 def check_cache_ram(self, safety_margin=0.1, prefix=''):586 # Check image caching requirements vs available memory587 b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes588 n = min(self.n, 30) # extrapolate from 30 random images589 for _ in range(n):590 im = cv2.imread(random.choice(self.im_files)) # sample image591 ratio = self.img_size / max(im.shape[0], im.shape[1]) # max(h, w) # ratio592 b += im.nbytes * ratio ** 2593 mem_required = b * self.n / n # GB required to cache dataset into RAM594 mem = psutil.virtual_memory()595 cache = mem_required * (1 + safety_margin) < mem.available # to cache or not to cache, that is the question596 if not cache:597 LOGGER.info(f"{prefix}{mem_required / gb:.1f}GB RAM required, "598 f"{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, "599 f"{'caching images ✅' if cache else 'not caching images ⚠️'}")600 return cache601 602 def cache_labels(self, path=Path('./labels.cache'), prefix=''):603 # Cache dataset labels, check images and read shapes604 x = {} # dict605 nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages606 desc = f"{prefix}Scanning {path.parent / path.stem}..."607 with Pool(NUM_THREADS) as pool:608 pbar = tqdm(pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))),609 desc=desc,610 total=len(self.im_files),611 bar_format=TQDM_BAR_FORMAT)612 for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:613 nm += nm_f614 nf += nf_f615 ne += ne_f616 nc += nc_f617 if im_file:618 x[im_file] = [lb, shape, segments]619 if msg:620 msgs.append(msg)621 pbar.desc = f"{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt"622 623 pbar.close()624 if msgs:625 LOGGER.info('\n'.join(msgs))626 if nf == 0:627 LOGGER.warning(f'{prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}')628 x['hash'] = get_hash(self.label_files + self.im_files)629 x['results'] = nf, nm, ne, nc, len(self.im_files)630 x['msgs'] = msgs # warnings631 x['version'] = self.cache_version # cache version632 try:633 np.save(path, x) # save cache for next time634 path.with_suffix('.cache.npy').rename(path) # remove .npy suffix635 LOGGER.info(f'{prefix}New cache created: {path}')636 except Exception as e:637 LOGGER.warning(f'{prefix}WARNING ⚠️ Cache directory {path.parent} is not writeable: {e}') # not writeable638 return x639 640 def __len__(self):641 return len(self.im_files)642 643 # def __iter__(self):644 # self.count = -1645 # print('ran dataset iter')646 # #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)647 # return self648 649 def __getitem__(self, index):650 index = self.indices[index] # linear, shuffled, or image_weights651 652 hyp = self.hyp653 mosaic = self.mosaic and random.random() < hyp['mosaic']654 if mosaic:655 # Load mosaic656 img, labels = self.load_mosaic(index)657 shapes = None658 659 # MixUp augmentation660 if random.random() < hyp['mixup']:661 img, labels = mixup(img, labels, *self.load_mosaic(random.randint(0, self.n - 1)))662 663 else:664 # Load image665 img, (h0, w0), (h, w) = self.load_image(index)666 667 # Letterbox668 shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape669 img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)670 shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling671 672 labels = self.labels[index].copy()673 if labels.size: # normalized xywh to pixel xyxy format674 labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])675 676 if self.augment:677 img, labels = random_perspective(img,678 labels,679 degrees=hyp['degrees'],680 translate=hyp['translate'],681 scale=hyp['scale'],682 shear=hyp['shear'],683 perspective=hyp['perspective'])684 685 nl = len(labels) # number of labels686 if nl:687 labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)688 689 if self.augment:690 # Albumentations691 img, labels = self.albumentations(img, labels)692 nl = len(labels) # update after albumentations693 694 # HSV color-space695 augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])696 697 # Flip up-down698 if random.random() < hyp['flipud']:699 img = np.flipud(img)700 if nl:701 labels[:, 2] = 1 - labels[:, 2]702 703 # Flip left-right704 if random.random() < hyp['fliplr']:705 img = np.fliplr(img)706 if nl:707 labels[:, 1] = 1 - labels[:, 1]708 709 # Cutouts710 # labels = cutout(img, labels, p=0.5)711 # nl = len(labels) # update after cutout712 713 labels_out = torch.zeros((nl, 6))714 if nl:715 labels_out[:, 1:] = torch.from_numpy(labels)716 717 # Convert718 img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB719 img = np.ascontiguousarray(img)720 721 return torch.from_numpy(img), labels_out, self.im_files[index], shapes722 723 def load_image(self, i):724 # Loads 1 image from dataset index 'i', returns (im, original hw, resized hw)725 im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i],726 if im is None: # not cached in RAM727 if fn.exists(): # load npy728 im = np.load(fn)729 else: # read image730 im = cv2.imread(f) # BGR731 assert im is not None, f'Image Not Found {f}'732 h0, w0 = im.shape[:2] # orig hw733 r = self.img_size / max(h0, w0) # ratio734 if r != 1: # if sizes are not equal735 interp = cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA736 im = cv2.resize(im, (int(w0 * r), int(h0 * r)), interpolation=interp)737 return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized738 return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized739 740 def cache_images_to_disk(self, i):741 # Saves an image as an *.npy file for faster loading742 f = self.npy_files[i]743 if not f.exists():744 np.save(f.as_posix(), cv2.imread(self.im_files[i]))745 746 def load_mosaic(self, index):747 # YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic748 labels4, segments4 = [], []749 s = self.img_size750 yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border) # mosaic center x, y751 indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices752 random.shuffle(indices)753 for i, index in enumerate(indices):754 # Load image755 img, _, (h, w) = self.load_image(index)756 757 # place img in img4758 if i == 0: # top left759 img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles760 x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)761 x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)762 elif i == 1: # top right763 x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc764 x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h765 elif i == 2: # bottom left766 x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)767 x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)768 elif i == 3: # bottom right769 x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)770 x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)771 772 img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]773 padw = x1a - x1b774 padh = y1a - y1b775 776 # Labels777 labels, segments = self.labels[index].copy(), self.segments[index].copy()778 if labels.size:779 labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format780 segments = [xyn2xy(x, w, h, padw, padh) for x in segments]781 labels4.append(labels)782 segments4.extend(segments)783 784 # Concat/clip labels785 labels4 = np.concatenate(labels4, 0)786 for x in (labels4[:, 1:], *segments4):787 np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()788 # img4, labels4 = replicate(img4, labels4) # replicate789 790 # Augment791 img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])792 img4, labels4 = random_perspective(img4,793 labels4,794 segments4,795 degrees=self.hyp['degrees'],796 translate=self.hyp['translate'],797 scale=self.hyp['scale'],798 shear=self.hyp['shear'],799 perspective=self.hyp['perspective'],800 border=self.mosaic_border) # border to remove801 802 return img4, labels4803 804 def load_mosaic9(self, index):805 # YOLOv5 9-mosaic loader. Loads 1 image + 8 random images into a 9-image mosaic806 labels9, segments9 = [], []807 s = self.img_size808 indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices809 random.shuffle(indices)810 hp, wp = -1, -1 # height, width previous811 for i, index in enumerate(indices):812 # Load image813 img, _, (h, w) = self.load_image(index)814 815 # place img in img9816 if i == 0: # center817 img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles818 h0, w0 = h, w819 c = s, s, s + w, s + h # xmin, ymin, xmax, ymax (base) coordinates820 elif i == 1: # top821 c = s, s - h, s + w, s822 elif i == 2: # top right823 c = s + wp, s - h, s + wp + w, s824 elif i == 3: # right825 c = s + w0, s, s + w0 + w, s + h826 elif i == 4: # bottom right827 c = s + w0, s + hp, s + w0 + w, s + hp + h828 elif i == 5: # bottom829 c = s + w0 - w, s + h0, s + w0, s + h0 + h830 elif i == 6: # bottom left831 c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h832 elif i == 7: # left833 c = s - w, s + h0 - h, s, s + h0834 elif i == 8: # top left835 c = s - w, s + h0 - hp - h, s, s + h0 - hp836 837 padx, pady = c[:2]838 x1, y1, x2, y2 = (max(x, 0) for x in c) # allocate coords839 840 # Labels841 labels, segments = self.labels[index].copy(), self.segments[index].copy()842 if labels.size:843 labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format844 segments = [xyn2xy(x, w, h, padx, pady) for x in segments]845 labels9.append(labels)846 segments9.extend(segments)847 848 # Image849 img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]850 hp, wp = h, w # height, width previous851 852 # Offset853 yc, xc = (int(random.uniform(0, s)) for _ in self.mosaic_border) # mosaic center x, y854 img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]855 856 # Concat/clip labels857 labels9 = np.concatenate(labels9, 0)858 labels9[:, [1, 3]] -= xc859 labels9[:, [2, 4]] -= yc860 c = np.array([xc, yc]) # centers861 segments9 = [x - c for x in segments9]862 863 for x in (labels9[:, 1:], *segments9):864 np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()865 # img9, labels9 = replicate(img9, labels9) # replicate866 867 # Augment868 img9, labels9, segments9 = copy_paste(img9, labels9, segments9, p=self.hyp['copy_paste'])869 img9, labels9 = random_perspective(img9,870 labels9,871 segments9,872 degrees=self.hyp['degrees'],873 translate=self.hyp['translate'],874 scale=self.hyp['scale'],875 shear=self.hyp['shear'],876 perspective=self.hyp['perspective'],877 border=self.mosaic_border) # border to remove878 879 return img9, labels9880 881 @staticmethod882 def collate_fn(batch):883 im, label, path, shapes = zip(*batch) # transposed884 for i, lb in enumerate(label):885 lb[:, 0] = i # add target image index for build_targets()886 return torch.stack(im, 0), torch.cat(label, 0), path, shapes887 888 @staticmethod889 def collate_fn4(batch):890 im, label, path, shapes = zip(*batch) # transposed891 n = len(shapes) // 4892 im4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]893 894 ho = torch.tensor([[0.0, 0, 0, 1, 0, 0]])895 wo = torch.tensor([[0.0, 0, 1, 0, 0, 0]])896 s = torch.tensor([[1, 1, 0.5, 0.5, 0.5, 0.5]]) # scale897 for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW898 i *= 4899 if random.random() < 0.5:900 im1 = F.interpolate(im[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear',901 align_corners=False)[0].type(im[i].type())902 lb = label[i]903 else:904 im1 = torch.cat((torch.cat((im[i], im[i + 1]), 1), torch.cat((im[i + 2], im[i + 3]), 1)), 2)905 lb = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s906 im4.append(im1)907 label4.append(lb)908 909 for i, lb in enumerate(label4):910 lb[:, 0] = i # add target image index for build_targets()911 912 return torch.stack(im4, 0), torch.cat(label4, 0), path4, shapes4913 914 915# Ancillary functions --------------------------------------------------------------------------------------------------916def flatten_recursive(path=DATASETS_DIR / 'coco128'):917 # Flatten a recursive directory by bringing all files to top level918 new_path = Path(f'{str(path)}_flat')919 if os.path.exists(new_path):920 shutil.rmtree(new_path) # delete output folder921 os.makedirs(new_path) # make new output folder922 for file in tqdm(glob.glob(f'{str(Path(path))}/**/*.*', recursive=True)):923 shutil.copyfile(file, new_path / Path(file).name)924 925 926def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders import *; extract_boxes()927 # Convert detection dataset into classification dataset, with one directory per class928 path = Path(path) # images dir929 shutil.rmtree(path / 'classification') if (path / 'classification').is_dir() else None # remove existing930 files = list(path.rglob('*.*'))931 n = len(files) # number of files932 for im_file in tqdm(files, total=n):933 if im_file.suffix[1:] in IMG_FORMATS:934 # image935 im = cv2.imread(str(im_file))[..., ::-1] # BGR to RGB936 h, w = im.shape[:2]937 938 # labels939 lb_file = Path(img2label_paths([str(im_file)])[0])940 if Path(lb_file).exists():941 with open(lb_file) as f:942 lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels943 944 for j, x in enumerate(lb):945 c = int(x[0]) # class946 f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename947 if not f.parent.is_dir():948 f.parent.mkdir(parents=True)949 950 b = x[1:] * [w, h, w, h] # box951 # b[2:] = b[2:].max() # rectangle to square952 b[2:] = b[2:] * 1.2 + 3 # pad953 b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(int)954 955 b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image956 b[[1, 3]] = np.clip(b[[1, 3]], 0, h)957 assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'958 959 960def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False):961 """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files962 Usage: from utils.dataloaders import *; autosplit()963 Arguments964 path: Path to images directory965 weights: Train, val, test weights (list, tuple)966 annotated_only: Only use images with an annotated txt file967 """968 path = Path(path) # images dir969 files = sorted(x for x in path.rglob('*.*') if x.suffix[1:].lower() in IMG_FORMATS) # image files only970 n = len(files) # number of files971 random.seed(0) # for reproducibility972 indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split973 974 txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files975 for x in txt:976 if (path.parent / x).exists():977 (path.parent / x).unlink() # remove existing978 979 print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)980 for i, img in tqdm(zip(indices, files), total=n):981 if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label982 with open(path.parent / txt[i], 'a') as f:983 f.write(f'./{img.relative_to(path.parent).as_posix()}' + '\n') # add image to txt file984 985 986def verify_image_label(args):987 # Verify one image-label pair988 im_file, lb_file, prefix = args989 nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments990 try:991 # verify images992 im = Image.open(im_file)993 im.verify() # PIL verify994 shape = exif_size(im) # image size995 assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'996 assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}'997 if im.format.lower() in ('jpg', 'jpeg'):998 with open(im_file, 'rb') as f:999 f.seek(-2, 2)1000 if f.read() != b'\xff\xd9': # corrupt JPEG1001 ImageOps.exif_transpose(Image.open(im_file)).save(im_file, 'JPEG', subsampling=0, quality=100)1002 msg = f'{prefix}WARNING ⚠️ {im_file}: corrupt JPEG restored and saved'1003 1004 # verify labels1005 if os.path.isfile(lb_file):1006 nf = 1 # label found1007 with open(lb_file) as f:1008 lb = [x.split() for x in f.read().strip().splitlines() if len(x)]1009 if any(len(x) > 6 for x in lb): # is segment1010 classes = np.array([x[0] for x in lb], dtype=np.float32)1011 segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in lb] # (cls, xy1...)1012 lb = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)1013 lb = np.array(lb, dtype=np.float32)1014 nl = len(lb)1015 if nl:1016 assert lb.shape[1] == 5, f'labels require 5 columns, {lb.shape[1]} columns detected'1017 assert (lb >= 0).all(), f'negative label values {lb[lb < 0]}'1018 assert (lb[:, 1:] <= 1).all(), f'non-normalized or out of bounds coordinates {lb[:, 1:][lb[:, 1:] > 1]}'1019 _, i = np.unique(lb, axis=0, return_index=True)1020 if len(i) < nl: # duplicate row check1021 lb = lb[i] # remove duplicates1022 if segments:1023 segments = [segments[x] for x in i]1024 msg = f'{prefix}WARNING ⚠️ {im_file}: {nl - len(i)} duplicate labels removed'1025 else:1026 ne = 1 # label empty1027 lb = np.zeros((0, 5), dtype=np.float32)1028 else:1029 nm = 1 # label missing1030 lb = np.zeros((0, 5), dtype=np.float32)1031 return im_file, lb, shape, segments, nm, nf, ne, nc, msg1032 except Exception as e:1033 nc = 11034 msg = f'{prefix}WARNING ⚠️ {im_file}: ignoring corrupt image/label: {e}'1035 return [None, None, None, None, nm, nf, ne, nc, msg]1036 1037 1038class HUBDatasetStats():1039 """ Class for generating HUB dataset JSON and `-hub` dataset directory1040 1041 Arguments1042 path: Path to data.yaml or data.zip (with data.yaml inside data.zip)1043 autodownload: Attempt to download dataset if not found locally1044 1045 Usage1046 from utils.dataloaders import HUBDatasetStats1047 stats = HUBDatasetStats('coco128.yaml', autodownload=True) # usage 11048 stats = HUBDatasetStats('path/to/coco128.zip') # usage 21049 stats.get_json(save=False)1050 stats.process_images()1051 """1052 1053 def __init__(self, path='coco128.yaml', autodownload=False):1054 # Initialize class1055 zipped, data_dir, yaml_path = self._unzip(Path(path))1056 try:1057 with open(check_yaml(yaml_path), errors='ignore') as f:1058 data = yaml.safe_load(f) # data dict1059 if zipped:1060 data['path'] = data_dir1061 except Exception as e:1062 raise Exception("error/HUB/dataset_stats/yaml_load") from e1063 1064 check_dataset(data, autodownload) # download dataset if missing1065 self.hub_dir = Path(data['path'] + '-hub')1066 self.im_dir = self.hub_dir / 'images'1067 self.im_dir.mkdir(parents=True, exist_ok=True) # makes /images1068 self.stats = {'nc': data['nc'], 'names': list(data['names'].values())} # statistics dictionary1069 self.data = data1070 1071 @staticmethod1072 def _find_yaml(dir):1073 # Return data.yaml file1074 files = list(dir.glob('*.yaml')) or list(dir.rglob('*.yaml')) # try root level first and then recursive1075 assert files, f'No *.yaml file found in {dir}'1076 if len(files) > 1:1077 files = [f for f in files if f.stem == dir.stem] # prefer *.yaml files that match dir name1078 assert files, f'Multiple *.yaml files found in {dir}, only 1 *.yaml file allowed'1079 assert len(files) == 1, f'Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}'1080 return files[0]1081 1082 def _unzip(self, path):1083 # Unzip data.zip1084 if not str(path).endswith('.zip'): # path is data.yaml1085 return False, None, path1086 assert Path(path).is_file(), f'Error unzipping {path}, file not found'1087 unzip_file(path, path=path.parent)1088 dir = path.with_suffix('') # dataset directory == zip name1089 assert dir.is_dir(), f'Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/'1090 return True, str(dir), self._find_yaml(dir) # zipped, data_dir, yaml_path1091 1092 def _hub_ops(self, f, max_dim=1920):1093 # HUB ops for 1 image 'f': resize and save at reduced quality in /dataset-hub for web/app viewing1094 f_new = self.im_dir / Path(f).name # dataset-hub image filename1095 try: # use PIL1096 im = Image.open(f)1097 r = max_dim / max(im.height, im.width) # ratio1098 if r < 1.0: # image too large1099 im = im.resize((int(im.width * r), int(im.height * r)))1100 im.save(f_new, 'JPEG', quality=50, optimize=True) # save1101 except Exception as e: # use OpenCV1102 LOGGER.info(f'WARNING ⚠️ HUB ops PIL failure {f}: {e}')1103 im = cv2.imread(f)1104 im_height, im_width = im.shape[:2]1105 r = max_dim / max(im_height, im_width) # ratio1106 if r < 1.0: # image too large1107 im = cv2.resize(im, (int(im_width * r), int(im_height * r)), interpolation=cv2.INTER_AREA)1108 cv2.imwrite(str(f_new), im)1109 1110 def get_json(self, save=False, verbose=False):1111 # Return dataset JSON for Ultralytics HUB1112 def _round(labels):1113 # Update labels to integer class and 6 decimal place floats1114 return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]1115 1116 for split in 'train', 'val', 'test':1117 if self.data.get(split) is None:1118 self.stats[split] = None # i.e. no test set1119 continue1120 dataset = LoadImagesAndLabels(self.data[split]) # load dataset1121 x = np.array([1122 np.bincount(label[:, 0].astype(int), minlength=self.data['nc'])1123 for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics')]) # shape(128x80)1124 self.stats[split] = {1125 'instance_stats': {1126 'total': int(x.sum()),1127 'per_class': x.sum(0).tolist()},1128 'image_stats': {1129 'total': dataset.n,1130 'unlabelled': int(np.all(x == 0, 1).sum()),1131 'per_class': (x > 0).sum(0).tolist()},1132 'labels': [{1133 str(Path(k).name): _round(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)]}1134 1135 # Save, print and return1136 if save:1137 stats_path = self.hub_dir / 'stats.json'1138 print(f'Saving {stats_path.resolve()}...')1139 with open(stats_path, 'w') as f:1140 json.dump(self.stats, f) # save stats.json1141 if verbose:1142 print(json.dumps(self.stats, indent=2, sort_keys=False))1143 return self.stats1144 1145 def process_images(self):1146 # Compress images for Ultralytics HUB1147 for split in 'train', 'val', 'test':1148 if self.data.get(split) is None:1149 continue1150 dataset = LoadImagesAndLabels(self.data[split]) # load dataset1151 desc = f'{split} images'1152 for _ in tqdm(ThreadPool(NUM_THREADS).imap(self._hub_ops, dataset.im_files), total=dataset.n, desc=desc):1153 pass1154 print(f'Done. All images saved to {self.im_dir}')1155 return self.im_dir1156 1157 1158# Classification dataloaders -------------------------------------------------------------------------------------------1159class ClassificationDataset(torchvision.datasets.ImageFolder):1160 """1161 YOLOv5 Classification Dataset.1162 Arguments1163 root: Dataset path1164 transform: torchvision transforms, used by default1165 album_transform: Albumentations transforms, used if installed1166 """1167 1168 def __init__(self, root, augment, imgsz, cache=False):1169 super().__init__(root=root)1170 self.torch_transforms = classify_transforms(imgsz)1171 self.album_transforms = classify_albumentations(augment, imgsz) if augment else None1172 self.cache_ram = cache is True or cache == 'ram'1173 self.cache_disk = cache == 'disk'1174 self.samples = [list(x) + [Path(x[0]).with_suffix('.npy'), None] for x in self.samples] # file, index, npy, im1175 1176 def __getitem__(self, i):1177 f, j, fn, im = self.samples[i] # filename, index, filename.with_suffix('.npy'), image1178 if self.cache_ram and im is None:1179 im = self.samples[i][3] = cv2.imread(f)1180 elif self.cache_disk:1181 if not fn.exists(): # load npy1182 np.save(fn.as_posix(), cv2.imread(f))1183 im = np.load(fn)1184 else: # read image1185 im = cv2.imread(f) # BGR1186 if self.album_transforms:1187 sample = self.album_transforms(image=cv2.cvtColor(im, cv2.COLOR_BGR2RGB))["image"]1188 else:1189 sample = self.torch_transforms(im)1190 return sample, j1191 1192 1193def create_classification_dataloader(path,1194 imgsz=224,1195 batch_size=16,1196 augment=True,1197 cache=False,1198 rank=-1,1199 workers=8,1200 shuffle=True):