|
- import glob
- import logging
- import os
- import random
- import hashlib
- import math
- import multiprocessing
- import numpy as np
- from itertools import repeat
- from PIL import Image, ExifTags
- from tqdm import tqdm
- from multiprocessing.pool import ThreadPool
- from pathlib import Path
-
- import mindspore.dataset as de
- from mindspore.dataset import vision
- from mindspore.dataset.vision import Inter
-
- from .general import segments2boxes, xywhn2xyxy, xyxy2xywh
- from .augumentations import load_image, load_mosaic, letterbox, random_perspective, \
- augment_hsv
-
- # Parameters
- help_url = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
- img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes
- vid_formats = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
-
- # Get orientation exif tag
- for orientation in ExifTags.TAGS.keys():
- if ExifTags.TAGS[orientation] == 'Orientation':
- break
-
-
- def get_hash(paths):
- # Returns a single hash value of a list of paths (files or dirs)
- size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes
- h = hashlib.md5(str(size).encode()) # hash sizes
- h.update(''.join(paths).encode()) # hash paths
- return h.hexdigest() # return hash
-
-
- def exif_size(img):
- # Returns exif-corrected PIL size
- s = img.size # (width, height)
- try:
- rotation = dict(img._getexif().items())[orientation]
- if rotation == 6: # rotation 270
- s = (s[1], s[0])
- elif rotation == 8: # rotation 90
- s = (s[1], s[0])
- except:
- pass
-
- return s
-
-
- def img2label_paths(img_paths):
- # Define label paths as a function of image paths
- sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings
- return ['txt'.join(x.replace(sa, sb, 1).rsplit(x.split('.')[-1], 1)) for x in img_paths]
-
-
- class LoadImagesAndLabels: # for training/testing
- def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
- cache_images=False, single_cls=False, stride=32, pad=0.0, max_box_per_img=160, prefix=''):
- self.img_size = img_size
- self.augment = augment
- self.hyp = hyp
- self.image_weights = image_weights
- self.rect = False if image_weights else rect
- self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
- self.mosaic_border = [-img_size // 2, -img_size // 2]
- self.stride = stride
- self.path = path
- self.max_box_per_img = max_box_per_img
- # self.albumentations = Albumentations() if augment else None
-
- try:
- f = [] # image files
- for p in path if isinstance(path, list) else [path]:
- p = Path(p) # os-agnostic
- if p.is_dir(): # dir
- f += glob.glob(str(p / '**' / '*.*'), recursive=True)
- # f = list(p.rglob('**/*.*')) # pathlib
- elif p.is_file(): # file
- with open(p, 'r') as t:
- t = t.read().strip().splitlines()
- parent = str(p.parent) + os.sep
- f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
- # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
- else:
- raise Exception(f'{prefix}{p} does not exist')
- self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in img_formats])
- # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in img_formats]) # pathlib
- assert self.img_files, f'{prefix}No images found'
- except Exception as e:
- raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}')
-
- # Check cache
- self.label_files = img2label_paths(self.img_files) # labels
- cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') # cached labels
- if cache_path.is_file():
- cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict
- assert cache['version'] == self.cache_version # matches current version
- assert cache['hash'] == get_hash(self.label_files + self.img_files) # identical hash
- else:
- cache, exists = self.cache_labels(cache_path, prefix), False # cache
-
- # Display cache
- nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total
- if exists:
- d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
- tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results
- assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {help_url}'
-
- # Read cache
- cache.pop('hash') # remove hash
- cache.pop('version') # remove version
- labels, shapes, self.segments = zip(*cache.values())
- self.labels = list(labels)
- self.shapes = np.array(shapes, dtype=np.float64)
- self.img_files = list(cache.keys()) # update
- self.label_files = img2label_paths(cache.keys()) # update
- if single_cls:
- for x in self.labels:
- x[:, 0] = 0
-
- n = len(shapes) # number of images
- bi = np.floor(np.arange(n) / batch_size).astype(np.int32) # batch index
- nb = bi[-1] + 1 # number of batches
- self.batch = bi # batch index of image
- self.n = n
- self.indices = range(n)
-
- # Rectangular Training
- if self.rect:
- # Sort by aspect ratio
- s = self.shapes # wh
- ar = s[:, 1] / s[:, 0] # aspect ratio
- irect = ar.argsort()
- self.img_files = [self.img_files[i] for i in irect]
- self.label_files = [self.label_files[i] for i in irect]
- self.labels = [self.labels[i] for i in irect]
- self.shapes = s[irect] # wh
- ar = ar[irect]
-
- # Set training image shapes
- shapes = [[1, 1]] * nb
- for i in range(nb):
- ari = ar[bi == i]
- mini, maxi = ari.min(), ari.max()
- if maxi < 1:
- shapes[i] = [maxi, 1]
- elif mini > 1:
- shapes[i] = [1, 1 / mini]
-
- self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
-
- # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM)
- self.imgs = [None] * n
- if cache_images:
- if cache_images == 'disk':
- self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + '_npy')
- self.img_npy = [self.im_cache_dir / Path(f).with_suffix('.npy').name for f in self.img_files]
- self.im_cache_dir.mkdir(parents=True, exist_ok=True)
- gb = 0 # Gigabytes of cached images
- self.img_hw0, self.img_hw = [None] * n, [None] * n
- results = ThreadPool(8).imap(lambda x: load_image(*x), zip(repeat(self), range(n)))
- pbar = tqdm(enumerate(results), total=n)
- for i, x in pbar:
- if cache_images == 'disk':
- if not self.img_npy[i].exists():
- np.save(self.img_npy[i].as_posix(), x[0])
- gb += self.img_npy[i].stat().st_size
- else:
- self.imgs[i], self.img_hw0[i], self.img_hw[i] = x
- gb += self.imgs[i].nbytes
- pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)'
- pbar.close()
-
- def cache_labels(self, path=Path('./labels.cache'), prefix=''):
- # Cache dataset labels, check images and read shapes
- x = {} # dict
- nm, nf, ne, nc = 0, 0, 0, 0 # number missing, found, empty, duplicate
- pbar = tqdm(zip(self.img_files, self.label_files), desc='Scanning images', total=len(self.img_files))
- for i, (im_file, lb_file) in enumerate(pbar):
- try:
- # verify images
- im = Image.open(im_file)
- im.verify() # PIL verify
- shape = exif_size(im) # image size
- segments = [] # instance segments
- assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
- assert im.format.lower() in img_formats, f'invalid image format {im.format}'
-
- # verify labels
- if os.path.isfile(lb_file):
- nf += 1 # label found
- with open(lb_file, 'r') as f:
- l = [x.split() for x in f.read().strip().splitlines()]
- if any([len(x) > 8 for x in l]): # is segment
- classes = np.array([x[0] for x in l], dtype=np.float32)
- segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in l] # (cls, xy1...)
- l = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
- l = np.array(l, dtype=np.float32)
- if len(l):
- assert l.shape[1] == 5, 'labels require 5 columns each'
- assert (l >= 0).all(), 'negative labels'
- assert (l[:, 1:] <= 1).all(), 'non-normalized or out of bounds coordinate labels'
- assert np.unique(l, axis=0).shape[0] == l.shape[0], 'duplicate labels'
- else:
- ne += 1 # label empty
- l = np.zeros((0, 5), dtype=np.float32)
- else:
- nm += 1 # label missing
- l = np.zeros((0, 5), dtype=np.float32)
- x[im_file] = [l, shape, segments]
- except Exception as e:
- nc += 1
- print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}')
-
- pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels... " \
- f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
- pbar.close()
-
- if nf == 0:
- print(f'{prefix}WARNING: No labels found in {path}. See {help_url}')
-
- x['hash'] = get_hash(self.label_files + self.img_files)
- x['results'] = nf, nm, ne, nc, i + 1
- x['version'] = 0.1 # cache version
- np.save(path, x) # save for next time
- logging.info(f'{prefix}New cache created: {path}')
- return x
-
- def __len__(self):
- return len(self.img_files)
-
- def __getitem__(self, index):
- index = self.indices[index] # linear, shuffled, or image_weights
-
- hyp = self.hyp
- mosaic = self.mosaic and random.random() < hyp['mosaic']
- if mosaic:
- # Load mosaic
- img, labels = load_mosaic(self, index)
- # shapes = None
- shapes = np.zeros((3, 2))
-
- # MixUp https://arxiv.org/pdf/1710.09412.pdf
- if random.random() < hyp['mixup']:
- img2, labels2 = load_mosaic(self, random.randint(0, len(self.labels) - 1))
- r = np.random.beta(8.0, 8.0) # mixup ratio, alpha=beta=8.0
- img = (img * r + img2 * (1 - r)).astype(np.uint8)
- labels = np.concatenate((labels, labels2), 0)
-
- else:
- # Load image
- img, (h0, w0), (h, w) = load_image(self, index)
-
- # Letterbox
- shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
- img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
- # shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
- shapes = np.array([[h0, w0],
- [h / h0, w / w0],
- [pad[0], pad[1]]]) # (3, 2)
-
- labels = self.labels[index].copy()
- if labels.size: # normalized xywh to pixel xyxy format
- labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
-
- if self.augment:
- # Augment imagespace
- if not mosaic:
- img, labels = random_perspective(img, labels,
- degrees=hyp['degrees'],
- translate=hyp['translate'],
- scale=hyp['scale'],
- shear=hyp['shear'],
- perspective=hyp['perspective'])
-
- # img, labels = self.albumentations(img, labels)
-
- # Augment colorspace
- augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
-
- # Apply cutouts
- # if random.random() < 0.9:
- # labels = cutout(img, labels)
-
- nL = len(labels) # number of labels
- if nL:
- labels[:, [1, 3]] = labels[:, [1, 3]].clip(0, img.shape[1] - 1e-3) # x1, x2
- labels[:, [2, 4]] = labels[:, [2, 4]].clip(0, img.shape[0] - 1e-3) # y1, y2
- labels[:, 1:5] = xyxy2xywh(labels[:, 1:5]) # convert xyxy to xywh
- labels[:, [2, 4]] /= img.shape[0] # normalized height 0-1
- labels[:, [1, 3]] /= img.shape[1] # normalized width 0-1
-
- if self.augment:
- # flip up-down
- if random.random() < hyp['flipud']:
- img = np.flipud(img)
- if nL:
- labels[:, 2] = 1 - labels[:, 2]
-
- # flip left-right
- if random.random() < hyp['fliplr']:
- img = np.fliplr(img)
- if nL:
- labels[:, 1] = 1 - labels[:, 1]
-
- _labels_out = np.zeros((nL, 6))
- if nL:
- _labels_out[:, 1:] = labels
-
- # create fixed label, avoid dynamic shape problem.
- labels_out = np.full((self.max_box_per_img, 6), -1, dtype=np.float32)
- if nL:
- labels_out[:min(nL, self.max_box_per_img), :] = _labels_out[:min(nL, self.max_box_per_img), :]
-
- # Convert
- img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
- img = np.ascontiguousarray(img)
-
- return img, labels_out, self.img_files[index], shapes
-
- @staticmethod
- def collate_fn(img, label, path, shapes, batch_info):
- for i, l in enumerate(label):
- l[:, 0] = i # add target image index for build_targets()
- return np.stack(img, 0), np.stack(label, 0), path, np.stack(shapes, 0)
-
- @staticmethod
- def collate_fn4(img, label, path, shapes, batch_info):
- n = len(img) // 4
- img4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]
-
- ho = np.array([[0., 0, 0, 1, 0, 0]])
- wo = np.array([[0., 0, 1, 0, 0, 0]])
- s = np.array([[1, 1, .5, .5, .5, .5]]) # scale
- for i in range(n): # BCHW
- i *= 4
- if random.random() < 0.5:
-
- _resize_shape = (img[i].shape[1] * 2, img[i].shape[2] * 2)
- # (c,h,w) -> (h,w,c) -> (c,h,w)
- im = vision.Resize(_resize_shape, Inter.BILINEAR)(img[i].transpose(1, 2, 0)).transpose(2, 0, 1)
- # im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2., mode='bilinear', align_corners=False)[
- # 0].type(img[i].type())
- l = label[i]
- else:
- im = np.concatenate((np.concatenate((img[i], img[i + 1]), 1),
- np.concatenate((img[i + 2], img[i + 3]), 1)), 2)
- l = np.concatenate((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
- img4.append(im)
- label4.append(l)
-
- for i, l in enumerate(label4):
- l[:, 0] = i # add target image index for build_targets()
-
- return np.stack(img4, 0), np.stack(label4, 0), path4, shapes4
-
-
- def create_dataloader(path, imgsz, batch_size, stride, opt, epoch_size=300, hyp=None, augment=False,
- cache=False, pad=0.0, rect=False, rank=0, rank_size=1, num_parallel_workers=8,
- shuffle=True, drop_remainder=True, image_weights=False, quad=False, max_box_per_img=160, prefix=''):
- # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
- dataset = LoadImagesAndLabels(path, imgsz, batch_size,
- augment=augment, # augment images
- hyp=hyp, # augmentation hyperparameters
- rect=rect, # rectangular training
- cache_images=cache,
- single_cls=opt.single_cls,
- stride=int(stride),
- pad=pad,
- image_weights=image_weights,
- max_box_per_img=max_box_per_img,
- prefix=prefix)
-
- cores = multiprocessing.cpu_count()
- num_parallel_workers = min(int(cores / rank_size), num_parallel_workers)
- dataset_column_names = ["img", "label_out", "img_files", "shapes"]
- if rank_size > 1:
- ds = de.GeneratorDataset(dataset, column_names=dataset_column_names,
- num_parallel_workers=min(8, num_parallel_workers), shuffle=shuffle,
- num_shards=rank_size, shard_id=rank)
- else:
- ds = de.GeneratorDataset(dataset, column_names=dataset_column_names,
- num_parallel_workers=min(32, num_parallel_workers), shuffle=shuffle)
- ds = ds.batch(batch_size,
- per_batch_map=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn,
- input_columns=dataset_column_names,
- drop_remainder=drop_remainder)
- ds = ds.repeat(epoch_size)
-
- per_epoch_size = int(len(dataset) / batch_size / rank_size) if drop_remainder else \
- math.ceil(len(dataset) / batch_size / rank_size)
-
- return ds, dataset, per_epoch_size
-
-
|