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.

1056 lines
43KB

  1. # Dataset utils and dataloaders
  2. import glob
  3. import logging
  4. import math
  5. import os
  6. import random
  7. import shutil
  8. import time
  9. from itertools import repeat
  10. from multiprocessing.pool import ThreadPool
  11. from pathlib import Path
  12. from threading import Thread
  13. import cv2
  14. import numpy as np
  15. import torch
  16. import torch.nn.functional as F
  17. from PIL import Image, ExifTags
  18. from torch.utils.data import Dataset
  19. from tqdm import tqdm
  20. from utils.general import xyxy2xywh, xywh2xyxy, xywhn2xyxy, xyn2xy, segment2box, segments2boxes, resample_segments, \
  21. clean_str
  22. from utils.torch_utils import torch_distributed_zero_first
  23. # Parameters
  24. help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
  25. img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp'] # acceptable image suffixes
  26. vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
  27. logger = logging.getLogger(__name__)
  28. # Get orientation exif tag
  29. for orientation in ExifTags.TAGS.keys():
  30. if ExifTags.TAGS[orientation] == 'Orientation':
  31. break
  32. def get_hash(files):
  33. # Returns a single hash value of a list of files
  34. return sum(os.path.getsize(f) for f in files if os.path.isfile(f))
  35. def exif_size(img):
  36. # Returns exif-corrected PIL size
  37. s = img.size # (width, height)
  38. try:
  39. rotation = dict(img._getexif().items())[orientation]
  40. if rotation == 6: # rotation 270
  41. s = (s[1], s[0])
  42. elif rotation == 8: # rotation 90
  43. s = (s[1], s[0])
  44. except:
  45. pass
  46. return s
  47. def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
  48. rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix=''):
  49. # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
  50. with torch_distributed_zero_first(rank):
  51. dataset = LoadImagesAndLabels(path, imgsz, batch_size,
  52. augment=augment, # augment images
  53. hyp=hyp, # augmentation hyperparameters
  54. rect=rect, # rectangular training
  55. cache_images=cache,
  56. single_cls=opt.single_cls,
  57. stride=int(stride),
  58. pad=pad,
  59. image_weights=image_weights,
  60. prefix=prefix)
  61. batch_size = min(batch_size, len(dataset))
  62. nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers]) # number of workers
  63. sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None
  64. loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader
  65. # Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader()
  66. dataloader = loader(dataset,
  67. batch_size=batch_size,
  68. num_workers=nw,
  69. sampler=sampler,
  70. pin_memory=True,
  71. collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn)
  72. return dataloader, dataset
  73. class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):
  74. """ Dataloader that reuses workers
  75. Uses same syntax as vanilla DataLoader
  76. """
  77. def __init__(self, *args, **kwargs):
  78. super().__init__(*args, **kwargs)
  79. object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
  80. self.iterator = super().__iter__()
  81. def __len__(self):
  82. return len(self.batch_sampler.sampler)
  83. def __iter__(self):
  84. for i in range(len(self)):
  85. yield next(self.iterator)
  86. class _RepeatSampler(object):
  87. """ Sampler that repeats forever
  88. Args:
  89. sampler (Sampler)
  90. """
  91. def __init__(self, sampler):
  92. self.sampler = sampler
  93. def __iter__(self):
  94. while True:
  95. yield from iter(self.sampler)
  96. class LoadImages: # for inference
  97. def __init__(self, path, img_size=640, stride=32):
  98. p = str(Path(path).absolute()) # os-agnostic absolute path
  99. if '*' in p:
  100. files = sorted(glob.glob(p, recursive=True)) # glob
  101. elif os.path.isdir(p):
  102. files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
  103. elif os.path.isfile(p):
  104. files = [p] # files
  105. else:
  106. raise Exception(f'ERROR: {p} does not exist')
  107. images = [x for x in files if x.split('.')[-1].lower() in img_formats]
  108. videos = [x for x in files if x.split('.')[-1].lower() in vid_formats]
  109. ni, nv = len(images), len(videos)
  110. self.img_size = img_size
  111. self.stride = stride
  112. self.files = images + videos
  113. self.nf = ni + nv # number of files
  114. self.video_flag = [False] * ni + [True] * nv
  115. self.mode = 'image'
  116. if any(videos):
  117. self.new_video(videos[0]) # new video
  118. else:
  119. self.cap = None
  120. assert self.nf > 0, f'No images or videos found in {p}. ' \
  121. f'Supported formats are:\nimages: {img_formats}\nvideos: {vid_formats}'
  122. def __iter__(self):
  123. self.count = 0
  124. return self
  125. def __next__(self):
  126. if self.count == self.nf:
  127. raise StopIteration
  128. path = self.files[self.count]
  129. if self.video_flag[self.count]:
  130. # Read video
  131. self.mode = 'video'
  132. ret_val, img0 = self.cap.read()
  133. if not ret_val:
  134. self.count += 1
  135. self.cap.release()
  136. if self.count == self.nf: # last video
  137. raise StopIteration
  138. else:
  139. path = self.files[self.count]
  140. self.new_video(path)
  141. ret_val, img0 = self.cap.read()
  142. self.frame += 1
  143. print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='')
  144. else:
  145. # Read image
  146. self.count += 1
  147. img0 = cv2.imread(path) # BGR
  148. assert img0 is not None, 'Image Not Found ' + path
  149. print(f'image {self.count}/{self.nf} {path}: ', end='')
  150. # Padded resize
  151. img = letterbox(img0, self.img_size, stride=self.stride)[0]
  152. # Convert
  153. img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
  154. img = np.ascontiguousarray(img)
  155. return path, img, img0, self.cap
  156. def new_video(self, path):
  157. self.frame = 0
  158. self.cap = cv2.VideoCapture(path)
  159. self.nframes = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
  160. def __len__(self):
  161. return self.nf # number of files
  162. class LoadWebcam: # for inference
  163. def __init__(self, pipe='0', img_size=640, stride=32):
  164. self.img_size = img_size
  165. self.stride = stride
  166. if pipe.isnumeric():
  167. pipe = eval(pipe) # local camera
  168. # pipe = 'rtsp://192.168.1.64/1' # IP camera
  169. # pipe = 'rtsp://username:password@192.168.1.64/1' # IP camera with login
  170. # pipe = 'http://wmccpinetop.axiscam.net/mjpg/video.mjpg' # IP golf camera
  171. self.pipe = pipe
  172. self.cap = cv2.VideoCapture(pipe) # video capture object
  173. self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
  174. def __iter__(self):
  175. self.count = -1
  176. return self
  177. def __next__(self):
  178. self.count += 1
  179. if cv2.waitKey(1) == ord('q'): # q to quit
  180. self.cap.release()
  181. cv2.destroyAllWindows()
  182. raise StopIteration
  183. # Read frame
  184. if self.pipe == 0: # local camera
  185. ret_val, img0 = self.cap.read()
  186. img0 = cv2.flip(img0, 1) # flip left-right
  187. else: # IP camera
  188. n = 0
  189. while True:
  190. n += 1
  191. self.cap.grab()
  192. if n % 30 == 0: # skip frames
  193. ret_val, img0 = self.cap.retrieve()
  194. if ret_val:
  195. break
  196. # Print
  197. assert ret_val, f'Camera Error {self.pipe}'
  198. img_path = 'webcam.jpg'
  199. print(f'webcam {self.count}: ', end='')
  200. # Padded resize
  201. img = letterbox(img0, self.img_size, stride=self.stride)[0]
  202. # Convert
  203. img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
  204. img = np.ascontiguousarray(img)
  205. return img_path, img, img0, None
  206. def __len__(self):
  207. return 0
  208. class LoadStreams: # multiple IP or RTSP cameras
  209. def __init__(self, sources='streams.txt', img_size=640, stride=32):
  210. self.mode = 'stream'
  211. self.img_size = img_size
  212. self.stride = stride
  213. if os.path.isfile(sources):
  214. with open(sources, 'r') as f:
  215. sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]
  216. else:
  217. sources = [sources]
  218. n = len(sources)
  219. self.imgs = [None] * n
  220. self.sources = [clean_str(x) for x in sources] # clean source names for later
  221. for i, s in enumerate(sources):
  222. # Start the thread to read frames from the video stream
  223. print(f'{i + 1}/{n}: {s}... ', end='')
  224. cap = cv2.VideoCapture(eval(s) if s.isnumeric() else s)
  225. assert cap.isOpened(), f'Failed to open {s}'
  226. w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
  227. h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
  228. fps = cap.get(cv2.CAP_PROP_FPS) % 100
  229. _, self.imgs[i] = cap.read() # guarantee first frame
  230. thread = Thread(target=self.update, args=([i, cap]), daemon=True)
  231. print(f' success ({w}x{h} at {fps:.2f} FPS).')
  232. thread.start()
  233. print('') # newline
  234. # check for common shapes
  235. s = np.stack([letterbox(x, self.img_size, stride=self.stride)[0].shape for x in self.imgs], 0) # shapes
  236. self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
  237. if not self.rect:
  238. print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.')
  239. def update(self, index, cap):
  240. # Read next stream frame in a daemon thread
  241. n = 0
  242. while cap.isOpened():
  243. n += 1
  244. # _, self.imgs[index] = cap.read()
  245. cap.grab()
  246. if n == 4: # read every 4th frame
  247. success, im = cap.retrieve()
  248. self.imgs[index] = im if success else self.imgs[index] * 0
  249. n = 0
  250. time.sleep(0.01) # wait time
  251. def __iter__(self):
  252. self.count = -1
  253. return self
  254. def __next__(self):
  255. self.count += 1
  256. img0 = self.imgs.copy()
  257. if cv2.waitKey(1) == ord('q'): # q to quit
  258. cv2.destroyAllWindows()
  259. raise StopIteration
  260. # Letterbox
  261. img = [letterbox(x, self.img_size, auto=self.rect, stride=self.stride)[0] for x in img0]
  262. # Stack
  263. img = np.stack(img, 0)
  264. # Convert
  265. img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB, to bsx3x416x416
  266. img = np.ascontiguousarray(img)
  267. return self.sources, img, img0, None
  268. def __len__(self):
  269. return 0 # 1E12 frames = 32 streams at 30 FPS for 30 years
  270. def img2label_paths(img_paths):
  271. # Define label paths as a function of image paths
  272. sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings
  273. return ['txt'.join(x.replace(sa, sb, 1).rsplit(x.split('.')[-1], 1)) for x in img_paths]
  274. class LoadImagesAndLabels(Dataset): # for training/testing
  275. def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
  276. cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
  277. self.img_size = img_size
  278. self.augment = augment
  279. self.hyp = hyp
  280. self.image_weights = image_weights
  281. self.rect = False if image_weights else rect
  282. self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
  283. self.mosaic_border = [-img_size // 2, -img_size // 2]
  284. self.stride = stride
  285. self.path = path
  286. try:
  287. f = [] # image files
  288. for p in path if isinstance(path, list) else [path]:
  289. p = Path(p) # os-agnostic
  290. if p.is_dir(): # dir
  291. f += glob.glob(str(p / '**' / '*.*'), recursive=True)
  292. # f = list(p.rglob('**/*.*')) # pathlib
  293. elif p.is_file(): # file
  294. with open(p, 'r') as t:
  295. t = t.read().strip().splitlines()
  296. parent = str(p.parent) + os.sep
  297. f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
  298. # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
  299. else:
  300. raise Exception(f'{prefix}{p} does not exist')
  301. self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in img_formats])
  302. # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
  303. assert self.img_files, f'{prefix}No images found'
  304. except Exception as e:
  305. raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}')
  306. # Check cache
  307. self.label_files = img2label_paths(self.img_files) # labels
  308. cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') # cached labels
  309. if cache_path.is_file():
  310. cache, exists = torch.load(cache_path), True # load
  311. if cache['hash'] != get_hash(self.label_files + self.img_files) or 'version' not in cache: # changed
  312. cache, exists = self.cache_labels(cache_path, prefix), False # re-cache
  313. else:
  314. cache, exists = self.cache_labels(cache_path, prefix), False # cache
  315. # Display cache
  316. nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
  317. if exists:
  318. d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  319. tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results
  320. assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {help_url}'
  321. # Read cache
  322. cache.pop('hash') # remove hash
  323. cache.pop('version') # remove version
  324. labels, shapes, self.segments = zip(*cache.values())
  325. self.labels = list(labels)
  326. self.shapes = np.array(shapes, dtype=np.float64)
  327. self.img_files = list(cache.keys()) # update
  328. self.label_files = img2label_paths(cache.keys()) # update
  329. if single_cls:
  330. for x in self.labels:
  331. x[:, 0] = 0
  332. n = len(shapes) # number of images
  333. bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
  334. nb = bi[-1] + 1 # number of batches
  335. self.batch = bi # batch index of image
  336. self.n = n
  337. self.indices = range(n)
  338. # Rectangular Training
  339. if self.rect:
  340. # Sort by aspect ratio
  341. s = self.shapes # wh
  342. ar = s[:, 1] / s[:, 0] # aspect ratio
  343. irect = ar.argsort()
  344. self.img_files = [self.img_files[i] for i in irect]
  345. self.label_files = [self.label_files[i] for i in irect]
  346. self.labels = [self.labels[i] for i in irect]
  347. self.shapes = s[irect] # wh
  348. ar = ar[irect]
  349. # Set training image shapes
  350. shapes = [[1, 1]] * nb
  351. for i in range(nb):
  352. ari = ar[bi == i]
  353. mini, maxi = ari.min(), ari.max()
  354. if maxi < 1:
  355. shapes[i] = [maxi, 1]
  356. elif mini > 1:
  357. shapes[i] = [1, 1 / mini]
  358. self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
  359. # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
  360. self.imgs = [None] * n
  361. if cache_images:
  362. gb = 0 # Gigabytes of cached images
  363. self.img_hw0, self.img_hw = [None] * n, [None] * n
  364. results = ThreadPool(8).imap(lambda x: load_image(*x), zip(repeat(self), range(n))) # 8 threads
  365. pbar = tqdm(enumerate(results), total=n)
  366. for i, x in pbar:
  367. self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # img, hw_original, hw_resized = load_image(self, i)
  368. gb += self.imgs[i].nbytes
  369. pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)'
  370. def cache_labels(self, path=Path('./labels.cache'), prefix=''):
  371. # Cache dataset labels, check images and read shapes
  372. x = {} # dict
  373. nm, nf, ne, nc = 0, 0, 0, 0 # number missing, found, empty, duplicate
  374. pbar = tqdm(zip(self.img_files, self.label_files), desc='Scanning images', total=len(self.img_files))
  375. for i, (im_file, lb_file) in enumerate(pbar):
  376. try:
  377. # verify images
  378. im = Image.open(im_file)
  379. im.verify() # PIL verify
  380. shape = exif_size(im) # image size
  381. segments = [] # instance segments
  382. assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
  383. assert im.format.lower() in img_formats, f'invalid image format {im.format}'
  384. # verify labels
  385. if os.path.isfile(lb_file):
  386. nf += 1 # label found
  387. with open(lb_file, 'r') as f:
  388. l = [x.split() for x in f.read().strip().splitlines()]
  389. if any([len(x) > 8 for x in l]): # is segment
  390. classes = np.array([x[0] for x in l], dtype=np.float32)
  391. segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
  392. l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
  393. l = np.array(l, dtype=np.float32)
  394. if len(l):
  395. assert l.shape[1] == 5, 'labels require 5 columns each'
  396. assert (l >= 0).all(), 'negative labels'
  397. assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels'
  398. assert np.unique(l, axis=0).shape[0] == l.shape[0], 'duplicate labels'
  399. else:
  400. ne += 1 # label empty
  401. l = np.zeros((0, 5), dtype=np.float32)
  402. else:
  403. nm += 1 # label missing
  404. l = np.zeros((0, 5), dtype=np.float32)
  405. x[im_file] = [l, shape, segments]
  406. except Exception as e:
  407. nc += 1
  408. print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}')
  409. pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels... " \
  410. f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
  411. if nf == 0:
  412. print(f'{prefix}WARNING: No labels found in {path}. See {help_url}')
  413. x['hash'] = get_hash(self.label_files + self.img_files)
  414. x['results'] = nf, nm, ne, nc, i + 1
  415. x['version'] = 0.1 # cache version
  416. torch.save(x, path) # save for next time
  417. logging.info(f'{prefix}New cache created: {path}')
  418. return x
  419. def __len__(self):
  420. return len(self.img_files)
  421. # def __iter__(self):
  422. # self.count = -1
  423. # print('ran dataset iter')
  424. # #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
  425. # return self
  426. def __getitem__(self, index):
  427. index = self.indices[index] # linear, shuffled, or image_weights
  428. hyp = self.hyp
  429. mosaic = self.mosaic and random.random() < hyp['mosaic']
  430. if mosaic:
  431. # Load mosaic
  432. img, labels = load_mosaic(self, index)
  433. shapes = None
  434. # MixUp https://arxiv.org/pdf/1710.09412.pdf
  435. if random.random() < hyp['mixup']:
  436. img2, labels2 = load_mosaic(self, random.randint(0, self.n - 1))
  437. r = np.random.beta(8.0, 8.0) # mixup ratio, alpha=beta=8.0
  438. img = (img * r + img2 * (1 - r)).astype(np.uint8)
  439. labels = np.concatenate((labels, labels2), 0)
  440. else:
  441. # Load image
  442. img, (h0, w0), (h, w) = load_image(self, index)
  443. # Letterbox
  444. shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
  445. img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
  446. shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
  447. labels = self.labels[index].copy()
  448. if labels.size: # normalized xywh to pixel xyxy format
  449. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
  450. if self.augment:
  451. # Augment imagespace
  452. if not mosaic:
  453. img, labels = random_perspective(img, labels,
  454. degrees=hyp['degrees'],
  455. translate=hyp['translate'],
  456. scale=hyp['scale'],
  457. shear=hyp['shear'],
  458. perspective=hyp['perspective'])
  459. # Augment colorspace
  460. augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
  461. # Apply cutouts
  462. # if random.random() < 0.9:
  463. # labels = cutout(img, labels)
  464. nL = len(labels) # number of labels
  465. if nL:
  466. labels[:, 1:5] = xyxy2xywh(labels[:, 1:5]) # convert xyxy to xywh
  467. labels[:, [2, 4]] /= img.shape[0] # normalized height 0-1
  468. labels[:, [1, 3]] /= img.shape[1] # normalized width 0-1
  469. if self.augment:
  470. # flip up-down
  471. if random.random() < hyp['flipud']:
  472. img = np.flipud(img)
  473. if nL:
  474. labels[:, 2] = 1 - labels[:, 2]
  475. # flip left-right
  476. if random.random() < hyp['fliplr']:
  477. img = np.fliplr(img)
  478. if nL:
  479. labels[:, 1] = 1 - labels[:, 1]
  480. labels_out = torch.zeros((nL, 6))
  481. if nL:
  482. labels_out[:, 1:] = torch.from_numpy(labels)
  483. # Convert
  484. img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
  485. img = np.ascontiguousarray(img)
  486. return torch.from_numpy(img), labels_out, self.img_files[index], shapes
  487. @staticmethod
  488. def collate_fn(batch):
  489. img, label, path, shapes = zip(*batch) # transposed
  490. for i, l in enumerate(label):
  491. l[:, 0] = i # add target image index for build_targets()
  492. return torch.stack(img, 0), torch.cat(label, 0), path, shapes
  493. @staticmethod
  494. def collate_fn4(batch):
  495. img, label, path, shapes = zip(*batch) # transposed
  496. n = len(shapes) // 4
  497. img4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]
  498. ho = torch.tensor([[0., 0, 0, 1, 0, 0]])
  499. wo = torch.tensor([[0., 0, 1, 0, 0, 0]])
  500. s = torch.tensor([[1, 1, .5, .5, .5, .5]]) # scale
  501. for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW
  502. i *= 4
  503. if random.random() < 0.5:
  504. im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2., mode='bilinear', align_corners=False)[
  505. 0].type(img[i].type())
  506. l = label[i]
  507. else:
  508. im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
  509. l = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
  510. img4.append(im)
  511. label4.append(l)
  512. for i, l in enumerate(label4):
  513. l[:, 0] = i # add target image index for build_targets()
  514. return torch.stack(img4, 0), torch.cat(label4, 0), path4, shapes4
  515. # Ancillary functions --------------------------------------------------------------------------------------------------
  516. def load_image(self, index):
  517. # loads 1 image from dataset, returns img, original hw, resized hw
  518. img = self.imgs[index]
  519. if img is None: # not cached
  520. path = self.img_files[index]
  521. img = cv2.imread(path) # BGR
  522. assert img is not None, 'Image Not Found ' + path
  523. h0, w0 = img.shape[:2] # orig hw
  524. r = self.img_size / max(h0, w0) # resize image to img_size
  525. if r != 1: # always resize down, only resize up if training with augmentation
  526. interp = cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR
  527. img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
  528. return img, (h0, w0), img.shape[:2] # img, hw_original, hw_resized
  529. else:
  530. return self.imgs[index], self.img_hw0[index], self.img_hw[index] # img, hw_original, hw_resized
  531. def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
  532. r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains
  533. hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
  534. dtype = img.dtype # uint8
  535. x = np.arange(0, 256, dtype=np.int16)
  536. lut_hue = ((x * r[0]) % 180).astype(dtype)
  537. lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
  538. lut_val = np.clip(x * r[2], 0, 255).astype(dtype)
  539. img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
  540. cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed
  541. def hist_equalize(img, clahe=True, bgr=False):
  542. # Equalize histogram on BGR image 'img' with img.shape(n,m,3) and range 0-255
  543. yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
  544. if clahe:
  545. c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
  546. yuv[:, :, 0] = c.apply(yuv[:, :, 0])
  547. else:
  548. yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0]) # equalize Y channel histogram
  549. return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR if bgr else cv2.COLOR_YUV2RGB) # convert YUV image to RGB
  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 = random_perspective(img4, labels4, segments4,
  590. degrees=self.hyp['degrees'],
  591. translate=self.hyp['translate'],
  592. scale=self.hyp['scale'],
  593. shear=self.hyp['shear'],
  594. perspective=self.hyp['perspective'],
  595. border=self.mosaic_border) # border to remove
  596. return img4, labels4
  597. def load_mosaic9(self, index):
  598. # loads images in a 9-mosaic
  599. labels9, segments9 = [], []
  600. s = self.img_size
  601. indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices
  602. for i, index in enumerate(indices):
  603. # Load image
  604. img, _, (h, w) = load_image(self, index)
  605. # place img in img9
  606. if i == 0: # center
  607. img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
  608. h0, w0 = h, w
  609. c = s, s, s + w, s + h # xmin, ymin, xmax, ymax (base) coordinates
  610. elif i == 1: # top
  611. c = s, s - h, s + w, s
  612. elif i == 2: # top right
  613. c = s + wp, s - h, s + wp + w, s
  614. elif i == 3: # right
  615. c = s + w0, s, s + w0 + w, s + h
  616. elif i == 4: # bottom right
  617. c = s + w0, s + hp, s + w0 + w, s + hp + h
  618. elif i == 5: # bottom
  619. c = s + w0 - w, s + h0, s + w0, s + h0 + h
  620. elif i == 6: # bottom left
  621. c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
  622. elif i == 7: # left
  623. c = s - w, s + h0 - h, s, s + h0
  624. elif i == 8: # top left
  625. c = s - w, s + h0 - hp - h, s, s + h0 - hp
  626. padx, pady = c[:2]
  627. x1, y1, x2, y2 = [max(x, 0) for x in c] # allocate coords
  628. # Labels
  629. labels, segments = self.labels[index].copy(), self.segments[index].copy()
  630. if labels.size:
  631. labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format
  632. segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
  633. labels9.append(labels)
  634. segments9.extend(segments)
  635. # Image
  636. img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]
  637. hp, wp = h, w # height, width previous
  638. # Offset
  639. yc, xc = [int(random.uniform(0, s)) for _ in self.mosaic_border] # mosaic center x, y
  640. img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]
  641. # Concat/clip labels
  642. labels9 = np.concatenate(labels9, 0)
  643. labels9[:, [1, 3]] -= xc
  644. labels9[:, [2, 4]] -= yc
  645. c = np.array([xc, yc]) # centers
  646. segments9 = [x - c for x in segments9]
  647. for x in (labels9[:, 1:], *segments9):
  648. np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
  649. # img9, labels9 = replicate(img9, labels9) # replicate
  650. # Augment
  651. img9, labels9 = random_perspective(img9, labels9, segments9,
  652. degrees=self.hyp['degrees'],
  653. translate=self.hyp['translate'],
  654. scale=self.hyp['scale'],
  655. shear=self.hyp['shear'],
  656. perspective=self.hyp['perspective'],
  657. border=self.mosaic_border) # border to remove
  658. return img9, labels9
  659. def replicate(img, labels):
  660. # Replicate labels
  661. h, w = img.shape[:2]
  662. boxes = labels[:, 1:].astype(int)
  663. x1, y1, x2, y2 = boxes.T
  664. s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels)
  665. for i in s.argsort()[:round(s.size * 0.5)]: # smallest indices
  666. x1b, y1b, x2b, y2b = boxes[i]
  667. bh, bw = y2b - y1b, x2b - x1b
  668. yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y
  669. x1a, y1a, x2a, y2a = [xc, yc, xc + bw, yc + bh]
  670. img[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
  671. labels = np.append(labels, [[labels[i, 0], x1a, y1a, x2a, y2a]], axis=0)
  672. return img, labels
  673. def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
  674. # Resize and pad image while meeting stride-multiple constraints
  675. shape = img.shape[:2] # current shape [height, width]
  676. if isinstance(new_shape, int):
  677. new_shape = (new_shape, new_shape)
  678. # Scale ratio (new / old)
  679. r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
  680. if not scaleup: # only scale down, do not scale up (for better test mAP)
  681. r = min(r, 1.0)
  682. # Compute padding
  683. ratio = r, r # width, height ratios
  684. new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
  685. dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
  686. if auto: # minimum rectangle
  687. dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding
  688. elif scaleFill: # stretch
  689. dw, dh = 0.0, 0.0
  690. new_unpad = (new_shape[1], new_shape[0])
  691. ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
  692. dw /= 2 # divide padding into 2 sides
  693. dh /= 2
  694. if shape[::-1] != new_unpad: # resize
  695. img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
  696. top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
  697. left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
  698. img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
  699. return img, ratio, (dw, dh)
  700. def random_perspective(img, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0,
  701. border=(0, 0)):
  702. # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
  703. # targets = [cls, xyxy]
  704. height = img.shape[0] + border[0] * 2 # shape(h,w,c)
  705. width = img.shape[1] + border[1] * 2
  706. # Center
  707. C = np.eye(3)
  708. C[0, 2] = -img.shape[1] / 2 # x translation (pixels)
  709. C[1, 2] = -img.shape[0] / 2 # y translation (pixels)
  710. # Perspective
  711. P = np.eye(3)
  712. P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y)
  713. P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x)
  714. # Rotation and Scale
  715. R = np.eye(3)
  716. a = random.uniform(-degrees, degrees)
  717. # a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations
  718. s = random.uniform(1 - scale, 1 + scale)
  719. # s = 2 ** random.uniform(-scale, scale)
  720. R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
  721. # Shear
  722. S = np.eye(3)
  723. S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)
  724. S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)
  725. # Translation
  726. T = np.eye(3)
  727. T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width # x translation (pixels)
  728. T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height # y translation (pixels)
  729. # Combined rotation matrix
  730. M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT
  731. if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
  732. if perspective:
  733. img = cv2.warpPerspective(img, M, dsize=(width, height), borderValue=(114, 114, 114))
  734. else: # affine
  735. img = cv2.warpAffine(img, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
  736. # Visualize
  737. # import matplotlib.pyplot as plt
  738. # ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
  739. # ax[0].imshow(img[:, :, ::-1]) # base
  740. # ax[1].imshow(img2[:, :, ::-1]) # warped
  741. # Transform label coordinates
  742. n = len(targets)
  743. if n:
  744. use_segments = any(x.any() for x in segments)
  745. new = np.zeros((n, 4))
  746. if use_segments: # warp segments
  747. segments = resample_segments(segments) # upsample
  748. for i, segment in enumerate(segments):
  749. xy = np.ones((len(segment), 3))
  750. xy[:, :2] = segment
  751. xy = xy @ M.T # transform
  752. xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2] # perspective rescale or affine
  753. # clip
  754. new[i] = segment2box(xy, width, height)
  755. else: # warp boxes
  756. xy = np.ones((n * 4, 3))
  757. xy[:, :2] = targets[:, [1, 2, 3, 4, 1, 4, 3, 2]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
  758. xy = xy @ M.T # transform
  759. xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine
  760. # create new boxes
  761. x = xy[:, [0, 2, 4, 6]]
  762. y = xy[:, [1, 3, 5, 7]]
  763. new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
  764. # clip
  765. new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)
  766. new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)
  767. # filter candidates
  768. i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01 if use_segments else 0.10)
  769. targets = targets[i]
  770. targets[:, 1:5] = new[i]
  771. return img, targets
  772. def box_candidates(box1, box2, wh_thr=2, ar_thr=20, area_thr=0.1, eps=1e-16): # box1(4,n), box2(4,n)
  773. # Compute candidate boxes: box1 before augment, box2 after augment, wh_thr (pixels), aspect_ratio_thr, area_ratio
  774. w1, h1 = box1[2] - box1[0], box1[3] - box1[1]
  775. w2, h2 = box2[2] - box2[0], box2[3] - box2[1]
  776. ar = np.maximum(w2 / (h2 + eps), h2 / (w2 + eps)) # aspect ratio
  777. return (w2 > wh_thr) & (h2 > wh_thr) & (w2 * h2 / (w1 * h1 + eps) > area_thr) & (ar < ar_thr) # candidates
  778. def cutout(image, labels):
  779. # Applies image cutout augmentation https://arxiv.org/abs/1708.04552
  780. h, w = image.shape[:2]
  781. def bbox_ioa(box1, box2):
  782. # Returns the intersection over box2 area given box1, box2. box1 is 4, box2 is nx4. boxes are x1y1x2y2
  783. box2 = box2.transpose()
  784. # Get the coordinates of bounding boxes
  785. b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
  786. b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
  787. # Intersection area
  788. inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \
  789. (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0)
  790. # box2 area
  791. box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + 1e-16
  792. # Intersection over box2 area
  793. return inter_area / box2_area
  794. # create random masks
  795. scales = [0.5] * 1 + [0.25] * 2 + [0.125] * 4 + [0.0625] * 8 + [0.03125] * 16 # image size fraction
  796. for s in scales:
  797. mask_h = random.randint(1, int(h * s))
  798. mask_w = random.randint(1, int(w * s))
  799. # box
  800. xmin = max(0, random.randint(0, w) - mask_w // 2)
  801. ymin = max(0, random.randint(0, h) - mask_h // 2)
  802. xmax = min(w, xmin + mask_w)
  803. ymax = min(h, ymin + mask_h)
  804. # apply random color mask
  805. image[ymin:ymax, xmin:xmax] = [random.randint(64, 191) for _ in range(3)]
  806. # return unobscured labels
  807. if len(labels) and s > 0.03:
  808. box = np.array([xmin, ymin, xmax, ymax], dtype=np.float32)
  809. ioa = bbox_ioa(box, labels[:, 1:5]) # intersection over area
  810. labels = labels[ioa < 0.60] # remove >60% obscured labels
  811. return labels
  812. def create_folder(path='./new'):
  813. # Create folder
  814. if os.path.exists(path):
  815. shutil.rmtree(path) # delete output folder
  816. os.makedirs(path) # make new output folder
  817. def flatten_recursive(path='../coco128'):
  818. # Flatten a recursive directory by bringing all files to top level
  819. new_path = Path(path + '_flat')
  820. create_folder(new_path)
  821. for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)):
  822. shutil.copyfile(file, new_path / Path(file).name)
  823. def extract_boxes(path='../coco128/'): # from utils.datasets import *; extract_boxes('../coco128')
  824. # Convert detection dataset into classification dataset, with one directory per class
  825. path = Path(path) # images dir
  826. shutil.rmtree(path / 'classifier') if (path / 'classifier').is_dir() else None # remove existing
  827. files = list(path.rglob('*.*'))
  828. n = len(files) # number of files
  829. for im_file in tqdm(files, total=n):
  830. if im_file.suffix[1:] in img_formats:
  831. # image
  832. im = cv2.imread(str(im_file))[..., ::-1] # BGR to RGB
  833. h, w = im.shape[:2]
  834. # labels
  835. lb_file = Path(img2label_paths([str(im_file)])[0])
  836. if Path(lb_file).exists():
  837. with open(lb_file, 'r') as f:
  838. lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels
  839. for j, x in enumerate(lb):
  840. c = int(x[0]) # class
  841. f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename
  842. if not f.parent.is_dir():
  843. f.parent.mkdir(parents=True)
  844. b = x[1:] * [w, h, w, h] # box
  845. # b[2:] = b[2:].max() # rectangle to square
  846. b[2:] = b[2:] * 1.2 + 3 # pad
  847. b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int)
  848. b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image
  849. b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
  850. assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'
  851. def autosplit(path='../coco128', weights=(0.9, 0.1, 0.0), annotated_only=False):
  852. """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
  853. Usage: from utils.datasets import *; autosplit('../coco128')
  854. Arguments
  855. path: Path to images directory
  856. weights: Train, val, test weights (list)
  857. annotated_only: Only use images with an annotated txt file
  858. """
  859. path = Path(path) # images dir
  860. files = sum([list(path.rglob(f"*.{img_ext}")) for img_ext in img_formats], []) # image files only
  861. n = len(files) # number of files
  862. indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
  863. txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files
  864. [(path / x).unlink() for x in txt if (path / x).exists()] # remove existing
  865. print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)
  866. for i, img in tqdm(zip(indices, files), total=n):
  867. if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
  868. with open(path / txt[i], 'a') as f:
  869. f.write(str(img) + '\n') # add image to txt file