|
|
|
""" |
|
Dataloaders and dataset utils |
|
""" |
|
|
|
import contextlib |
|
import glob |
|
import hashlib |
|
import json |
|
import math |
|
import os |
|
import random |
|
import shutil |
|
import time |
|
from itertools import repeat |
|
from multiprocessing.pool import Pool, ThreadPool |
|
from pathlib import Path |
|
from threading import Thread |
|
from urllib.parse import urlparse |
|
|
|
import numpy as np |
|
import psutil |
|
import torch |
|
import torch.nn.functional as F |
|
import torchvision |
|
import yaml |
|
from PIL import ExifTags, Image, ImageOps |
|
from torch.utils.data import DataLoader, Dataset, dataloader, distributed |
|
from tqdm import tqdm |
|
|
|
from utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste, |
|
letterbox, mixup, random_perspective) |
|
from utils.general import (DATASETS_DIR, LOGGER, NUM_THREADS, TQDM_BAR_FORMAT, check_dataset, check_requirements, |
|
check_yaml, clean_str, cv2, is_colab, is_kaggle, segments2boxes, unzip_file, xyn2xy, |
|
xywh2xyxy, xywhn2xyxy, xyxy2xywhn) |
|
from utils.torch_utils import torch_distributed_zero_first |
|
|
|
|
|
HELP_URL = 'See https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data' |
|
IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm' |
|
VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' |
|
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) |
|
RANK = int(os.getenv('RANK', -1)) |
|
PIN_MEMORY = str(os.getenv('PIN_MEMORY', True)).lower() == 'true' |
|
|
|
|
|
for orientation in ExifTags.TAGS.keys(): |
|
if ExifTags.TAGS[orientation] == 'Orientation': |
|
break |
|
|
|
|
|
def get_hash(paths): |
|
|
|
size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) |
|
h = hashlib.md5(str(size).encode()) |
|
h.update(''.join(paths).encode()) |
|
return h.hexdigest() |
|
|
|
|
|
def exif_size(img): |
|
|
|
s = img.size |
|
with contextlib.suppress(Exception): |
|
rotation = dict(img._getexif().items())[orientation] |
|
if rotation in [6, 8]: |
|
s = (s[1], s[0]) |
|
return s |
|
|
|
|
|
def exif_transpose(image): |
|
""" |
|
Transpose a PIL image accordingly if it has an EXIF Orientation tag. |
|
Inplace version of https://github.com/python-pillow/Pillow/blob/master/src/PIL/ImageOps.py exif_transpose() |
|
|
|
:param image: The image to transpose. |
|
:return: An image. |
|
""" |
|
exif = image.getexif() |
|
orientation = exif.get(0x0112, 1) |
|
if orientation > 1: |
|
method = { |
|
2: Image.FLIP_LEFT_RIGHT, |
|
3: Image.ROTATE_180, |
|
4: Image.FLIP_TOP_BOTTOM, |
|
5: Image.TRANSPOSE, |
|
6: Image.ROTATE_270, |
|
7: Image.TRANSVERSE, |
|
8: Image.ROTATE_90}.get(orientation) |
|
if method is not None: |
|
image = image.transpose(method) |
|
del exif[0x0112] |
|
image.info["exif"] = exif.tobytes() |
|
return image |
|
|
|
|
|
def seed_worker(worker_id): |
|
|
|
worker_seed = torch.initial_seed() % 2 ** 32 |
|
np.random.seed(worker_seed) |
|
random.seed(worker_seed) |
|
|
|
|
|
|
|
class LoadImages: |
|
|
|
def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vid_stride=1): |
|
if isinstance(path, str) and Path(path).suffix == ".txt": |
|
path = Path(path).read_text().rsplit() |
|
files = [] |
|
for p in sorted(path) if isinstance(path, (list, tuple)) else [path]: |
|
p = str(Path(p).resolve()) |
|
if '*' in p: |
|
files.extend(sorted(glob.glob(p, recursive=True))) |
|
elif os.path.isdir(p): |
|
files.extend(sorted(glob.glob(os.path.join(p, '*.*')))) |
|
elif os.path.isfile(p): |
|
files.append(p) |
|
else: |
|
raise FileNotFoundError(f'{p} does not exist') |
|
|
|
images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS] |
|
videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS] |
|
ni, nv = len(images), len(videos) |
|
|
|
self.img_size = img_size |
|
self.stride = stride |
|
self.files = images + videos |
|
self.nf = ni + nv |
|
self.video_flag = [False] * ni + [True] * nv |
|
self.mode = 'image' |
|
self.auto = auto |
|
self.transforms = transforms |
|
self.vid_stride = vid_stride |
|
if any(videos): |
|
self._new_video(videos[0]) |
|
else: |
|
self.cap = None |
|
assert self.nf > 0, f'No images or videos found in {p}. ' \ |
|
f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}' |
|
|
|
def __iter__(self): |
|
self.count = 0 |
|
return self |
|
|
|
def __next__(self): |
|
if self.count == self.nf: |
|
raise StopIteration |
|
path = self.files[self.count] |
|
|
|
if self.video_flag[self.count]: |
|
|
|
self.mode = 'video' |
|
for _ in range(self.vid_stride): |
|
self.cap.grab() |
|
ret_val, im0 = self.cap.retrieve() |
|
while not ret_val: |
|
self.count += 1 |
|
self.cap.release() |
|
if self.count == self.nf: |
|
raise StopIteration |
|
path = self.files[self.count] |
|
self._new_video(path) |
|
ret_val, im0 = self.cap.read() |
|
|
|
self.frame += 1 |
|
|
|
s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ' |
|
|
|
else: |
|
|
|
self.count += 1 |
|
im0 = cv2.imread(path) |
|
assert im0 is not None, f'Image Not Found {path}' |
|
s = f'image {self.count}/{self.nf} {path}: ' |
|
|
|
if self.transforms: |
|
im = self.transforms(im0) |
|
else: |
|
im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0] |
|
im = im.transpose((2, 0, 1))[::-1] |
|
im = np.ascontiguousarray(im) |
|
|
|
return path, im, im0, self.cap, s |
|
|
|
def _new_video(self, path): |
|
|
|
self.frame = 0 |
|
self.cap = cv2.VideoCapture(path) |
|
self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT) / self.vid_stride) |
|
self.orientation = int(self.cap.get(cv2.CAP_PROP_ORIENTATION_META)) |
|
|
|
|
|
def _cv2_rotate(self, im): |
|
|
|
if self.orientation == 0: |
|
return cv2.rotate(im, cv2.ROTATE_90_CLOCKWISE) |
|
elif self.orientation == 180: |
|
return cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE) |
|
elif self.orientation == 90: |
|
return cv2.rotate(im, cv2.ROTATE_180) |
|
return im |
|
|
|
def __len__(self): |
|
return self.nf |
|
|
|
def img2label_paths(img_paths): |
|
|
|
sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}' |
|
return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths] |
|
|
|
|
|
def flatten_recursive(path=DATASETS_DIR / 'coco128'): |
|
|
|
new_path = Path(f'{str(path)}_flat') |
|
if os.path.exists(new_path): |
|
shutil.rmtree(new_path) |
|
os.makedirs(new_path) |
|
for file in tqdm(glob.glob(f'{str(Path(path))}/**/*.*', recursive=True)): |
|
shutil.copyfile(file, new_path / Path(file).name) |
|
|
|
|
|
def extract_boxes(path=DATASETS_DIR / 'coco128'): |
|
|
|
path = Path(path) |
|
shutil.rmtree(path / 'classification') if (path / 'classification').is_dir() else None |
|
files = list(path.rglob('*.*')) |
|
n = len(files) |
|
for im_file in tqdm(files, total=n): |
|
if im_file.suffix[1:] in IMG_FORMATS: |
|
|
|
im = cv2.imread(str(im_file))[..., ::-1] |
|
h, w = im.shape[:2] |
|
|
|
|
|
lb_file = Path(img2label_paths([str(im_file)])[0]) |
|
if Path(lb_file).exists(): |
|
with open(lb_file) as f: |
|
lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) |
|
|
|
for j, x in enumerate(lb): |
|
c = int(x[0]) |
|
f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' |
|
if not f.parent.is_dir(): |
|
f.parent.mkdir(parents=True) |
|
|
|
b = x[1:] * [w, h, w, h] |
|
|
|
b[2:] = b[2:] * 1.2 + 3 |
|
b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(int) |
|
|
|
b[[0, 2]] = np.clip(b[[0, 2]], 0, w) |
|
b[[1, 3]] = np.clip(b[[1, 3]], 0, h) |
|
assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}' |
|
|
|
|
|
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 |
|
Usage: from utils.dataloaders import *; autosplit() |
|
Arguments |
|
path: Path to images directory |
|
weights: Train, val, test weights (list, tuple) |
|
annotated_only: Only use images with an annotated txt file |
|
""" |
|
path = Path(path) |
|
files = sorted(x for x in path.rglob('*.*') if x.suffix[1:].lower() in IMG_FORMATS) |
|
n = len(files) |
|
random.seed(0) |
|
indices = random.choices([0, 1, 2], weights=weights, k=n) |
|
|
|
txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] |
|
for x in txt: |
|
if (path.parent / x).exists(): |
|
(path.parent / x).unlink() |
|
|
|
print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only) |
|
for i, img in tqdm(zip(indices, files), total=n): |
|
if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): |
|
with open(path.parent / txt[i], 'a') as f: |
|
f.write(f'./{img.relative_to(path.parent).as_posix()}' + '\n') |
|
|
|
|
|
def verify_image_label(args): |
|
|
|
im_file, lb_file, prefix = args |
|
nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] |
|
try: |
|
|
|
im = Image.open(im_file) |
|
im.verify() |
|
shape = exif_size(im) |
|
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}' |
|
if im.format.lower() in ('jpg', 'jpeg'): |
|
with open(im_file, 'rb') as f: |
|
f.seek(-2, 2) |
|
if f.read() != b'\xff\xd9': |
|
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' |
|
|
|
|
|
if os.path.isfile(lb_file): |
|
nf = 1 |
|
with open(lb_file) as f: |
|
lb = [x.split() for x in f.read().strip().splitlines() if len(x)] |
|
if any(len(x) > 6 for x in lb): |
|
classes = np.array([x[0] for x in lb], dtype=np.float32) |
|
segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in lb] |
|
lb = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) |
|
lb = np.array(lb, dtype=np.float32) |
|
nl = len(lb) |
|
if nl: |
|
assert lb.shape[1] == 5, f'labels require 5 columns, {lb.shape[1]} columns detected' |
|
assert (lb >= 0).all(), f'negative label values {lb[lb < 0]}' |
|
assert (lb[:, 1:] <= 1).all(), f'non-normalized or out of bounds coordinates {lb[:, 1:][lb[:, 1:] > 1]}' |
|
_, i = np.unique(lb, axis=0, return_index=True) |
|
if len(i) < nl: |
|
lb = lb[i] |
|
if segments: |
|
segments = [segments[x] for x in i] |
|
msg = f'{prefix}WARNING ⚠️ {im_file}: {nl - len(i)} duplicate labels removed' |
|
else: |
|
ne = 1 |
|
lb = np.zeros((0, 5), dtype=np.float32) |
|
else: |
|
nm = 1 |
|
lb = np.zeros((0, 5), dtype=np.float32) |
|
return im_file, lb, shape, segments, nm, nf, ne, nc, msg |
|
except Exception as e: |
|
nc = 1 |
|
msg = f'{prefix}WARNING ⚠️ {im_file}: ignoring corrupt image/label: {e}' |
|
return [None, None, None, None, nm, nf, ne, nc, msg] |
|
|