Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

934 rindas
37KB

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