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.

997 lines
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_requirements, check_file, check_dataset, xywh2xyxy, xywhn2xyxy, xyxy2xywhn, \
  27. xyn2xy, segments2boxes, clean_str
  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).absolute()) # 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. def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
  300. cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
  301. self.img_size = img_size
  302. self.augment = augment
  303. self.hyp = hyp
  304. self.image_weights = image_weights
  305. self.rect = False if image_weights else rect
  306. self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
  307. self.mosaic_border = [-img_size // 2, -img_size // 2]
  308. self.stride = stride
  309. self.path = path
  310. self.albumentations = Albumentations() if augment else None
  311. try:
  312. f = [] # image files
  313. for p in path if isinstance(path, list) else [path]:
  314. p = Path(p) # os-agnostic
  315. if p.is_dir(): # dir
  316. f += glob.glob(str(p / '**' / '*.*'), recursive=True)
  317. # f = list(p.rglob('**/*.*')) # pathlib
  318. elif p.is_file(): # file
  319. with open(p, 'r') as t:
  320. t = t.read().strip().splitlines()
  321. parent = str(p.parent) + os.sep
  322. f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
  323. # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
  324. else:
  325. raise Exception(f'{prefix}{p} does not exist')
  326. self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS])
  327. # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
  328. assert self.img_files, f'{prefix}No images found'
  329. except Exception as e:
  330. raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')
  331. # Check cache
  332. self.label_files = img2label_paths(self.img_files) # labels
  333. cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')
  334. try:
  335. cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict
  336. assert cache['version'] == 0.4 and cache['hash'] == get_hash(self.label_files + self.img_files)
  337. except:
  338. cache, exists = self.cache_labels(cache_path, prefix), False # cache
  339. # Display cache
  340. nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
  341. if exists:
  342. d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  343. tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results
  344. if cache['msgs']:
  345. logging.info('\n'.join(cache['msgs'])) # display warnings
  346. assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}'
  347. # Read cache
  348. [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items
  349. labels, shapes, self.segments = zip(*cache.values())
  350. self.labels = list(labels)
  351. self.shapes = np.array(shapes, dtype=np.float64)
  352. self.img_files = list(cache.keys()) # update
  353. self.label_files = img2label_paths(cache.keys()) # update
  354. if single_cls:
  355. for x in self.labels:
  356. x[:, 0] = 0
  357. n = len(shapes) # number of images
  358. bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
  359. nb = bi[-1] + 1 # number of batches
  360. self.batch = bi # batch index of image
  361. self.n = n
  362. self.indices = range(n)
  363. # Rectangular Training
  364. if self.rect:
  365. # Sort by aspect ratio
  366. s = self.shapes # wh
  367. ar = s[:, 1] / s[:, 0] # aspect ratio
  368. irect = ar.argsort()
  369. self.img_files = [self.img_files[i] for i in irect]
  370. self.label_files = [self.label_files[i] for i in irect]
  371. self.labels = [self.labels[i] for i in irect]
  372. self.shapes = s[irect] # wh
  373. ar = ar[irect]
  374. # Set training image shapes
  375. shapes = [[1, 1]] * nb
  376. for i in range(nb):
  377. ari = ar[bi == i]
  378. mini, maxi = ari.min(), ari.max()
  379. if maxi < 1:
  380. shapes[i] = [maxi, 1]
  381. elif mini > 1:
  382. shapes[i] = [1, 1 / mini]
  383. self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
  384. # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
  385. self.imgs, self.img_npy = [None] * n, [None] * n
  386. if cache_images:
  387. if cache_images == 'disk':
  388. self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + '_npy')
  389. self.img_npy = [self.im_cache_dir / Path(f).with_suffix('.npy').name for f in self.img_files]
  390. self.im_cache_dir.mkdir(parents=True, exist_ok=True)
  391. gb = 0 # Gigabytes of cached images
  392. self.img_hw0, self.img_hw = [None] * n, [None] * n
  393. results = ThreadPool(NUM_THREADS).imap(lambda x: load_image(*x), zip(repeat(self), range(n)))
  394. pbar = tqdm(enumerate(results), total=n)
  395. for i, x in pbar:
  396. if cache_images == 'disk':
  397. if not self.img_npy[i].exists():
  398. np.save(self.img_npy[i].as_posix(), x[0])
  399. gb += self.img_npy[i].stat().st_size
  400. else:
  401. self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)
  402. gb += self.imgs[i].nbytes
  403. pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})'
  404. pbar.close()
  405. def cache_labels(self, path=Path('./labels.cache'), prefix=''):
  406. # Cache dataset labels, check images and read shapes
  407. x = {} # dict
  408. nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages
  409. desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..."
  410. with Pool(NUM_THREADS) as pool:
  411. pbar = tqdm(pool.imap_unordered(verify_image_label, zip(self.img_files, self.label_files, repeat(prefix))),
  412. desc=desc, total=len(self.img_files))
  413. for im_file, l, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
  414. nm += nm_f
  415. nf += nf_f
  416. ne += ne_f
  417. nc += nc_f
  418. if im_file:
  419. x[im_file] = [l, shape, segments]
  420. if msg:
  421. msgs.append(msg)
  422. pbar.desc = f"{desc}{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  423. pbar.close()
  424. if msgs:
  425. logging.info('\n'.join(msgs))
  426. if nf == 0:
  427. logging.info(f'{prefix}WARNING: No labels found in {path}. See {HELP_URL}')
  428. x['hash'] = get_hash(self.label_files + self.img_files)
  429. x['results'] = nf, nm, ne, nc, len(self.img_files)
  430. x['msgs'] = msgs # warnings
  431. x['version'] = 0.4 # cache version
  432. try:
  433. np.save(path, x) # save cache for next time
  434. path.with_suffix('.cache.npy').rename(path) # remove .npy suffix
  435. logging.info(f'{prefix}New cache created: {path}')
  436. except Exception as e:
  437. logging.info(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # path not writeable
  438. return x
  439. def __len__(self):
  440. return len(self.img_files)
  441. # def __iter__(self):
  442. # self.count = -1
  443. # print('ran dataset iter')
  444. # #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
  445. # return self
  446. def __getitem__(self, index):
  447. index = self.indices[index] # linear, shuffled, or image_weights
  448. hyp = self.hyp
  449. mosaic = self.mosaic and random.random() < hyp['mosaic']
  450. if mosaic:
  451. # Load mosaic
  452. img, labels = load_mosaic(self, index)
  453. shapes = None
  454. # MixUp augmentation
  455. if random.random() < hyp['mixup']:
  456. img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1)))
  457. else:
  458. # Load image
  459. img, (h0, w0), (h, w) = load_image(self, index)
  460. # Letterbox
  461. shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
  462. img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
  463. shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
  464. labels = self.labels[index].copy()
  465. if labels.size: # normalized xywh to pixel xyxy format
  466. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
  467. if self.augment:
  468. img, labels = random_perspective(img, labels,
  469. degrees=hyp['degrees'],
  470. translate=hyp['translate'],
  471. scale=hyp['scale'],
  472. shear=hyp['shear'],
  473. perspective=hyp['perspective'])
  474. nl = len(labels) # number of labels
  475. if nl:
  476. labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)
  477. if self.augment:
  478. # Albumentations
  479. img, labels = self.albumentations(img, labels)
  480. nl = len(labels) # update after albumentations
  481. # HSV color-space
  482. augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
  483. # Flip up-down
  484. if random.random() < hyp['flipud']:
  485. img = np.flipud(img)
  486. if nl:
  487. labels[:, 2] = 1 - labels[:, 2]
  488. # Flip left-right
  489. if random.random() < hyp['fliplr']:
  490. img = np.fliplr(img)
  491. if nl:
  492. labels[:, 1] = 1 - labels[:, 1]
  493. # Cutouts
  494. # labels = cutout(img, labels, p=0.5)
  495. labels_out = torch.zeros((nl, 6))
  496. if nl:
  497. labels_out[:, 1:] = torch.from_numpy(labels)
  498. # Convert
  499. img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
  500. img = np.ascontiguousarray(img)
  501. return torch.from_numpy(img), labels_out, self.img_files[index], shapes
  502. @staticmethod
  503. def collate_fn(batch):
  504. img, label, path, shapes = zip(*batch) # transposed
  505. for i, l in enumerate(label):
  506. l[:, 0] = i # add target image index for build_targets()
  507. return torch.stack(img, 0), torch.cat(label, 0), path, shapes
  508. @staticmethod
  509. def collate_fn4(batch):
  510. img, label, path, shapes = zip(*batch) # transposed
  511. n = len(shapes) // 4
  512. img4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]
  513. ho = torch.tensor([[0., 0, 0, 1, 0, 0]])
  514. wo = torch.tensor([[0., 0, 1, 0, 0, 0]])
  515. s = torch.tensor([[1, 1, .5, .5, .5, .5]]) # scale
  516. for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW
  517. i *= 4
  518. if random.random() < 0.5:
  519. im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2., mode='bilinear', align_corners=False)[
  520. 0].type(img[i].type())
  521. l = label[i]
  522. else:
  523. im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
  524. l = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
  525. img4.append(im)
  526. label4.append(l)
  527. for i, l in enumerate(label4):
  528. l[:, 0] = i # add target image index for build_targets()
  529. return torch.stack(img4, 0), torch.cat(label4, 0), path4, shapes4
  530. # Ancillary functions --------------------------------------------------------------------------------------------------
  531. def load_image(self, i):
  532. # loads 1 image from dataset index 'i', returns im, original hw, resized hw
  533. im = self.imgs[i]
  534. if im is None: # not cached in ram
  535. npy = self.img_npy[i]
  536. if npy and npy.exists(): # load npy
  537. im = np.load(npy)
  538. else: # read image
  539. path = self.img_files[i]
  540. im = cv2.imread(path) # BGR
  541. assert im is not None, 'Image Not Found ' + path
  542. h0, w0 = im.shape[:2] # orig hw
  543. r = self.img_size / max(h0, w0) # ratio
  544. if r != 1: # if sizes are not equal
  545. im = cv2.resize(im, (int(w0 * r), int(h0 * r)),
  546. interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR)
  547. return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
  548. else:
  549. return self.imgs[i], self.img_hw0[i], self.img_hw[i] # im, hw_original, hw_resized
  550. def load_mosaic(self, index):
  551. # loads images in a 4-mosaic
  552. labels4, segments4 = [], []
  553. s = self.img_size
  554. yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border] # mosaic center x, y
  555. indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
  556. for i, index in enumerate(indices):
  557. # Load image
  558. img, _, (h, w) = load_image(self, index)
  559. # place img in img4
  560. if i == 0: # top left
  561. img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
  562. x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
  563. x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
  564. elif i == 1: # top right
  565. x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
  566. x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
  567. elif i == 2: # bottom left
  568. x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
  569. x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
  570. elif i == 3: # bottom right
  571. x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
  572. x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
  573. img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
  574. padw = x1a - x1b
  575. padh = y1a - y1b
  576. # Labels
  577. labels, segments = self.labels[index].copy(), self.segments[index].copy()
  578. if labels.size:
  579. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
  580. segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
  581. labels4.append(labels)
  582. segments4.extend(segments)
  583. # Concat/clip labels
  584. labels4 = np.concatenate(labels4, 0)
  585. for x in (labels4[:, 1:], *segments4):
  586. np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
  587. # img4, labels4 = replicate(img4, labels4) # replicate
  588. # Augment
  589. img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])
  590. img4, labels4 = random_perspective(img4, labels4, segments4,
  591. degrees=self.hyp['degrees'],
  592. translate=self.hyp['translate'],
  593. scale=self.hyp['scale'],
  594. shear=self.hyp['shear'],
  595. perspective=self.hyp['perspective'],
  596. border=self.mosaic_border) # border to remove
  597. return img4, labels4
  598. def load_mosaic9(self, index):
  599. # loads images in a 9-mosaic
  600. labels9, segments9 = [], []
  601. s = self.img_size
  602. indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices
  603. for i, index in enumerate(indices):
  604. # Load image
  605. img, _, (h, w) = load_image(self, index)
  606. # place img in img9
  607. if i == 0: # center
  608. img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
  609. h0, w0 = h, w
  610. c = s, s, s + w, s + h # xmin, ymin, xmax, ymax (base) coordinates
  611. elif i == 1: # top
  612. c = s, s - h, s + w, s
  613. elif i == 2: # top right
  614. c = s + wp, s - h, s + wp + w, s
  615. elif i == 3: # right
  616. c = s + w0, s, s + w0 + w, s + h
  617. elif i == 4: # bottom right
  618. c = s + w0, s + hp, s + w0 + w, s + hp + h
  619. elif i == 5: # bottom
  620. c = s + w0 - w, s + h0, s + w0, s + h0 + h
  621. elif i == 6: # bottom left
  622. c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
  623. elif i == 7: # left
  624. c = s - w, s + h0 - h, s, s + h0
  625. elif i == 8: # top left
  626. c = s - w, s + h0 - hp - h, s, s + h0 - hp
  627. padx, pady = c[:2]
  628. x1, y1, x2, y2 = [max(x, 0) for x in c] # allocate coords
  629. # Labels
  630. labels, segments = self.labels[index].copy(), self.segments[index].copy()
  631. if labels.size:
  632. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format
  633. segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
  634. labels9.append(labels)
  635. segments9.extend(segments)
  636. # Image
  637. img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]
  638. hp, wp = h, w # height, width previous
  639. # Offset
  640. yc, xc = [int(random.uniform(0, s)) for _ in self.mosaic_border] # mosaic center x, y
  641. img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]
  642. # Concat/clip labels
  643. labels9 = np.concatenate(labels9, 0)
  644. labels9[:, [1, 3]] -= xc
  645. labels9[:, [2, 4]] -= yc
  646. c = np.array([xc, yc]) # centers
  647. segments9 = [x - c for x in segments9]
  648. for x in (labels9[:, 1:], *segments9):
  649. np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
  650. # img9, labels9 = replicate(img9, labels9) # replicate
  651. # Augment
  652. img9, labels9 = random_perspective(img9, labels9, segments9,
  653. degrees=self.hyp['degrees'],
  654. translate=self.hyp['translate'],
  655. scale=self.hyp['scale'],
  656. shear=self.hyp['shear'],
  657. perspective=self.hyp['perspective'],
  658. border=self.mosaic_border) # border to remove
  659. return img9, labels9
  660. def create_folder(path='./new'):
  661. # Create folder
  662. if os.path.exists(path):
  663. shutil.rmtree(path) # delete output folder
  664. os.makedirs(path) # make new output folder
  665. def flatten_recursive(path='../datasets/coco128'):
  666. # Flatten a recursive directory by bringing all files to top level
  667. new_path = Path(path + '_flat')
  668. create_folder(new_path)
  669. for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)):
  670. shutil.copyfile(file, new_path / Path(file).name)
  671. def extract_boxes(path='../datasets/coco128'): # from utils.datasets import *; extract_boxes()
  672. # Convert detection dataset into classification dataset, with one directory per class
  673. path = Path(path) # images dir
  674. shutil.rmtree(path / 'classifier') if (path / 'classifier').is_dir() else None # remove existing
  675. files = list(path.rglob('*.*'))
  676. n = len(files) # number of files
  677. for im_file in tqdm(files, total=n):
  678. if im_file.suffix[1:] in IMG_FORMATS:
  679. # image
  680. im = cv2.imread(str(im_file))[..., ::-1] # BGR to RGB
  681. h, w = im.shape[:2]
  682. # labels
  683. lb_file = Path(img2label_paths([str(im_file)])[0])
  684. if Path(lb_file).exists():
  685. with open(lb_file, 'r') as f:
  686. lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels
  687. for j, x in enumerate(lb):
  688. c = int(x[0]) # class
  689. f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename
  690. if not f.parent.is_dir():
  691. f.parent.mkdir(parents=True)
  692. b = x[1:] * [w, h, w, h] # box
  693. # b[2:] = b[2:].max() # rectangle to square
  694. b[2:] = b[2:] * 1.2 + 3 # pad
  695. b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int)
  696. b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image
  697. b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
  698. assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'
  699. def autosplit(path='../datasets/coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False):
  700. """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
  701. Usage: from utils.datasets import *; autosplit()
  702. Arguments
  703. path: Path to images directory
  704. weights: Train, val, test weights (list, tuple)
  705. annotated_only: Only use images with an annotated txt file
  706. """
  707. path = Path(path) # images dir
  708. files = sum([list(path.rglob(f"*.{img_ext}")) for img_ext in IMG_FORMATS], []) # image files only
  709. n = len(files) # number of files
  710. random.seed(0) # for reproducibility
  711. indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
  712. txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files
  713. [(path.parent / x).unlink(missing_ok=True) for x in txt] # remove existing
  714. print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)
  715. for i, img in tqdm(zip(indices, files), total=n):
  716. if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
  717. with open(path.parent / txt[i], 'a') as f:
  718. f.write('./' + img.relative_to(path.parent).as_posix() + '\n') # add image to txt file
  719. def verify_image_label(args):
  720. # Verify one image-label pair
  721. im_file, lb_file, prefix = args
  722. nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments
  723. try:
  724. # verify images
  725. im = Image.open(im_file)
  726. im.verify() # PIL verify
  727. shape = exif_size(im) # image size
  728. assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
  729. assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}'
  730. if im.format.lower() in ('jpg', 'jpeg'):
  731. with open(im_file, 'rb') as f:
  732. f.seek(-2, 2)
  733. if f.read() != b'\xff\xd9': # corrupt JPEG
  734. Image.open(im_file).save(im_file, format='JPEG', subsampling=0, quality=100) # re-save image
  735. msg = f'{prefix}WARNING: corrupt JPEG restored and saved {im_file}'
  736. # verify labels
  737. if os.path.isfile(lb_file):
  738. nf = 1 # label found
  739. with open(lb_file, 'r') as f:
  740. l = [x.split() for x in f.read().strip().splitlines() if len(x)]
  741. if any([len(x) > 8 for x in l]): # is segment
  742. classes = np.array([x[0] for x in l], dtype=np.float32)
  743. segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
  744. l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
  745. l = np.array(l, dtype=np.float32)
  746. if len(l):
  747. assert l.shape[1] == 5, 'labels require 5 columns each'
  748. assert (l >= 0).all(), 'negative labels'
  749. assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels'
  750. assert np.unique(l, axis=0).shape[0] == l.shape[0], 'duplicate labels'
  751. else:
  752. ne = 1 # label empty
  753. l = np.zeros((0, 5), dtype=np.float32)
  754. else:
  755. nm = 1 # label missing
  756. l = np.zeros((0, 5), dtype=np.float32)
  757. return im_file, l, shape, segments, nm, nf, ne, nc, msg
  758. except Exception as e:
  759. nc = 1
  760. msg = f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}'
  761. return [None, None, None, None, nm, nf, ne, nc, msg]
  762. def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profile=False, hub=False):
  763. """ Return dataset statistics dictionary with images and instances counts per split per class
  764. To run in parent directory: export PYTHONPATH="$PWD/yolov5"
  765. Usage1: from utils.datasets import *; dataset_stats('coco128.yaml', autodownload=True)
  766. Usage2: from utils.datasets import *; dataset_stats('../datasets/coco128_with_yaml.zip')
  767. Arguments
  768. path: Path to data.yaml or data.zip (with data.yaml inside data.zip)
  769. autodownload: Attempt to download dataset if not found locally
  770. verbose: Print stats dictionary
  771. """
  772. def round_labels(labels):
  773. # Update labels to integer class and 6 decimal place floats
  774. return [[int(c), *[round(x, 4) for x in points]] for c, *points in labels]
  775. def unzip(path):
  776. # Unzip data.zip TODO: CONSTRAINT: path/to/abc.zip MUST unzip to 'path/to/abc/'
  777. if str(path).endswith('.zip'): # path is data.zip
  778. assert Path(path).is_file(), f'Error unzipping {path}, file not found'
  779. assert os.system(f'unzip -q {path} -d {path.parent}') == 0, f'Error unzipping {path}'
  780. dir = path.with_suffix('') # dataset directory
  781. return True, str(dir), next(dir.rglob('*.yaml')) # zipped, data_dir, yaml_path
  782. else: # path is data.yaml
  783. return False, None, path
  784. def hub_ops(f, max_dim=1920):
  785. # HUB ops for 1 image 'f'
  786. im = Image.open(f)
  787. r = max_dim / max(im.height, im.width) # ratio
  788. if r < 1.0: # image too large
  789. im = im.resize((int(im.width * r), int(im.height * r)))
  790. im.save(im_dir / Path(f).name, quality=75) # save
  791. zipped, data_dir, yaml_path = unzip(Path(path))
  792. with open(check_file(yaml_path), errors='ignore') as f:
  793. data = yaml.safe_load(f) # data dict
  794. if zipped:
  795. data['path'] = data_dir # TODO: should this be dir.resolve()?
  796. check_dataset(data, autodownload) # download dataset if missing
  797. hub_dir = Path(data['path'] + ('-hub' if hub else ''))
  798. stats = {'nc': data['nc'], 'names': data['names']} # statistics dictionary
  799. for split in 'train', 'val', 'test':
  800. if data.get(split) is None:
  801. stats[split] = None # i.e. no test set
  802. continue
  803. x = []
  804. dataset = LoadImagesAndLabels(data[split]) # load dataset
  805. for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'):
  806. x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc']))
  807. x = np.array(x) # shape(128x80)
  808. stats[split] = {'instance_stats': {'total': int(x.sum()), 'per_class': x.sum(0).tolist()},
  809. 'image_stats': {'total': dataset.n, 'unlabelled': int(np.all(x == 0, 1).sum()),
  810. 'per_class': (x > 0).sum(0).tolist()},
  811. 'labels': [{str(Path(k).name): round_labels(v.tolist())} for k, v in
  812. zip(dataset.img_files, dataset.labels)]}
  813. if hub:
  814. im_dir = hub_dir / 'images'
  815. im_dir.mkdir(parents=True, exist_ok=True)
  816. for _ in tqdm(ThreadPool(NUM_THREADS).imap(hub_ops, dataset.img_files), total=dataset.n, desc='HUB Ops'):
  817. pass
  818. # Profile
  819. stats_path = hub_dir / 'stats.json'
  820. if profile:
  821. for _ in range(1):
  822. file = stats_path.with_suffix('.npy')
  823. t1 = time.time()
  824. np.save(file, stats)
  825. t2 = time.time()
  826. x = np.load(file, allow_pickle=True)
  827. print(f'stats.npy times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
  828. file = stats_path.with_suffix('.json')
  829. t1 = time.time()
  830. with open(file, 'w') as f:
  831. json.dump(stats, f) # save stats *.json
  832. t2 = time.time()
  833. with open(file, 'r') as f:
  834. x = json.load(f) # load hyps dict
  835. print(f'stats.json times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
  836. # Save, print and return
  837. if hub:
  838. print(f'Saving {stats_path.resolve()}...')
  839. with open(stats_path, 'w') as f:
  840. json.dump(stats, f) # save stats.json
  841. if verbose:
  842. print(json.dumps(stats, indent=2, sort_keys=False))
  843. return stats