|
|
@@ -56,7 +56,7 @@ def exif_size(img): |
|
|
|
|
|
|
|
|
|
|
|
def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False, |
|
|
|
rank=-1, world_size=1, workers=8, image_weights=False, quad=False): |
|
|
|
rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix=''): |
|
|
|
# Make sure only the first process in DDP process the dataset first, and the following others can use the cache |
|
|
|
with torch_distributed_zero_first(rank): |
|
|
|
dataset = LoadImagesAndLabels(path, imgsz, batch_size, |
|
|
@@ -67,8 +67,8 @@ def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=Fa |
|
|
|
single_cls=opt.single_cls, |
|
|
|
stride=int(stride), |
|
|
|
pad=pad, |
|
|
|
rank=rank, |
|
|
|
image_weights=image_weights) |
|
|
|
image_weights=image_weights, |
|
|
|
prefix=prefix) |
|
|
|
|
|
|
|
batch_size = min(batch_size, len(dataset)) |
|
|
|
nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers]) # number of workers |
|
|
@@ -129,7 +129,7 @@ class LoadImages: # for inference |
|
|
|
elif os.path.isfile(p): |
|
|
|
files = [p] # files |
|
|
|
else: |
|
|
|
raise Exception('ERROR: %s does not exist' % p) |
|
|
|
raise Exception(f'ERROR: {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] |
|
|
@@ -144,8 +144,8 @@ class LoadImages: # for inference |
|
|
|
self.new_video(videos[0]) # new video |
|
|
|
else: |
|
|
|
self.cap = None |
|
|
|
assert self.nf > 0, 'No images or videos found in %s. Supported formats are:\nimages: %s\nvideos: %s' % \ |
|
|
|
(p, img_formats, vid_formats) |
|
|
|
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 |
|
|
@@ -171,14 +171,14 @@ class LoadImages: # for inference |
|
|
|
ret_val, img0 = self.cap.read() |
|
|
|
|
|
|
|
self.frame += 1 |
|
|
|
print('video %g/%g (%g/%g) %s: ' % (self.count + 1, self.nf, self.frame, self.nframes, path), end='') |
|
|
|
print(f'video {self.count + 1}/{self.nf} ({self.frame}/{self.nframes}) {path}: ', end='') |
|
|
|
|
|
|
|
else: |
|
|
|
# Read image |
|
|
|
self.count += 1 |
|
|
|
img0 = cv2.imread(path) # BGR |
|
|
|
assert img0 is not None, 'Image Not Found ' + path |
|
|
|
print('image %g/%g %s: ' % (self.count, self.nf, path), end='') |
|
|
|
print(f'image {self.count}/{self.nf} {path}: ', end='') |
|
|
|
|
|
|
|
# Padded resize |
|
|
|
img = letterbox(img0, new_shape=self.img_size)[0] |
|
|
@@ -238,9 +238,9 @@ class LoadWebcam: # for inference |
|
|
|
break |
|
|
|
|
|
|
|
# Print |
|
|
|
assert ret_val, 'Camera Error %s' % self.pipe |
|
|
|
assert ret_val, f'Camera Error {self.pipe}' |
|
|
|
img_path = 'webcam.jpg' |
|
|
|
print('webcam %g: ' % self.count, end='') |
|
|
|
print(f'webcam {self.count}: ', end='') |
|
|
|
|
|
|
|
# Padded resize |
|
|
|
img = letterbox(img0, new_shape=self.img_size)[0] |
|
|
@@ -271,15 +271,15 @@ class LoadStreams: # multiple IP or RTSP cameras |
|
|
|
self.sources = [clean_str(x) for x in sources] # clean source names for later |
|
|
|
for i, s in enumerate(sources): |
|
|
|
# Start the thread to read frames from the video stream |
|
|
|
print('%g/%g: %s... ' % (i + 1, n, s), end='') |
|
|
|
print(f'{i + 1}/{n}: {s}... ', end='') |
|
|
|
cap = cv2.VideoCapture(eval(s) if s.isnumeric() else s) |
|
|
|
assert cap.isOpened(), 'Failed to open %s' % s |
|
|
|
assert cap.isOpened(), f'Failed to open {s}' |
|
|
|
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
|
|
|
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) |
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS) % 100 |
|
|
|
_, self.imgs[i] = cap.read() # guarantee first frame |
|
|
|
thread = Thread(target=self.update, args=([i, cap]), daemon=True) |
|
|
|
print(' success (%gx%g at %.2f FPS).' % (w, h, fps)) |
|
|
|
print(f' success ({w}x{h} at {fps:.2f} FPS).') |
|
|
|
thread.start() |
|
|
|
print('') # newline |
|
|
|
|
|
|
@@ -336,7 +336,7 @@ def img2label_paths(img_paths): |
|
|
|
|
|
|
|
class LoadImagesAndLabels(Dataset): # 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, rank=-1): |
|
|
|
cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''): |
|
|
|
self.img_size = img_size |
|
|
|
self.augment = augment |
|
|
|
self.hyp = hyp |
|
|
@@ -358,11 +358,11 @@ class LoadImagesAndLabels(Dataset): # for training/testing |
|
|
|
parent = str(p.parent) + os.sep |
|
|
|
f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path |
|
|
|
else: |
|
|
|
raise Exception('%s does not exist' % p) |
|
|
|
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]) |
|
|
|
assert self.img_files, 'No images found' |
|
|
|
assert self.img_files, f'{prefix}No images found' |
|
|
|
except Exception as e: |
|
|
|
raise Exception('Error loading data from %s: %s\nSee %s' % (path, e, help_url)) |
|
|
|
raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}') |
|
|
|
|
|
|
|
# Check cache |
|
|
|
self.label_files = img2label_paths(self.img_files) # labels |
|
|
@@ -370,15 +370,15 @@ class LoadImagesAndLabels(Dataset): # for training/testing |
|
|
|
if cache_path.is_file(): |
|
|
|
cache = torch.load(cache_path) # load |
|
|
|
if cache['hash'] != get_hash(self.label_files + self.img_files) or 'results' not in cache: # changed |
|
|
|
cache = self.cache_labels(cache_path) # re-cache |
|
|
|
cache = self.cache_labels(cache_path, prefix) # re-cache |
|
|
|
else: |
|
|
|
cache = self.cache_labels(cache_path) # cache |
|
|
|
cache = self.cache_labels(cache_path, prefix) # cache |
|
|
|
|
|
|
|
# Display cache |
|
|
|
[nf, nm, ne, nc, n] = cache.pop('results') # found, missing, empty, corrupted, total |
|
|
|
desc = f"Scanning '{cache_path}' for images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted" |
|
|
|
tqdm(None, desc=desc, total=n, initial=n) |
|
|
|
assert nf > 0 or not augment, f'No labels found in {cache_path}. Can not train without labels. See {help_url}' |
|
|
|
tqdm(None, desc=prefix + desc, total=n, initial=n) |
|
|
|
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 |
|
|
@@ -432,9 +432,9 @@ class LoadImagesAndLabels(Dataset): # for training/testing |
|
|
|
for i, x in pbar: |
|
|
|
self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # img, hw_original, hw_resized = load_image(self, i) |
|
|
|
gb += self.imgs[i].nbytes |
|
|
|
pbar.desc = 'Caching images (%.1fGB)' % (gb / 1E9) |
|
|
|
pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB)' |
|
|
|
|
|
|
|
def cache_labels(self, path=Path('./labels.cache')): |
|
|
|
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 |
|
|
@@ -466,18 +466,18 @@ class LoadImagesAndLabels(Dataset): # for training/testing |
|
|
|
x[im_file] = [l, shape] |
|
|
|
except Exception as e: |
|
|
|
nc += 1 |
|
|
|
print('WARNING: Ignoring corrupted image and/or label %s: %s' % (im_file, e)) |
|
|
|
print(f'{prefix}WARNING: Ignoring corrupted image and/or label {im_file}: {e}') |
|
|
|
|
|
|
|
pbar.desc = f"Scanning '{path.parent / path.stem}' for images and labels... " \ |
|
|
|
pbar.desc = f"{prefix}Scanning '{path.parent / path.stem}' for images and labels... " \ |
|
|
|
f"{nf} found, {nm} missing, {ne} empty, {nc} corrupted" |
|
|
|
|
|
|
|
if nf == 0: |
|
|
|
print(f'WARNING: No labels found in {path}. See {help_url}') |
|
|
|
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] |
|
|
|
torch.save(x, path) # save for next time |
|
|
|
logging.info(f"New cache created: {path}") |
|
|
|
logging.info(f'{prefix}New cache created: {path}') |
|
|
|
return x |
|
|
|
|
|
|
|
def __len__(self): |