ultralytics 8.0.80 single-line docstring fixes (#2060)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
GitHub
@ -66,7 +66,7 @@ class Compose:
|
||||
|
||||
|
||||
class BaseMixTransform:
|
||||
"""This implementation is from mmyolo"""
|
||||
"""This implementation is from mmyolo."""
|
||||
|
||||
def __init__(self, dataset, pre_transform=None, p=0.0) -> None:
|
||||
self.dataset = dataset
|
||||
@ -77,12 +77,12 @@ class BaseMixTransform:
|
||||
if random.uniform(0, 1) > self.p:
|
||||
return labels
|
||||
|
||||
# get index of one or three other images
|
||||
# Get index of one or three other images
|
||||
indexes = self.get_indexes()
|
||||
if isinstance(indexes, int):
|
||||
indexes = [indexes]
|
||||
|
||||
# get images information will be used for Mosaic or MixUp
|
||||
# Get images information will be used for Mosaic or MixUp
|
||||
mix_labels = [self.dataset.get_label_info(i) for i in indexes]
|
||||
|
||||
if self.pre_transform is not None:
|
||||
@ -132,7 +132,7 @@ class Mosaic(BaseMixTransform):
|
||||
img = labels_patch['img']
|
||||
h, w = labels_patch.pop('resized_shape')
|
||||
|
||||
# place img in img4
|
||||
# Place img in img4
|
||||
if i == 0: # top left
|
||||
img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
|
||||
x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
|
||||
@ -158,7 +158,7 @@ class Mosaic(BaseMixTransform):
|
||||
return final_labels
|
||||
|
||||
def _update_labels(self, labels, padw, padh):
|
||||
"""Update labels"""
|
||||
"""Update labels."""
|
||||
nh, nw = labels['img'].shape[:2]
|
||||
labels['instances'].convert_bbox(format='xyxy')
|
||||
labels['instances'].denormalize(nw, nh)
|
||||
@ -193,7 +193,7 @@ class MixUp(BaseMixTransform):
|
||||
return random.randint(0, len(self.dataset) - 1)
|
||||
|
||||
def _mix_transform(self, labels):
|
||||
# Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
|
||||
"""Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf."""
|
||||
r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0
|
||||
labels2 = labels['mix_labels'][0]
|
||||
labels['img'] = (labels['img'] * r + labels2['img'] * (1 - r)).astype(np.uint8)
|
||||
@ -217,12 +217,12 @@ class RandomPerspective:
|
||||
self.scale = scale
|
||||
self.shear = shear
|
||||
self.perspective = perspective
|
||||
# mosaic border
|
||||
# Mosaic border
|
||||
self.border = border
|
||||
self.pre_transform = pre_transform
|
||||
|
||||
def affine_transform(self, img, border):
|
||||
# Center
|
||||
"""Center."""
|
||||
C = np.eye(3, dtype=np.float32)
|
||||
|
||||
C[0, 2] = -img.shape[1] / 2 # x translation (pixels)
|
||||
@ -253,7 +253,7 @@ class RandomPerspective:
|
||||
|
||||
# Combined rotation matrix
|
||||
M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT
|
||||
# affine image
|
||||
# Affine image
|
||||
if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
|
||||
if self.perspective:
|
||||
img = cv2.warpPerspective(img, M, dsize=self.size, borderValue=(114, 114, 114))
|
||||
@ -281,7 +281,7 @@ class RandomPerspective:
|
||||
xy = xy @ M.T # transform
|
||||
xy = (xy[:, :2] / xy[:, 2:3] if self.perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine
|
||||
|
||||
# create new boxes
|
||||
# Create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
y = xy[:, [1, 3, 5, 7]]
|
||||
return np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1)), dtype=bboxes.dtype).reshape(4, n).T
|
||||
@ -348,7 +348,7 @@ class RandomPerspective:
|
||||
img = labels['img']
|
||||
cls = labels['cls']
|
||||
instances = labels.pop('instances')
|
||||
# make sure the coord formats are right
|
||||
# Make sure the coord formats are right
|
||||
instances.convert_bbox(format='xyxy')
|
||||
instances.denormalize(*img.shape[:2][::-1])
|
||||
|
||||
@ -362,19 +362,19 @@ class RandomPerspective:
|
||||
|
||||
segments = instances.segments
|
||||
keypoints = instances.keypoints
|
||||
# update bboxes if there are segments.
|
||||
# Update bboxes if there are segments.
|
||||
if len(segments):
|
||||
bboxes, segments = self.apply_segments(segments, M)
|
||||
|
||||
if keypoints is not None:
|
||||
keypoints = self.apply_keypoints(keypoints, M)
|
||||
new_instances = Instances(bboxes, segments, keypoints, bbox_format='xyxy', normalized=False)
|
||||
# clip
|
||||
# Clip
|
||||
new_instances.clip(*self.size)
|
||||
|
||||
# filter instances
|
||||
# Filter instances
|
||||
instances.scale(scale_w=scale, scale_h=scale, bbox_only=True)
|
||||
# make the bboxes have the same scale with new_bboxes
|
||||
# Make the bboxes have the same scale with new_bboxes
|
||||
i = self.box_candidates(box1=instances.bboxes.T,
|
||||
box2=new_instances.bboxes.T,
|
||||
area_thr=0.01 if len(segments) else 0.10)
|
||||
@ -441,7 +441,7 @@ class RandomFlip:
|
||||
if self.direction == 'horizontal' and random.random() < self.p:
|
||||
img = np.fliplr(img)
|
||||
instances.fliplr(w)
|
||||
# for keypoints
|
||||
# For keypoints
|
||||
if self.flip_idx is not None and instances.keypoints is not None:
|
||||
instances.keypoints = np.ascontiguousarray(instances.keypoints[:, self.flip_idx, :])
|
||||
labels['img'] = np.ascontiguousarray(img)
|
||||
@ -450,7 +450,7 @@ class RandomFlip:
|
||||
|
||||
|
||||
class LetterBox:
|
||||
"""Resize image and padding for detection, instance segmentation, pose"""
|
||||
"""Resize image and padding for detection, instance segmentation, pose."""
|
||||
|
||||
def __init__(self, new_shape=(640, 640), auto=False, scaleFill=False, scaleup=True, stride=32):
|
||||
self.new_shape = new_shape
|
||||
@ -505,7 +505,7 @@ class LetterBox:
|
||||
return img
|
||||
|
||||
def _update_labels(self, labels, ratio, padw, padh):
|
||||
"""Update labels"""
|
||||
"""Update labels."""
|
||||
labels['instances'].convert_bbox(format='xyxy')
|
||||
labels['instances'].denormalize(*labels['img'].shape[:2][::-1])
|
||||
labels['instances'].scale(*ratio)
|
||||
@ -519,7 +519,7 @@ class CopyPaste:
|
||||
self.p = p
|
||||
|
||||
def __call__(self, labels):
|
||||
# Implement Copy-Paste augmentation https://arxiv.org/abs/2012.07177, labels as nx5 np.array(cls, xyxy)
|
||||
"""Implement Copy-Paste augmentation https://arxiv.org/abs/2012.07177, labels as nx5 np.array(cls, xyxy)."""
|
||||
im = labels['img']
|
||||
cls = labels['cls']
|
||||
h, w = im.shape[:2]
|
||||
@ -531,7 +531,7 @@ class CopyPaste:
|
||||
_, w, _ = im.shape # height, width, channels
|
||||
im_new = np.zeros(im.shape, np.uint8)
|
||||
|
||||
# calculate ioa first then select indexes randomly
|
||||
# Calculate ioa first then select indexes randomly
|
||||
ins_flip = deepcopy(instances)
|
||||
ins_flip.fliplr(w)
|
||||
|
||||
@ -641,7 +641,7 @@ class Format:
|
||||
labels['bboxes'] = torch.from_numpy(instances.bboxes) if nl else torch.zeros((nl, 4))
|
||||
if self.return_keypoint:
|
||||
labels['keypoints'] = torch.from_numpy(instances.keypoints)
|
||||
# then we can use collate_fn
|
||||
# Then we can use collate_fn
|
||||
if self.batch_idx:
|
||||
labels['batch_idx'] = torch.zeros(nl)
|
||||
return labels
|
||||
@ -654,7 +654,7 @@ class Format:
|
||||
return img
|
||||
|
||||
def _format_segments(self, instances, cls, w, h):
|
||||
"""convert polygon points to bitmap"""
|
||||
"""convert polygon points to bitmap."""
|
||||
segments = instances.segments
|
||||
if self.mask_overlap:
|
||||
masks, sorted_idx = polygons2masks_overlap((h, w), segments, downsample_ratio=self.mask_ratio)
|
||||
|
||||
@ -70,7 +70,7 @@ class BaseDataset(Dataset):
|
||||
|
||||
self.ni = len(self.labels)
|
||||
|
||||
# rect stuff
|
||||
# Rect stuff
|
||||
self.rect = rect
|
||||
self.batch_size = batch_size
|
||||
self.stride = stride
|
||||
@ -79,13 +79,13 @@ class BaseDataset(Dataset):
|
||||
assert self.batch_size is not None
|
||||
self.set_rectangle()
|
||||
|
||||
# cache stuff
|
||||
# Cache stuff
|
||||
self.ims = [None] * self.ni
|
||||
self.npy_files = [Path(f).with_suffix('.npy') for f in self.im_files]
|
||||
if cache:
|
||||
self.cache_images(cache)
|
||||
|
||||
# transforms
|
||||
# Transforms
|
||||
self.transforms = self.build_transforms(hyp=hyp)
|
||||
|
||||
def get_img_files(self, img_path):
|
||||
@ -96,13 +96,13 @@ class BaseDataset(Dataset):
|
||||
p = Path(p) # os-agnostic
|
||||
if p.is_dir(): # dir
|
||||
f += glob.glob(str(p / '**' / '*.*'), recursive=True)
|
||||
# f = list(p.rglob('*.*')) # pathlib
|
||||
# F = list(p.rglob('*.*')) # pathlib
|
||||
elif p.is_file(): # file
|
||||
with open(p) 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)
|
||||
# F += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
|
||||
else:
|
||||
raise FileNotFoundError(f'{self.prefix}{p} does not exist')
|
||||
im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
|
||||
@ -113,7 +113,7 @@ class BaseDataset(Dataset):
|
||||
return im_files
|
||||
|
||||
def update_labels(self, include_class: Optional[list]):
|
||||
"""include_class, filter labels to include only these classes (optional)"""
|
||||
"""include_class, filter labels to include only these classes (optional)."""
|
||||
include_class_array = np.array(include_class).reshape(1, -1)
|
||||
for i in range(len(self.labels)):
|
||||
if include_class is not None:
|
||||
@ -129,7 +129,7 @@ class BaseDataset(Dataset):
|
||||
self.labels[i]['cls'][:, 0] = 0
|
||||
|
||||
def load_image(self, i):
|
||||
# Loads 1 image from dataset index 'i', returns (im, resized hw)
|
||||
"""Loads 1 image from dataset index 'i', returns (im, resized hw)."""
|
||||
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
|
||||
if im is None: # not cached in RAM
|
||||
if fn.exists(): # load npy
|
||||
@ -147,7 +147,7 @@ class BaseDataset(Dataset):
|
||||
return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized
|
||||
|
||||
def cache_images(self, cache):
|
||||
# cache images to memory or disk
|
||||
"""Cache images to memory or disk."""
|
||||
gb = 0 # Gigabytes of cached images
|
||||
self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni
|
||||
fcn = self.cache_images_to_disk if cache == 'disk' else self.load_image
|
||||
@ -164,7 +164,7 @@ class BaseDataset(Dataset):
|
||||
pbar.close()
|
||||
|
||||
def cache_images_to_disk(self, i):
|
||||
# Saves an image as an *.npy file for faster loading
|
||||
"""Saves an image as an *.npy file for faster loading."""
|
||||
f = self.npy_files[i]
|
||||
if not f.exists():
|
||||
np.save(f.as_posix(), cv2.imread(self.im_files[i]))
|
||||
@ -211,17 +211,17 @@ class BaseDataset(Dataset):
|
||||
return len(self.labels)
|
||||
|
||||
def update_labels_info(self, label):
|
||||
"""custom your label format here"""
|
||||
"""custom your label format here."""
|
||||
return label
|
||||
|
||||
def build_transforms(self, hyp=None):
|
||||
"""Users can custom augmentations here
|
||||
like:
|
||||
if self.augment:
|
||||
# training transforms
|
||||
# Training transforms
|
||||
return Compose([])
|
||||
else:
|
||||
# val transforms
|
||||
# Val transforms
|
||||
return Compose([])
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@ -104,7 +104,7 @@ def build_dataloader(cfg, batch, img_path, data_info, stride=32, rect=False, ran
|
||||
generator=generator), dataset
|
||||
|
||||
|
||||
# build classification
|
||||
# Build classification
|
||||
# TODO: using cfg like `build_dataloader`
|
||||
def build_classification_dataloader(path,
|
||||
imgsz=224,
|
||||
@ -114,7 +114,7 @@ def build_classification_dataloader(path,
|
||||
rank=-1,
|
||||
workers=8,
|
||||
shuffle=True):
|
||||
# Returns Dataloader object to be used with YOLOv5 Classifier
|
||||
"""Returns Dataloader object to be used with YOLOv5 Classifier."""
|
||||
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
|
||||
dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache)
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
|
||||
@ -70,7 +70,7 @@ class LoadStreams:
|
||||
self.threads[i].start()
|
||||
LOGGER.info('') # newline
|
||||
|
||||
# check for common shapes
|
||||
# Check for common shapes
|
||||
s = np.stack([LetterBox(imgsz, auto, stride=stride)(image=x).shape for x in self.imgs])
|
||||
self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
|
||||
self.auto = auto and self.rect
|
||||
@ -81,7 +81,7 @@ class LoadStreams:
|
||||
LOGGER.warning('WARNING ⚠️ Stream shapes differ. For optimal performance supply similarly-shaped streams.')
|
||||
|
||||
def update(self, i, cap, stream):
|
||||
# Read stream `i` frames in daemon thread
|
||||
"""Read stream `i` frames in daemon thread."""
|
||||
n, f = 0, self.frames[i] # frame number, frame array
|
||||
while cap.isOpened() and n < f:
|
||||
n += 1
|
||||
@ -123,7 +123,7 @@ class LoadStreams:
|
||||
class LoadScreenshots:
|
||||
# YOLOv8 screenshot dataloader, i.e. `yolo predict source=screen`
|
||||
def __init__(self, source, imgsz=640, stride=32, auto=True, transforms=None):
|
||||
# source = [screen_number left top width height] (pixels)
|
||||
"""source = [screen_number left top width height] (pixels)."""
|
||||
check_requirements('mss')
|
||||
import mss # noqa
|
||||
|
||||
@ -156,7 +156,7 @@ class LoadScreenshots:
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
# mss screen capture: get raw pixels from the screen as np array
|
||||
"""mss screen capture: get raw pixels from the screen as np array."""
|
||||
im0 = np.array(self.sct.grab(self.monitor))[:, :, :3] # [:, :, :3] BGRA to BGR
|
||||
s = f'screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: '
|
||||
|
||||
@ -256,7 +256,7 @@ class LoadImages:
|
||||
return path, im, im0, self.cap, s
|
||||
|
||||
def _new_video(self, path):
|
||||
# Create a new video capture object
|
||||
"""Create a new video capture object."""
|
||||
self.frame = 0
|
||||
self.cap = cv2.VideoCapture(path)
|
||||
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
|
||||
@ -266,7 +266,7 @@ class LoadImages:
|
||||
# self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0)
|
||||
|
||||
def _cv2_rotate(self, im):
|
||||
# Rotate a cv2 video manually
|
||||
"""Rotate a cv2 video manually."""
|
||||
if self.orientation == 0:
|
||||
return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE)
|
||||
elif self.orientation == 180:
|
||||
@ -291,7 +291,7 @@ class LoadPilAndNumpy:
|
||||
self.auto = auto
|
||||
self.transforms = transforms
|
||||
self.mode = 'image'
|
||||
# generate fake paths
|
||||
# Generate fake paths
|
||||
self.bs = len(self.im0)
|
||||
|
||||
@staticmethod
|
||||
|
||||
@ -55,19 +55,19 @@ class Albumentations:
|
||||
|
||||
|
||||
def normalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD, inplace=False):
|
||||
# Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = (x - mean) / std
|
||||
"""Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = (x - mean) / std."""
|
||||
return TF.normalize(x, mean, std, inplace=inplace)
|
||||
|
||||
|
||||
def denormalize(x, mean=IMAGENET_MEAN, std=IMAGENET_STD):
|
||||
# Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = x * std + mean
|
||||
"""Denormalize RGB images x per ImageNet stats in BCHW format, i.e. = x * std + mean."""
|
||||
for i in range(3):
|
||||
x[:, i] = x[:, i] * std[i] + mean[i]
|
||||
return x
|
||||
|
||||
|
||||
def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
|
||||
# HSV color-space augmentation
|
||||
"""HSV color-space augmentation."""
|
||||
if hgain or sgain or vgain:
|
||||
r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains
|
||||
hue, sat, val = cv2.split(cv2.cvtColor(im, cv2.COLOR_BGR2HSV))
|
||||
@ -83,7 +83,7 @@ def augment_hsv(im, hgain=0.5, sgain=0.5, vgain=0.5):
|
||||
|
||||
|
||||
def hist_equalize(im, clahe=True, bgr=False):
|
||||
# Equalize histogram on BGR image 'im' with im.shape(n,m,3) and range 0-255
|
||||
"""Equalize histogram on BGR image 'im' with im.shape(n,m,3) and range 0-255."""
|
||||
yuv = cv2.cvtColor(im, cv2.COLOR_BGR2YUV if bgr else cv2.COLOR_RGB2YUV)
|
||||
if clahe:
|
||||
c = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
@ -94,7 +94,7 @@ def hist_equalize(im, clahe=True, bgr=False):
|
||||
|
||||
|
||||
def replicate(im, labels):
|
||||
# Replicate labels
|
||||
"""Replicate labels."""
|
||||
h, w = im.shape[:2]
|
||||
boxes = labels[:, 1:].astype(int)
|
||||
x1, y1, x2, y2 = boxes.T
|
||||
@ -213,7 +213,7 @@ def random_perspective(im,
|
||||
xy = xy @ M.T # transform
|
||||
xy = xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2] # perspective rescale or affine
|
||||
|
||||
# clip
|
||||
# Clip
|
||||
new[i] = segment2box(xy, width, height)
|
||||
|
||||
else: # warp boxes
|
||||
@ -222,16 +222,16 @@ def random_perspective(im,
|
||||
xy = xy @ M.T # transform
|
||||
xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]).reshape(n, 8) # perspective rescale or affine
|
||||
|
||||
# create new boxes
|
||||
# Create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
y = xy[:, [1, 3, 5, 7]]
|
||||
new = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
|
||||
|
||||
# clip
|
||||
# Clip
|
||||
new[:, [0, 2]] = new[:, [0, 2]].clip(0, width)
|
||||
new[:, [1, 3]] = new[:, [1, 3]].clip(0, height)
|
||||
|
||||
# filter candidates
|
||||
# Filter candidates
|
||||
i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01 if use_segments else 0.10)
|
||||
targets = targets[i]
|
||||
targets[:, 1:5] = new[i]
|
||||
@ -240,13 +240,13 @@ def random_perspective(im,
|
||||
|
||||
|
||||
def copy_paste(im, labels, segments, p=0.5):
|
||||
# Implement Copy-Paste augmentation https://arxiv.org/abs/2012.07177, labels as nx5 np.array(cls, xyxy)
|
||||
"""Implement Copy-Paste augmentation https://arxiv.org/abs/2012.07177, labels as nx5 np.array(cls, xyxy)."""
|
||||
n = len(segments)
|
||||
if p and n:
|
||||
h, w, c = im.shape # height, width, channels
|
||||
im_new = np.zeros(im.shape, np.uint8)
|
||||
|
||||
# calculate ioa first then select indexes randomly
|
||||
# Calculate ioa first then select indexes randomly
|
||||
boxes = np.stack([w - labels[:, 3], labels[:, 2], w - labels[:, 1], labels[:, 4]], axis=-1) # (n, 4)
|
||||
ioa = bbox_ioa(boxes, labels[:, 1:5]) # intersection over area
|
||||
indexes = np.nonzero((ioa < 0.30).all(1))[0] # (N, )
|
||||
@ -265,7 +265,7 @@ def copy_paste(im, labels, segments, p=0.5):
|
||||
|
||||
|
||||
def cutout(im, labels, p=0.5):
|
||||
# Applies image cutout augmentation https://arxiv.org/abs/1708.04552
|
||||
"""Applies image cutout augmentation https://arxiv.org/abs/1708.04552."""
|
||||
if random.random() < p:
|
||||
h, w = im.shape[:2]
|
||||
scales = [0.5] * 1 + [0.25] * 2 + [0.125] * 4 + [0.0625] * 8 + [0.03125] * 16 # image size fraction
|
||||
@ -273,16 +273,16 @@ def cutout(im, labels, p=0.5):
|
||||
mask_h = random.randint(1, int(h * s)) # create random masks
|
||||
mask_w = random.randint(1, int(w * s))
|
||||
|
||||
# box
|
||||
# Box
|
||||
xmin = max(0, random.randint(0, w) - mask_w // 2)
|
||||
ymin = max(0, random.randint(0, h) - mask_h // 2)
|
||||
xmax = min(w, xmin + mask_w)
|
||||
ymax = min(h, ymin + mask_h)
|
||||
|
||||
# apply random color mask
|
||||
# Apply random color mask
|
||||
im[ymin:ymax, xmin:xmax] = [random.randint(64, 191) for _ in range(3)]
|
||||
|
||||
# return unobscured labels
|
||||
# Return unobscured labels
|
||||
if len(labels) and s > 0.03:
|
||||
box = np.array([[xmin, ymin, xmax, ymax]], dtype=np.float32)
|
||||
ioa = bbox_ioa(box, xywhn2xyxy(labels[:, 1:5], w, h))[0] # intersection over area
|
||||
@ -292,7 +292,7 @@ def cutout(im, labels, p=0.5):
|
||||
|
||||
|
||||
def mixup(im, labels, im2, labels2):
|
||||
# Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
|
||||
"""Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf."""
|
||||
r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0
|
||||
im = (im * r + im2 * (1 - r)).astype(np.uint8)
|
||||
labels = np.concatenate((labels, labels2), 0)
|
||||
@ -350,7 +350,7 @@ def classify_albumentations(
|
||||
|
||||
|
||||
def classify_transforms(size=224):
|
||||
# Transforms to apply if albumentations not installed
|
||||
"""Transforms to apply if albumentations not installed."""
|
||||
assert isinstance(size, int), f'ERROR: classify_transforms size {size} must be integer, not (list, tuple)'
|
||||
# T.Compose([T.ToTensor(), T.Resize(size), T.CenterCrop(size), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
|
||||
return T.Compose([CenterCrop(size), ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])
|
||||
|
||||
@ -50,7 +50,7 @@ for orientation in ExifTags.TAGS.keys():
|
||||
|
||||
|
||||
def get_hash(paths):
|
||||
# Returns a single hash value of a list of paths (files or dirs)
|
||||
"""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.sha256(str(size).encode()) # hash sizes
|
||||
h.update(''.join(paths).encode()) # hash paths
|
||||
@ -58,7 +58,7 @@ def get_hash(paths):
|
||||
|
||||
|
||||
def exif_size(img):
|
||||
# Returns exif-corrected PIL size
|
||||
"""Returns exif-corrected PIL size."""
|
||||
s = img.size # (width, height)
|
||||
with contextlib.suppress(Exception):
|
||||
rotation = dict(img._getexif().items())[orientation]
|
||||
@ -94,7 +94,7 @@ def exif_transpose(image):
|
||||
|
||||
|
||||
def seed_worker(worker_id):
|
||||
# Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader
|
||||
"""Set dataloader worker seed https://pytorch.org/docs/stable/notes/randomness.html#dataloader."""
|
||||
worker_seed = torch.initial_seed() % 2 ** 32
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
@ -156,7 +156,7 @@ def create_dataloader(path,
|
||||
|
||||
|
||||
class InfiniteDataLoader(dataloader.DataLoader):
|
||||
""" Dataloader that reuses workers
|
||||
"""Dataloader that reuses workers
|
||||
|
||||
Uses same syntax as vanilla DataLoader
|
||||
"""
|
||||
@ -175,7 +175,7 @@ class InfiniteDataLoader(dataloader.DataLoader):
|
||||
|
||||
|
||||
class _RepeatSampler:
|
||||
""" Sampler that repeats forever
|
||||
"""Sampler that repeats forever
|
||||
|
||||
Args:
|
||||
sampler (Dataset.sampler): The sampler to repeat.
|
||||
@ -192,7 +192,7 @@ class _RepeatSampler:
|
||||
class LoadScreenshots:
|
||||
# YOLOv5 screenshot dataloader, i.e. `python detect.py --source "screen 0 100 100 512 256"`
|
||||
def __init__(self, source, img_size=640, stride=32, auto=True, transforms=None):
|
||||
# source = [screen_number left top width height] (pixels)
|
||||
"""source = [screen_number left top width height] (pixels)."""
|
||||
check_requirements('mss')
|
||||
import mss
|
||||
|
||||
@ -224,7 +224,7 @@ class LoadScreenshots:
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
# mss screen capture: get raw pixels from the screen as np array
|
||||
"""mss screen capture: get raw pixels from the screen as np array."""
|
||||
im0 = np.array(self.sct.grab(self.monitor))[:, :, :3] # [:, :, :3] BGRA to BGR
|
||||
s = f'screen {self.screen} (LTWH): {self.left},{self.top},{self.width},{self.height}: '
|
||||
|
||||
@ -320,7 +320,7 @@ class LoadImages:
|
||||
return path, im, im0, self.cap, s
|
||||
|
||||
def _new_video(self, path):
|
||||
# Create a new video capture object
|
||||
"""Create a new video capture object."""
|
||||
self.frame = 0
|
||||
self.cap = cv2.VideoCapture(path)
|
||||
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride)
|
||||
@ -328,7 +328,7 @@ class LoadImages:
|
||||
# self.cap.set(cv2.CAP_PROP_ORIENTATION_AUTO, 0) # disable https://github.com/ultralytics/yolov5/issues/8493
|
||||
|
||||
def _cv2_rotate(self, im):
|
||||
# Rotate a cv2 video manually
|
||||
"""Rotate a cv2 video manually."""
|
||||
if self.orientation == 0:
|
||||
return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE)
|
||||
elif self.orientation == 180:
|
||||
@ -379,7 +379,7 @@ class LoadStreams:
|
||||
self.threads[i].start()
|
||||
LOGGER.info('') # newline
|
||||
|
||||
# check for common shapes
|
||||
# Check for common shapes
|
||||
s = np.stack([letterbox(x, img_size, stride=stride, auto=auto)[0].shape for x in self.imgs])
|
||||
self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
|
||||
self.auto = auto and self.rect
|
||||
@ -388,7 +388,7 @@ class LoadStreams:
|
||||
LOGGER.warning('WARNING ⚠️ Stream shapes differ. For optimal performance supply similarly-shaped streams.')
|
||||
|
||||
def update(self, i, cap, stream):
|
||||
# Read stream `i` frames in daemon thread
|
||||
"""Read stream `i` frames in daemon thread."""
|
||||
n, f = 0, self.frames[i] # frame number, frame array
|
||||
while cap.isOpened() and n < f:
|
||||
n += 1
|
||||
@ -428,13 +428,13 @@ class LoadStreams:
|
||||
|
||||
|
||||
def img2label_paths(img_paths):
|
||||
# Define label paths as a function of image paths
|
||||
"""Define label paths as a function of image paths."""
|
||||
sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}' # /images/, /labels/ substrings
|
||||
return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
|
||||
|
||||
|
||||
class LoadImagesAndLabels(Dataset):
|
||||
# YOLOv5 train_loader/val_loader, loads images and labels for training and validation
|
||||
"""YOLOv5 train_loader/val_loader, loads images and labels for training and validation."""
|
||||
cache_version = 0.6 # dataset labels *.cache version
|
||||
rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4]
|
||||
|
||||
@ -590,7 +590,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
pbar.close()
|
||||
|
||||
def check_cache_ram(self, safety_margin=0.1, prefix=''):
|
||||
# Check image caching requirements vs available memory
|
||||
"""Check image caching requirements vs available memory."""
|
||||
b, gb = 0, 1 << 30 # bytes of cached images, bytes per gigabytes
|
||||
n = min(self.n, 30) # extrapolate from 30 random images
|
||||
for _ in range(n):
|
||||
@ -648,12 +648,6 @@ class LoadImagesAndLabels(Dataset):
|
||||
def __len__(self):
|
||||
return len(self.im_files)
|
||||
|
||||
# def __iter__(self):
|
||||
# self.count = -1
|
||||
# print('ran dataset iter')
|
||||
# #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
|
||||
# return self
|
||||
|
||||
def __getitem__(self, index):
|
||||
index = self.indices[index] # linear, shuffled, or image_weights
|
||||
|
||||
@ -729,7 +723,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
return torch.from_numpy(img), labels_out, self.im_files[index], shapes
|
||||
|
||||
def load_image(self, i):
|
||||
# Loads 1 image from dataset index 'i', returns (im, original hw, resized hw)
|
||||
"""Loads 1 image from dataset index 'i', returns (im, original hw, resized hw)."""
|
||||
im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i],
|
||||
if im is None: # not cached in RAM
|
||||
if fn.exists(): # load npy
|
||||
@ -746,13 +740,13 @@ class LoadImagesAndLabels(Dataset):
|
||||
return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized
|
||||
|
||||
def cache_images_to_disk(self, i):
|
||||
# Saves an image as an *.npy file for faster loading
|
||||
"""Saves an image as an *.npy file for faster loading."""
|
||||
f = self.npy_files[i]
|
||||
if not f.exists():
|
||||
np.save(f.as_posix(), cv2.imread(self.im_files[i]))
|
||||
|
||||
def load_mosaic(self, index):
|
||||
# YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
|
||||
"""YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic."""
|
||||
labels4, segments4 = [], []
|
||||
s = self.img_size
|
||||
yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border) # mosaic center x, y
|
||||
@ -762,7 +756,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
# Load image
|
||||
img, _, (h, w) = self.load_image(index)
|
||||
|
||||
# place img in img4
|
||||
# Place img in img4
|
||||
if i == 0: # top left
|
||||
img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
|
||||
x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
|
||||
@ -810,7 +804,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
return img4, labels4
|
||||
|
||||
def load_mosaic9(self, index):
|
||||
# YOLOv5 9-mosaic loader. Loads 1 image + 8 random images into a 9-image mosaic
|
||||
"""YOLOv5 9-mosaic loader. Loads 1 image + 8 random images into a 9-image mosaic."""
|
||||
labels9, segments9 = [], []
|
||||
s = self.img_size
|
||||
indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices
|
||||
@ -820,7 +814,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
# Load image
|
||||
img, _, (h, w) = self.load_image(index)
|
||||
|
||||
# place img in img9
|
||||
# Place img in img9
|
||||
if i == 0: # center
|
||||
img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
|
||||
h0, w0 = h, w
|
||||
@ -888,7 +882,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch):
|
||||
# YOLOv8 collate function, outputs dict
|
||||
"""YOLOv8 collate function, outputs dict."""
|
||||
im, label, path, shapes = zip(*batch) # transposed
|
||||
for i, lb in enumerate(label):
|
||||
lb[:, 0] = i # add target image index for build_targets()
|
||||
@ -904,7 +898,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
|
||||
@staticmethod
|
||||
def collate_fn_old(batch):
|
||||
# YOLOv5 original collate function
|
||||
"""YOLOv5 original collate function."""
|
||||
im, label, path, shapes = zip(*batch) # transposed
|
||||
for i, lb in enumerate(label):
|
||||
lb[:, 0] = i # add target image index for build_targets()
|
||||
@ -913,7 +907,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
|
||||
# Ancillary functions --------------------------------------------------------------------------------------------------
|
||||
def flatten_recursive(path=DATASETS_DIR / 'coco128'):
|
||||
# Flatten a recursive directory by bringing all files to top level
|
||||
"""Flatten a recursive directory by bringing all files to top level."""
|
||||
new_path = Path(f'{str(path)}_flat')
|
||||
if os.path.exists(new_path):
|
||||
shutil.rmtree(new_path) # delete output folder
|
||||
@ -930,11 +924,11 @@ def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders impo
|
||||
n = len(files) # number of files
|
||||
for im_file in tqdm(files, total=n):
|
||||
if im_file.suffix[1:] in IMG_FORMATS:
|
||||
# image
|
||||
# Image
|
||||
im = cv2.imread(str(im_file))[..., ::-1] # BGR to RGB
|
||||
h, w = im.shape[:2]
|
||||
|
||||
# labels
|
||||
# Labels
|
||||
lb_file = Path(img2label_paths([str(im_file)])[0])
|
||||
if Path(lb_file).exists():
|
||||
with open(lb_file) as f:
|
||||
@ -947,7 +941,7 @@ def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders impo
|
||||
f.parent.mkdir(parents=True)
|
||||
|
||||
b = x[1:] * [w, h, w, h] # box
|
||||
# b[2:] = b[2:].max() # rectangle to square
|
||||
# B[2:] = b[2:].max() # rectangle to square
|
||||
b[2:] = b[2:] * 1.2 + 3 # pad
|
||||
b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(int)
|
||||
|
||||
@ -957,7 +951,7 @@ def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders impo
|
||||
|
||||
|
||||
def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False):
|
||||
""" Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
|
||||
"""Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
|
||||
Usage: from utils.dataloaders import *; autosplit()
|
||||
Arguments
|
||||
path: Path to images directory
|
||||
@ -983,11 +977,11 @@ def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), ann
|
||||
|
||||
|
||||
def verify_image_label(args):
|
||||
# Verify one image-label pair
|
||||
"""Verify one image-label pair."""
|
||||
im_file, lb_file, prefix = args
|
||||
nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments
|
||||
try:
|
||||
# verify images
|
||||
# Verify images
|
||||
im = Image.open(im_file)
|
||||
im.verify() # PIL verify
|
||||
shape = exif_size(im) # image size
|
||||
@ -1000,7 +994,7 @@ def verify_image_label(args):
|
||||
ImageOps.exif_transpose(Image.open(im_file)).save(im_file, 'JPEG', subsampling=0, quality=100)
|
||||
msg = f'{prefix}WARNING ⚠️ {im_file}: corrupt JPEG restored and saved'
|
||||
|
||||
# verify labels
|
||||
# Verify labels
|
||||
if os.path.isfile(lb_file):
|
||||
nf = 1 # label found
|
||||
with open(lb_file) as f:
|
||||
@ -1077,7 +1071,7 @@ def create_classification_dataloader(path,
|
||||
rank=-1,
|
||||
workers=8,
|
||||
shuffle=True):
|
||||
# Returns Dataloader object to be used with YOLOv5 Classifier
|
||||
"""Returns Dataloader object to be used with YOLOv5 Classifier."""
|
||||
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
|
||||
dataset = ClassificationDataset(root=path, imgsz=imgsz, augment=augment, cache=cache)
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
|
||||
@ -193,7 +193,7 @@ class YOLODataset(BaseDataset):
|
||||
self.transforms = self.build_transforms(hyp)
|
||||
|
||||
def update_labels_info(self, label):
|
||||
"""custom your label format here"""
|
||||
"""custom your label format here."""
|
||||
# NOTE: cls is not with bboxes now, classification and semantic segmentation need an independent cls label
|
||||
# we can make it also support classification and semantic segmentation by add or remove some dict keys there.
|
||||
bboxes = label.pop('bboxes')
|
||||
|
||||
@ -39,7 +39,7 @@ class MixAndRectDataset:
|
||||
"""
|
||||
labels = deepcopy(self.dataset[index])
|
||||
for transform in self.dataset.transforms.tolist():
|
||||
# mosaic and mixup
|
||||
# Mosaic and mixup
|
||||
if hasattr(transform, 'get_indexes'):
|
||||
indexes = transform.get_indexes(self.dataset)
|
||||
if not isinstance(indexes, collections.abc.Sequence):
|
||||
|
||||
@ -37,13 +37,13 @@ for orientation in ExifTags.TAGS.keys():
|
||||
|
||||
|
||||
def img2label_paths(img_paths):
|
||||
# Define label paths as a function of image paths
|
||||
"""Define label paths as a function of image paths."""
|
||||
sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}' # /images/, /labels/ substrings
|
||||
return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
|
||||
|
||||
|
||||
def get_hash(paths):
|
||||
# Returns a single hash value of a list of paths (files or dirs)
|
||||
"""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.sha256(str(size).encode()) # hash sizes
|
||||
h.update(''.join(paths).encode()) # hash paths
|
||||
@ -51,7 +51,7 @@ def get_hash(paths):
|
||||
|
||||
|
||||
def exif_size(img):
|
||||
# Returns exif-corrected PIL size
|
||||
"""Returns exif-corrected PIL size."""
|
||||
s = img.size # (width, height)
|
||||
with contextlib.suppress(Exception):
|
||||
rotation = dict(img._getexif().items())[orientation]
|
||||
@ -61,12 +61,12 @@ def exif_size(img):
|
||||
|
||||
|
||||
def verify_image_label(args):
|
||||
# Verify one image-label pair
|
||||
"""Verify one image-label pair."""
|
||||
im_file, lb_file, prefix, keypoint, num_cls, nkpt, ndim = args
|
||||
# number (missing, found, empty, corrupt), message, segments, keypoints
|
||||
# Number (missing, found, empty, corrupt), message, segments, keypoints
|
||||
nm, nf, ne, nc, msg, segments, keypoints = 0, 0, 0, 0, '', [], None
|
||||
try:
|
||||
# verify images
|
||||
# Verify images
|
||||
im = Image.open(im_file)
|
||||
im.verify() # PIL verify
|
||||
shape = exif_size(im) # image size
|
||||
@ -80,7 +80,7 @@ def verify_image_label(args):
|
||||
ImageOps.exif_transpose(Image.open(im_file)).save(im_file, 'JPEG', subsampling=0, quality=100)
|
||||
msg = f'{prefix}WARNING ⚠️ {im_file}: corrupt JPEG restored and saved'
|
||||
|
||||
# verify labels
|
||||
# Verify labels
|
||||
if os.path.isfile(lb_file):
|
||||
nf = 1 # label found
|
||||
with open(lb_file) as f:
|
||||
@ -191,7 +191,7 @@ def polygons2masks_overlap(imgsz, segments, downsample_ratio=1):
|
||||
|
||||
|
||||
def check_det_dataset(dataset, autodownload=True):
|
||||
# Download, check and/or unzip dataset if not found locally
|
||||
"""Download, check and/or unzip dataset if not found locally."""
|
||||
data = check_file(dataset)
|
||||
|
||||
# Download (optional)
|
||||
@ -321,7 +321,7 @@ class HUBDatasetStats():
|
||||
"""
|
||||
|
||||
def __init__(self, path='coco128.yaml', autodownload=False):
|
||||
# Initialize class
|
||||
"""Initialize class."""
|
||||
zipped, data_dir, yaml_path = self._unzip(Path(path))
|
||||
try:
|
||||
# data = yaml_load(check_yaml(yaml_path)) # data dict
|
||||
@ -339,7 +339,7 @@ class HUBDatasetStats():
|
||||
|
||||
@staticmethod
|
||||
def _find_yaml(dir):
|
||||
# Return data.yaml file
|
||||
"""Return data.yaml file."""
|
||||
files = list(dir.glob('*.yaml')) or list(dir.rglob('*.yaml')) # try root level first and then recursive
|
||||
assert files, f'No *.yaml file found in {dir}'
|
||||
if len(files) > 1:
|
||||
@ -349,7 +349,7 @@ class HUBDatasetStats():
|
||||
return files[0]
|
||||
|
||||
def _unzip(self, path):
|
||||
# Unzip data.zip
|
||||
"""Unzip data.zip."""
|
||||
if not str(path).endswith('.zip'): # path is data.yaml
|
||||
return False, None, path
|
||||
assert Path(path).is_file(), f'Error unzipping {path}, file not found'
|
||||
@ -362,12 +362,12 @@ class HUBDatasetStats():
|
||||
compress_one_image(f, self.im_dir / Path(f).name) # save to dataset-hub
|
||||
|
||||
def get_json(self, save=False, verbose=False):
|
||||
# Return dataset JSON for Ultralytics HUB
|
||||
"""Return dataset JSON for Ultralytics HUB."""
|
||||
# from ultralytics.yolo.data import YOLODataset
|
||||
from ultralytics.yolo.data.dataloaders.v5loader import LoadImagesAndLabels
|
||||
|
||||
def _round(labels):
|
||||
# Update labels to integer class and 6 decimal place floats
|
||||
"""Update labels to integer class and 6 decimal place floats."""
|
||||
return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]
|
||||
|
||||
for split in 'train', 'val', 'test':
|
||||
@ -400,7 +400,7 @@ class HUBDatasetStats():
|
||||
return self.stats
|
||||
|
||||
def process_images(self):
|
||||
# Compress images for Ultralytics HUB
|
||||
"""Compress images for Ultralytics HUB."""
|
||||
# from ultralytics.yolo.data import YOLODataset
|
||||
from ultralytics.yolo.data.dataloaders.v5loader import LoadImagesAndLabels
|
||||
|
||||
|
||||
Reference in New Issue
Block a user