You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1002 line
42KB

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  2. """
  3. Dataloaders and dataset utils
  4. """
  5. import glob
  6. import hashlib
  7. import json
  8. import logging
  9. import os
  10. import random
  11. import shutil
  12. import time
  13. from itertools import repeat
  14. from multiprocessing.pool import ThreadPool, Pool
  15. from pathlib import Path
  16. from threading import Thread
  17. import cv2
  18. import numpy as np
  19. import torch
  20. import torch.nn.functional as F
  21. import yaml
  22. from PIL import Image, ExifTags
  23. from torch.utils.data import Dataset
  24. from tqdm import tqdm
  25. from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective
  26. from utils.general import check_dataset, check_requirements, check_yaml, clean_str, segments2boxes, \
  27. xywh2xyxy, xywhn2xyxy, xyxy2xywhn, xyn2xy
  28. from utils.torch_utils import torch_distributed_zero_first
  29. # Parameters
  30. HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
  31. IMG_FORMATS = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes
  32. VID_FORMATS = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
  33. NUM_THREADS = min(8, os.cpu_count()) # number of multiprocessing threads
  34. # Get orientation exif tag
  35. for orientation in ExifTags.TAGS.keys():
  36. if ExifTags.TAGS[orientation] == 'Orientation':
  37. break
  38. def get_hash(paths):
  39. # Returns a single hash value of a list of paths (files or dirs)
  40. size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes
  41. h = hashlib.md5(str(size).encode()) # hash sizes
  42. h.update(''.join(paths).encode()) # hash paths
  43. return h.hexdigest() # return hash
  44. def exif_size(img):
  45. # Returns exif-corrected PIL size
  46. s = img.size # (width, height)
  47. try:
  48. rotation = dict(img._getexif().items())[orientation]
  49. if rotation == 6: # rotation 270
  50. s = (s[1], s[0])
  51. elif rotation == 8: # rotation 90
  52. s = (s[1], s[0])
  53. except:
  54. pass
  55. return s
  56. def exif_transpose(image):
  57. """
  58. Transpose a PIL image accordingly if it has an EXIF Orientation tag.
  59. From https://github.com/python-pillow/Pillow/blob/master/src/PIL/ImageOps.py
  60. :param image: The image to transpose.
  61. :return: An image.
  62. """
  63. exif = image.getexif()
  64. orientation = exif.get(0x0112, 1) # default 1
  65. if orientation > 1:
  66. method = {2: Image.FLIP_LEFT_RIGHT,
  67. 3: Image.ROTATE_180,
  68. 4: Image.FLIP_TOP_BOTTOM,
  69. 5: Image.TRANSPOSE,
  70. 6: Image.ROTATE_270,
  71. 7: Image.TRANSVERSE,
  72. 8: Image.ROTATE_90,
  73. }.get(orientation)
  74. if method is not None:
  75. image = image.transpose(method)
  76. del exif[0x0112]
  77. image.info["exif"] = exif.tobytes()
  78. return image
  79. def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0,
  80. rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix=''):
  81. # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
  82. with torch_distributed_zero_first(rank):
  83. dataset = LoadImagesAndLabels(path, imgsz, batch_size,
  84. augment=augment, # augment images
  85. hyp=hyp, # augmentation hyperparameters
  86. rect=rect, # rectangular training
  87. cache_images=cache,
  88. single_cls=single_cls,
  89. stride=int(stride),
  90. pad=pad,
  91. image_weights=image_weights,
  92. prefix=prefix)
  93. batch_size = min(batch_size, len(dataset))
  94. nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, workers]) # number of workers
  95. sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None
  96. loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader
  97. # Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()
  98. dataloader = loader(dataset,
  99. batch_size=batch_size,
  100. num_workers=nw,
  101. sampler=sampler,
  102. pin_memory=True,
  103. collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)
  104. return dataloader, dataset
  105. class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):
  106. """ Dataloader that reuses workers
  107. Uses same syntax as vanilla DataLoader
  108. """
  109. def __init__(self, *args, **kwargs):
  110. super().__init__(*args, **kwargs)
  111. object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
  112. self.iterator = super().__iter__()
  113. def __len__(self):
  114. return len(self.batch_sampler.sampler)
  115. def __iter__(self):
  116. for i in range(len(self)):
  117. yield next(self.iterator)
  118. class _RepeatSampler(object):
  119. """ Sampler that repeats forever
  120. Args:
  121. sampler (Sampler)
  122. """
  123. def __init__(self, sampler):
  124. self.sampler = sampler
  125. def __iter__(self):
  126. while True:
  127. yield from iter(self.sampler)
  128. class LoadImages: # for inference
  129. def __init__(self, path, img_size=640, stride=32, auto=True):
  130. p = str(Path(path).resolve()) # os-agnostic absolute path
  131. if '*' in p:
  132. files = sorted(glob.glob(p, recursive=True)) # glob
  133. elif os.path.isdir(p):
  134. files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
  135. elif os.path.isfile(p):
  136. files = [p] # files
  137. else:
  138. raise Exception(f'ERROR: {p} does not exist')
  139. images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
  140. videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
  141. ni, nv = len(images), len(videos)
  142. self.img_size = img_size
  143. self.stride = stride
  144. self.files = images + videos
  145. self.nf = ni + nv # number of files
  146. self.video_flag = [False] * ni + [True] * nv
  147. self.mode = 'image'
  148. self.auto = auto
  149. if any(videos):
  150. self.new_video(videos[0]) # new video
  151. else:
  152. self.cap = None
  153. assert self.nf > 0, f'No images or videos found in {p}. ' \
  154. f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}'
  155. def __iter__(self):
  156. self.count = 0
  157. return self
  158. def __next__(self):
  159. if self.count == self.nf:
  160. raise StopIteration
  161. path = self.files[self.count]
  162. if self.video_flag[self.count]:
  163. # Read video
  164. self.mode = 'video'
  165. ret_val, img0 = self.cap.read()
  166. if not ret_val:
  167. self.count += 1
  168. self.cap.release()
  169. if self.count == self.nf: # last video
  170. raise StopIteration
  171. else:
  172. path = self.files[self.count]
  173. self.new_video(path)
  174. ret_val, img0 = self.cap.read()
  175. self.frame += 1
  176. print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ', end='')
  177. else:
  178. # Read image
  179. self.count += 1
  180. img0 = cv2.imread(path) # BGR
  181. assert img0 is not None, 'Image Not Found ' + path
  182. print(f'image {self.count}/{self.nf} {path}: ', end='')
  183. # Padded resize
  184. img = letterbox(img0, self.img_size, stride=self.stride, auto=self.auto)[0]
  185. # Convert
  186. img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
  187. img = np.ascontiguousarray(img)
  188. return path, img, img0, self.cap
  189. def new_video(self, path):
  190. self.frame = 0
  191. self.cap = cv2.VideoCapture(path)
  192. self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
  193. def __len__(self):
  194. return self.nf # number of files
  195. class LoadWebcam: # for inference
  196. def __init__(self, pipe='0', img_size=640, stride=32):
  197. self.img_size = img_size
  198. self.stride = stride
  199. self.pipe = eval(pipe) if pipe.isnumeric() else pipe
  200. self.cap = cv2.VideoCapture(self.pipe) # video capture object
  201. self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
  202. def __iter__(self):
  203. self.count = -1
  204. return self
  205. def __next__(self):
  206. self.count += 1
  207. if cv2.waitKey(1) == ord('q'): # q to quit
  208. self.cap.release()
  209. cv2.destroyAllWindows()
  210. raise StopIteration
  211. # Read frame
  212. ret_val, img0 = self.cap.read()
  213. img0 = cv2.flip(img0, 1) # flip left-right
  214. # Print
  215. assert ret_val, f'Camera Error {self.pipe}'
  216. img_path = 'webcam.jpg'
  217. print(f'webcam {self.count}: ', end='')
  218. # Padded resize
  219. img = letterbox(img0, self.img_size, stride=self.stride)[0]
  220. # Convert
  221. img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
  222. img = np.ascontiguousarray(img)
  223. return img_path, img, img0, None
  224. def __len__(self):
  225. return 0
  226. class LoadStreams: # multiple IP or RTSP cameras
  227. def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True):
  228. self.mode = 'stream'
  229. self.img_size = img_size
  230. self.stride = stride
  231. if os.path.isfile(sources):
  232. with open(sources, 'r') as f:
  233. sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]
  234. else:
  235. sources = [sources]
  236. n = len(sources)
  237. self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
  238. self.sources = [clean_str(x) for x in sources] # clean source names for later
  239. self.auto = auto
  240. for i, s in enumerate(sources): # index, source
  241. # Start thread to read frames from video stream
  242. print(f'{i + 1}/{n}: {s}... ', end='')
  243. if 'youtube.com/' in s or 'youtu.be/' in s: # if source is YouTube video
  244. check_requirements(('pafy', 'youtube_dl'))
  245. import pafy
  246. s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL
  247. s = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam
  248. cap = cv2.VideoCapture(s)
  249. assert cap.isOpened(), f'Failed to open {s}'
  250. w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  251. h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  252. self.fps[i] = max(cap.get(cv2.CAP_PROP_FPS) % 100, 0) or 30.0 # 30 FPS fallback
  253. self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback
  254. _, self.imgs[i] = cap.read() # guarantee first frame
  255. self.threads[i] = Thread(target=self.update, args=([i, cap]), daemon=True)
  256. print(f" success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")
  257. self.threads[i].start()
  258. print('') # newline
  259. # check for common shapes
  260. s = np.stack([letterbox(x, self.img_size, stride=self.stride, auto=self.auto)[0].shape for x in self.imgs])
  261. self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
  262. if not self.rect:
  263. print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
  264. def update(self, i, cap):
  265. # Read stream `i` frames in daemon thread
  266. n, f, read = 0, self.frames[i], 1 # frame number, frame array, inference every 'read' frame
  267. while cap.isOpened() and n < f:
  268. n += 1
  269. # _, self.imgs[index] = cap.read()
  270. cap.grab()
  271. if n % read == 0:
  272. success, im = cap.retrieve()
  273. self.imgs[i] = im if success else self.imgs[i] * 0
  274. time.sleep(1 / self.fps[i]) # wait time
  275. def __iter__(self):
  276. self.count = -1
  277. return self
  278. def __next__(self):
  279. self.count += 1
  280. if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit
  281. cv2.destroyAllWindows()
  282. raise StopIteration
  283. # Letterbox
  284. img0 = self.imgs.copy()
  285. img = [letterbox(x, self.img_size, stride=self.stride, auto=self.rect and self.auto)[0] for x in img0]
  286. # Stack
  287. img = np.stack(img, 0)
  288. # Convert
  289. img = img[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW
  290. img = np.ascontiguousarray(img)
  291. return self.sources, img, img0, None
  292. def __len__(self):
  293. return len(self.sources) # 1E12 frames = 32 streams at 30 FPS for 30 years
  294. def img2label_paths(img_paths):
  295. # Define label paths as a function of image paths
  296. sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings
  297. return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
  298. class LoadImagesAndLabels(Dataset): # for training/testing
  299. cache_version = 0.5 # dataset labels *.cache version
  300. def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
  301. cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
  302. self.img_size = img_size
  303. self.augment = augment
  304. self.hyp = hyp
  305. self.image_weights = image_weights
  306. self.rect = False if image_weights else rect
  307. self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
  308. self.mosaic_border = [-img_size // 2, -img_size // 2]
  309. self.stride = stride
  310. self.path = path
  311. self.albumentations = Albumentations() if augment else None
  312. try:
  313. f = [] # image files
  314. for p in path if isinstance(path, list) else [path]:
  315. p = Path(p) # os-agnostic
  316. if p.is_dir(): # dir
  317. f += glob.glob(str(p / '**' / '*.*'), recursive=True)
  318. # f = list(p.rglob('**/*.*')) # pathlib
  319. elif p.is_file(): # file
  320. with open(p, 'r') as t:
  321. t = t.read().strip().splitlines()
  322. parent = str(p.parent) + os.sep
  323. f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
  324. # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
  325. else:
  326. raise Exception(f'{prefix}{p} does not exist')
  327. self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS])
  328. # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
  329. assert self.img_files, f'{prefix}No images found'
  330. except Exception as e:
  331. raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')
  332. # Check cache
  333. self.label_files = img2label_paths(self.img_files) # labels
  334. cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')
  335. try:
  336. cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict
  337. assert cache['version'] == self.cache_version # same version
  338. assert cache['hash'] == get_hash(self.label_files + self.img_files) # same hash
  339. except:
  340. cache, exists = self.cache_labels(cache_path, prefix), False # cache
  341. # Display cache
  342. nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
  343. if exists:
  344. d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  345. tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results
  346. if cache['msgs']:
  347. logging.info('\n'.join(cache['msgs'])) # display warnings
  348. assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}'
  349. # Read cache
  350. [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items
  351. labels, shapes, self.segments = zip(*cache.values())
  352. self.labels = list(labels)
  353. self.shapes = np.array(shapes, dtype=np.float64)
  354. self.img_files = list(cache.keys()) # update
  355. self.label_files = img2label_paths(cache.keys()) # update
  356. if single_cls:
  357. for x in self.labels:
  358. x[:, 0] = 0
  359. n = len(shapes) # number of images
  360. bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
  361. nb = bi[-1] + 1 # number of batches
  362. self.batch = bi # batch index of image
  363. self.n = n
  364. self.indices = range(n)
  365. # Rectangular Training
  366. if self.rect:
  367. # Sort by aspect ratio
  368. s = self.shapes # wh
  369. ar = s[:, 1] / s[:, 0] # aspect ratio
  370. irect = ar.argsort()
  371. self.img_files = [self.img_files[i] for i in irect]
  372. self.label_files = [self.label_files[i] for i in irect]
  373. self.labels = [self.labels[i] for i in irect]
  374. self.shapes = s[irect] # wh
  375. ar = ar[irect]
  376. # Set training image shapes
  377. shapes = [[1, 1]] * nb
  378. for i in range(nb):
  379. ari = ar[bi == i]
  380. mini, maxi = ari.min(), ari.max()
  381. if maxi < 1:
  382. shapes[i] = [maxi, 1]
  383. elif mini > 1:
  384. shapes[i] = [1, 1 / mini]
  385. self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
  386. # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
  387. self.imgs, self.img_npy = [None] * n, [None] * n
  388. if cache_images:
  389. if cache_images == 'disk':
  390. self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + '_npy')
  391. self.img_npy = [self.im_cache_dir / Path(f).with_suffix('.npy').name for f in self.img_files]
  392. self.im_cache_dir.mkdir(parents=True, exist_ok=True)
  393. gb = 0 # Gigabytes of cached images
  394. self.img_hw0, self.img_hw = [None] * n, [None] * n
  395. results = ThreadPool(NUM_THREADS).imap(lambda x: load_image(*x), zip(repeat(self), range(n)))
  396. pbar = tqdm(enumerate(results), total=n)
  397. for i, x in pbar:
  398. if cache_images == 'disk':
  399. if not self.img_npy[i].exists():
  400. np.save(self.img_npy[i].as_posix(), x[0])
  401. gb += self.img_npy[i].stat().st_size
  402. else:
  403. self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)
  404. gb += self.imgs[i].nbytes
  405. pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})'
  406. pbar.close()
  407. def cache_labels(self, path=Path('./labels.cache'), prefix=''):
  408. # Cache dataset labels, check images and read shapes
  409. x = {} # dict
  410. nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages
  411. desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..."
  412. with Pool(NUM_THREADS) as pool:
  413. pbar = tqdm(pool.imap(verify_image_label, zip(self.img_files, self.label_files, repeat(prefix))),
  414. desc=desc, total=len(self.img_files))
  415. for im_file, l, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
  416. nm += nm_f
  417. nf += nf_f
  418. ne += ne_f
  419. nc += nc_f
  420. if im_file:
  421. x[im_file] = [l, shape, segments]
  422. if msg:
  423. msgs.append(msg)
  424. pbar.desc = f"{desc}{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  425. pbar.close()
  426. if msgs:
  427. logging.info('\n'.join(msgs))
  428. if nf == 0:
  429. logging.info(f'{prefix}WARNING: No labels found in {path}. See {HELP_URL}')
  430. x['hash'] = get_hash(self.label_files + self.img_files)
  431. x['results'] = nf, nm, ne, nc, len(self.img_files)
  432. x['msgs'] = msgs # warnings
  433. x['version'] = self.cache_version # cache version
  434. try:
  435. np.save(path, x) # save cache for next time
  436. path.with_suffix('.cache.npy').rename(path) # remove .npy suffix
  437. logging.info(f'{prefix}New cache created: {path}')
  438. except Exception as e:
  439. logging.info(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # path not writeable
  440. return x
  441. def __len__(self):
  442. return len(self.img_files)
  443. # def __iter__(self):
  444. # self.count = -1
  445. # print('ran dataset iter')
  446. # #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
  447. # return self
  448. def __getitem__(self, index):
  449. index = self.indices[index] # linear, shuffled, or image_weights
  450. hyp = self.hyp
  451. mosaic = self.mosaic and random.random() < hyp['mosaic']
  452. if mosaic:
  453. # Load mosaic
  454. img, labels = load_mosaic(self, index)
  455. shapes = None
  456. # MixUp augmentation
  457. if random.random() < hyp['mixup']:
  458. img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1)))
  459. else:
  460. # Load image
  461. img, (h0, w0), (h, w) = load_image(self, index)
  462. # Letterbox
  463. shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
  464. img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
  465. shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
  466. labels = self.labels[index].copy()
  467. if labels.size: # normalized xywh to pixel xyxy format
  468. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
  469. if self.augment:
  470. img, labels = random_perspective(img, labels,
  471. degrees=hyp['degrees'],
  472. translate=hyp['translate'],
  473. scale=hyp['scale'],
  474. shear=hyp['shear'],
  475. perspective=hyp['perspective'])
  476. nl = len(labels) # number of labels
  477. if nl:
  478. labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)
  479. if self.augment:
  480. # Albumentations
  481. img, labels = self.albumentations(img, labels)
  482. nl = len(labels) # update after albumentations
  483. # HSV color-space
  484. augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
  485. # Flip up-down
  486. if random.random() < hyp['flipud']:
  487. img = np.flipud(img)
  488. if nl:
  489. labels[:, 2] = 1 - labels[:, 2]
  490. # Flip left-right
  491. if random.random() < hyp['fliplr']:
  492. img = np.fliplr(img)
  493. if nl:
  494. labels[:, 1] = 1 - labels[:, 1]
  495. # Cutouts
  496. # labels = cutout(img, labels, p=0.5)
  497. labels_out = torch.zeros((nl, 6))
  498. if nl:
  499. labels_out[:, 1:] = torch.from_numpy(labels)
  500. # Convert
  501. img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
  502. img = np.ascontiguousarray(img)
  503. return torch.from_numpy(img), labels_out, self.img_files[index], shapes
  504. @staticmethod
  505. def collate_fn(batch):
  506. img, label, path, shapes = zip(*batch) # transposed
  507. for i, l in enumerate(label):
  508. l[:, 0] = i # add target image index for build_targets()
  509. return torch.stack(img, 0), torch.cat(label, 0), path, shapes
  510. @staticmethod
  511. def collate_fn4(batch):
  512. img, label, path, shapes = zip(*batch) # transposed
  513. n = len(shapes) // 4
  514. img4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]
  515. ho = torch.tensor([[0., 0, 0, 1, 0, 0]])
  516. wo = torch.tensor([[0., 0, 1, 0, 0, 0]])
  517. s = torch.tensor([[1, 1, .5, .5, .5, .5]]) # scale
  518. for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW
  519. i *= 4
  520. if random.random() < 0.5:
  521. im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2., mode='bilinear', align_corners=False)[
  522. 0].type(img[i].type())
  523. l = label[i]
  524. else:
  525. im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
  526. l = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
  527. img4.append(im)
  528. label4.append(l)
  529. for i, l in enumerate(label4):
  530. l[:, 0] = i # add target image index for build_targets()
  531. return torch.stack(img4, 0), torch.cat(label4, 0), path4, shapes4
  532. # Ancillary functions --------------------------------------------------------------------------------------------------
  533. def load_image(self, i):
  534. # loads 1 image from dataset index 'i', returns im, original hw, resized hw
  535. im = self.imgs[i]
  536. if im is None: # not cached in ram
  537. npy = self.img_npy[i]
  538. if npy and npy.exists(): # load npy
  539. im = np.load(npy)
  540. else: # read image
  541. path = self.img_files[i]
  542. im = cv2.imread(path) # BGR
  543. assert im is not None, 'Image Not Found ' + path
  544. h0, w0 = im.shape[:2] # orig hw
  545. r = self.img_size / max(h0, w0) # ratio
  546. if r != 1: # if sizes are not equal
  547. im = cv2.resize(im, (int(w0 * r), int(h0 * r)),
  548. interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR)
  549. return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
  550. else:
  551. return self.imgs[i], self.img_hw0[i], self.img_hw[i] # im, hw_original, hw_resized
  552. def load_mosaic(self, index):
  553. # loads images in a 4-mosaic
  554. labels4, segments4 = [], []
  555. s = self.img_size
  556. yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border] # mosaic center x, y
  557. indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
  558. random.shuffle(indices)
  559. for i, index in enumerate(indices):
  560. # Load image
  561. img, _, (h, w) = load_image(self, index)
  562. # place img in img4
  563. if i == 0: # top left
  564. img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
  565. x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
  566. x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
  567. elif i == 1: # top right
  568. x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
  569. x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
  570. elif i == 2: # bottom left
  571. x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
  572. x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
  573. elif i == 3: # bottom right
  574. x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
  575. x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
  576. img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
  577. padw = x1a - x1b
  578. padh = y1a - y1b
  579. # Labels
  580. labels, segments = self.labels[index].copy(), self.segments[index].copy()
  581. if labels.size:
  582. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
  583. segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
  584. labels4.append(labels)
  585. segments4.extend(segments)
  586. # Concat/clip labels
  587. labels4 = np.concatenate(labels4, 0)
  588. for x in (labels4[:, 1:], *segments4):
  589. np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
  590. # img4, labels4 = replicate(img4, labels4) # replicate
  591. # Augment
  592. img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])
  593. img4, labels4 = random_perspective(img4, labels4, segments4,
  594. degrees=self.hyp['degrees'],
  595. translate=self.hyp['translate'],
  596. scale=self.hyp['scale'],
  597. shear=self.hyp['shear'],
  598. perspective=self.hyp['perspective'],
  599. border=self.mosaic_border) # border to remove
  600. return img4, labels4
  601. def load_mosaic9(self, index):
  602. # loads images in a 9-mosaic
  603. labels9, segments9 = [], []
  604. s = self.img_size
  605. indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices
  606. random.shuffle(indices)
  607. for i, index in enumerate(indices):
  608. # Load image
  609. img, _, (h, w) = load_image(self, index)
  610. # place img in img9
  611. if i == 0: # center
  612. img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
  613. h0, w0 = h, w
  614. c = s, s, s + w, s + h # xmin, ymin, xmax, ymax (base) coordinates
  615. elif i == 1: # top
  616. c = s, s - h, s + w, s
  617. elif i == 2: # top right
  618. c = s + wp, s - h, s + wp + w, s
  619. elif i == 3: # right
  620. c = s + w0, s, s + w0 + w, s + h
  621. elif i == 4: # bottom right
  622. c = s + w0, s + hp, s + w0 + w, s + hp + h
  623. elif i == 5: # bottom
  624. c = s + w0 - w, s + h0, s + w0, s + h0 + h
  625. elif i == 6: # bottom left
  626. c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
  627. elif i == 7: # left
  628. c = s - w, s + h0 - h, s, s + h0
  629. elif i == 8: # top left
  630. c = s - w, s + h0 - hp - h, s, s + h0 - hp
  631. padx, pady = c[:2]
  632. x1, y1, x2, y2 = [max(x, 0) for x in c] # allocate coords
  633. # Labels
  634. labels, segments = self.labels[index].copy(), self.segments[index].copy()
  635. if labels.size:
  636. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format
  637. segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
  638. labels9.append(labels)
  639. segments9.extend(segments)
  640. # Image
  641. img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]
  642. hp, wp = h, w # height, width previous
  643. # Offset
  644. yc, xc = [int(random.uniform(0, s)) for _ in self.mosaic_border] # mosaic center x, y
  645. img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]
  646. # Concat/clip labels
  647. labels9 = np.concatenate(labels9, 0)
  648. labels9[:, [1, 3]] -= xc
  649. labels9[:, [2, 4]] -= yc
  650. c = np.array([xc, yc]) # centers
  651. segments9 = [x - c for x in segments9]
  652. for x in (labels9[:, 1:], *segments9):
  653. np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
  654. # img9, labels9 = replicate(img9, labels9) # replicate
  655. # Augment
  656. img9, labels9 = random_perspective(img9, labels9, segments9,
  657. degrees=self.hyp['degrees'],
  658. translate=self.hyp['translate'],
  659. scale=self.hyp['scale'],
  660. shear=self.hyp['shear'],
  661. perspective=self.hyp['perspective'],
  662. border=self.mosaic_border) # border to remove
  663. return img9, labels9
  664. def create_folder(path='./new'):
  665. # Create folder
  666. if os.path.exists(path):
  667. shutil.rmtree(path) # delete output folder
  668. os.makedirs(path) # make new output folder
  669. def flatten_recursive(path='../datasets/coco128'):
  670. # Flatten a recursive directory by bringing all files to top level
  671. new_path = Path(path + '_flat')
  672. create_folder(new_path)
  673. for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)):
  674. shutil.copyfile(file, new_path / Path(file).name)
  675. def extract_boxes(path='../datasets/coco128'): # from utils.datasets import *; extract_boxes()
  676. # Convert detection dataset into classification dataset, with one directory per class
  677. path = Path(path) # images dir
  678. shutil.rmtree(path / 'classifier') if (path / 'classifier').is_dir() else None # remove existing
  679. files = list(path.rglob('*.*'))
  680. n = len(files) # number of files
  681. for im_file in tqdm(files, total=n):
  682. if im_file.suffix[1:] in IMG_FORMATS:
  683. # image
  684. im = cv2.imread(str(im_file))[..., ::-1] # BGR to RGB
  685. h, w = im.shape[:2]
  686. # labels
  687. lb_file = Path(img2label_paths([str(im_file)])[0])
  688. if Path(lb_file).exists():
  689. with open(lb_file, 'r') as f:
  690. lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels
  691. for j, x in enumerate(lb):
  692. c = int(x[0]) # class
  693. f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename
  694. if not f.parent.is_dir():
  695. f.parent.mkdir(parents=True)
  696. b = x[1:] * [w, h, w, h] # box
  697. # b[2:] = b[2:].max() # rectangle to square
  698. b[2:] = b[2:] * 1.2 + 3 # pad
  699. b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int)
  700. b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image
  701. b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
  702. assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'
  703. def autosplit(path='../datasets/coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False):
  704. """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
  705. Usage: from utils.datasets import *; autosplit()
  706. Arguments
  707. path: Path to images directory
  708. weights: Train, val, test weights (list, tuple)
  709. annotated_only: Only use images with an annotated txt file
  710. """
  711. path = Path(path) # images dir
  712. files = sum([list(path.rglob(f"*.{img_ext}")) for img_ext in IMG_FORMATS], []) # image files only
  713. n = len(files) # number of files
  714. random.seed(0) # for reproducibility
  715. indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
  716. txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files
  717. [(path.parent / x).unlink(missing_ok=True) for x in txt] # remove existing
  718. print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)
  719. for i, img in tqdm(zip(indices, files), total=n):
  720. if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
  721. with open(path.parent / txt[i], 'a') as f:
  722. f.write('./' + img.relative_to(path.parent).as_posix() + '\n') # add image to txt file
  723. def verify_image_label(args):
  724. # Verify one image-label pair
  725. im_file, lb_file, prefix = args
  726. nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments
  727. try:
  728. # verify images
  729. im = Image.open(im_file)
  730. im.verify() # PIL verify
  731. shape = exif_size(im) # image size
  732. assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
  733. assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}'
  734. if im.format.lower() in ('jpg', 'jpeg'):
  735. with open(im_file, 'rb') as f:
  736. f.seek(-2, 2)
  737. if f.read() != b'\xff\xd9': # corrupt JPEG
  738. Image.open(im_file).save(im_file, format='JPEG', subsampling=0, quality=100) # re-save image
  739. msg = f'{prefix}WARNING: corrupt JPEG restored and saved {im_file}'
  740. # verify labels
  741. if os.path.isfile(lb_file):
  742. nf = 1 # label found
  743. with open(lb_file, 'r') as f:
  744. l = [x.split() for x in f.read().strip().splitlines() if len(x)]
  745. if any([len(x) > 8 for x in l]): # is segment
  746. classes = np.array([x[0] for x in l], dtype=np.float32)
  747. segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
  748. l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
  749. l = np.array(l, dtype=np.float32)
  750. if len(l):
  751. assert l.shape[1] == 5, 'labels require 5 columns each'
  752. assert (l >= 0).all(), 'negative labels'
  753. assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels'
  754. assert np.unique(l, axis=0).shape[0] == l.shape[0], 'duplicate labels'
  755. else:
  756. ne = 1 # label empty
  757. l = np.zeros((0, 5), dtype=np.float32)
  758. else:
  759. nm = 1 # label missing
  760. l = np.zeros((0, 5), dtype=np.float32)
  761. return im_file, l, shape, segments, nm, nf, ne, nc, msg
  762. except Exception as e:
  763. nc = 1
  764. msg = f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}'
  765. return [None, None, None, None, nm, nf, ne, nc, msg]
  766. def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profile=False, hub=False):
  767. """ Return dataset statistics dictionary with images and instances counts per split per class
  768. To run in parent directory: export PYTHONPATH="$PWD/yolov5"
  769. Usage1: from utils.datasets import *; dataset_stats('coco128.yaml', autodownload=True)
  770. Usage2: from utils.datasets import *; dataset_stats('../datasets/coco128_with_yaml.zip')
  771. Arguments
  772. path: Path to data.yaml or data.zip (with data.yaml inside data.zip)
  773. autodownload: Attempt to download dataset if not found locally
  774. verbose: Print stats dictionary
  775. """
  776. def round_labels(labels):
  777. # Update labels to integer class and 6 decimal place floats
  778. return [[int(c), *[round(x, 4) for x in points]] for c, *points in labels]
  779. def unzip(path):
  780. # Unzip data.zip TODO: CONSTRAINT: path/to/abc.zip MUST unzip to 'path/to/abc/'
  781. if str(path).endswith('.zip'): # path is data.zip
  782. assert Path(path).is_file(), f'Error unzipping {path}, file not found'
  783. assert os.system(f'unzip -q {path} -d {path.parent}') == 0, f'Error unzipping {path}'
  784. dir = path.with_suffix('') # dataset directory
  785. return True, str(dir), next(dir.rglob('*.yaml')) # zipped, data_dir, yaml_path
  786. else: # path is data.yaml
  787. return False, None, path
  788. def hub_ops(f, max_dim=1920):
  789. # HUB ops for 1 image 'f'
  790. im = Image.open(f)
  791. r = max_dim / max(im.height, im.width) # ratio
  792. if r < 1.0: # image too large
  793. im = im.resize((int(im.width * r), int(im.height * r)))
  794. im.save(im_dir / Path(f).name, quality=75) # save
  795. zipped, data_dir, yaml_path = unzip(Path(path))
  796. with open(check_yaml(yaml_path), errors='ignore') as f:
  797. data = yaml.safe_load(f) # data dict
  798. if zipped:
  799. data['path'] = data_dir # TODO: should this be dir.resolve()?
  800. check_dataset(data, autodownload) # download dataset if missing
  801. hub_dir = Path(data['path'] + ('-hub' if hub else ''))
  802. stats = {'nc': data['nc'], 'names': data['names']} # statistics dictionary
  803. for split in 'train', 'val', 'test':
  804. if data.get(split) is None:
  805. stats[split] = None # i.e. no test set
  806. continue
  807. x = []
  808. dataset = LoadImagesAndLabels(data[split]) # load dataset
  809. for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'):
  810. x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc']))
  811. x = np.array(x) # shape(128x80)
  812. stats[split] = {'instance_stats': {'total': int(x.sum()), 'per_class': x.sum(0).tolist()},
  813. 'image_stats': {'total': dataset.n, 'unlabelled': int(np.all(x == 0, 1).sum()),
  814. 'per_class': (x > 0).sum(0).tolist()},
  815. 'labels': [{str(Path(k).name): round_labels(v.tolist())} for k, v in
  816. zip(dataset.img_files, dataset.labels)]}
  817. if hub:
  818. im_dir = hub_dir / 'images'
  819. im_dir.mkdir(parents=True, exist_ok=True)
  820. for _ in tqdm(ThreadPool(NUM_THREADS).imap(hub_ops, dataset.img_files), total=dataset.n, desc='HUB Ops'):
  821. pass
  822. # Profile
  823. stats_path = hub_dir / 'stats.json'
  824. if profile:
  825. for _ in range(1):
  826. file = stats_path.with_suffix('.npy')
  827. t1 = time.time()
  828. np.save(file, stats)
  829. t2 = time.time()
  830. x = np.load(file, allow_pickle=True)
  831. print(f'stats.npy times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
  832. file = stats_path.with_suffix('.json')
  833. t1 = time.time()
  834. with open(file, 'w') as f:
  835. json.dump(stats, f) # save stats *.json
  836. t2 = time.time()
  837. with open(file, 'r') as f:
  838. x = json.load(f) # load hyps dict
  839. print(f'stats.json times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
  840. # Save, print and return
  841. if hub:
  842. print(f'Saving {stats_path.resolve()}...')
  843. with open(stats_path, 'w') as f:
  844. json.dump(stats, f) # save stats.json
  845. if verbose:
  846. print(json.dumps(stats, indent=2, sort_keys=False))
  847. return stats