megadetector 5.0.5__py3-none-any.whl → 5.0.7__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of megadetector might be problematic. Click here for more details.

Files changed (132) hide show
  1. api/batch_processing/data_preparation/manage_local_batch.py +302 -263
  2. api/batch_processing/data_preparation/manage_video_batch.py +81 -2
  3. api/batch_processing/postprocessing/add_max_conf.py +1 -0
  4. api/batch_processing/postprocessing/categorize_detections_by_size.py +50 -19
  5. api/batch_processing/postprocessing/compare_batch_results.py +110 -60
  6. api/batch_processing/postprocessing/load_api_results.py +56 -70
  7. api/batch_processing/postprocessing/md_to_coco.py +1 -1
  8. api/batch_processing/postprocessing/md_to_labelme.py +2 -1
  9. api/batch_processing/postprocessing/postprocess_batch_results.py +240 -81
  10. api/batch_processing/postprocessing/render_detection_confusion_matrix.py +625 -0
  11. api/batch_processing/postprocessing/repeat_detection_elimination/find_repeat_detections.py +71 -23
  12. api/batch_processing/postprocessing/repeat_detection_elimination/remove_repeat_detections.py +1 -1
  13. api/batch_processing/postprocessing/repeat_detection_elimination/repeat_detections_core.py +227 -75
  14. api/batch_processing/postprocessing/subset_json_detector_output.py +132 -5
  15. api/batch_processing/postprocessing/top_folders_to_bottom.py +1 -1
  16. api/synchronous/api_core/animal_detection_api/detection/run_detector_batch.py +2 -2
  17. classification/prepare_classification_script.py +191 -191
  18. data_management/coco_to_yolo.py +68 -45
  19. data_management/databases/integrity_check_json_db.py +7 -5
  20. data_management/generate_crops_from_cct.py +3 -3
  21. data_management/get_image_sizes.py +8 -6
  22. data_management/importers/add_timestamps_to_icct.py +79 -0
  23. data_management/importers/animl_results_to_md_results.py +160 -0
  24. data_management/importers/auckland_doc_test_to_json.py +4 -4
  25. data_management/importers/auckland_doc_to_json.py +1 -1
  26. data_management/importers/awc_to_json.py +5 -5
  27. data_management/importers/bellevue_to_json.py +5 -5
  28. data_management/importers/carrizo_shrubfree_2018.py +5 -5
  29. data_management/importers/carrizo_trail_cam_2017.py +5 -5
  30. data_management/importers/cct_field_adjustments.py +2 -3
  31. data_management/importers/channel_islands_to_cct.py +4 -4
  32. data_management/importers/ena24_to_json.py +5 -5
  33. data_management/importers/helena_to_cct.py +10 -10
  34. data_management/importers/idaho-camera-traps.py +12 -12
  35. data_management/importers/idfg_iwildcam_lila_prep.py +8 -8
  36. data_management/importers/jb_csv_to_json.py +4 -4
  37. data_management/importers/missouri_to_json.py +1 -1
  38. data_management/importers/noaa_seals_2019.py +1 -1
  39. data_management/importers/pc_to_json.py +5 -5
  40. data_management/importers/prepare-noaa-fish-data-for-lila.py +4 -4
  41. data_management/importers/prepare_zsl_imerit.py +5 -5
  42. data_management/importers/rspb_to_json.py +4 -4
  43. data_management/importers/save_the_elephants_survey_A.py +5 -5
  44. data_management/importers/save_the_elephants_survey_B.py +6 -6
  45. data_management/importers/snapshot_safari_importer.py +9 -9
  46. data_management/importers/snapshot_serengeti_lila.py +9 -9
  47. data_management/importers/timelapse_csv_set_to_json.py +5 -7
  48. data_management/importers/ubc_to_json.py +4 -4
  49. data_management/importers/umn_to_json.py +4 -4
  50. data_management/importers/wellington_to_json.py +1 -1
  51. data_management/importers/wi_to_json.py +2 -2
  52. data_management/importers/zamba_results_to_md_results.py +181 -0
  53. data_management/labelme_to_coco.py +35 -7
  54. data_management/labelme_to_yolo.py +229 -0
  55. data_management/lila/add_locations_to_island_camera_traps.py +1 -1
  56. data_management/lila/add_locations_to_nacti.py +147 -0
  57. data_management/lila/create_lila_blank_set.py +474 -0
  58. data_management/lila/create_lila_test_set.py +2 -1
  59. data_management/lila/create_links_to_md_results_files.py +106 -0
  60. data_management/lila/download_lila_subset.py +46 -21
  61. data_management/lila/generate_lila_per_image_labels.py +23 -14
  62. data_management/lila/get_lila_annotation_counts.py +17 -11
  63. data_management/lila/lila_common.py +14 -11
  64. data_management/lila/test_lila_metadata_urls.py +116 -0
  65. data_management/ocr_tools.py +829 -0
  66. data_management/resize_coco_dataset.py +13 -11
  67. data_management/yolo_output_to_md_output.py +84 -12
  68. data_management/yolo_to_coco.py +38 -20
  69. detection/process_video.py +36 -14
  70. detection/pytorch_detector.py +23 -8
  71. detection/run_detector.py +76 -19
  72. detection/run_detector_batch.py +178 -63
  73. detection/run_inference_with_yolov5_val.py +326 -57
  74. detection/run_tiled_inference.py +153 -43
  75. detection/video_utils.py +34 -8
  76. md_utils/ct_utils.py +172 -1
  77. md_utils/md_tests.py +372 -51
  78. md_utils/path_utils.py +167 -39
  79. md_utils/process_utils.py +26 -7
  80. md_utils/split_locations_into_train_val.py +215 -0
  81. md_utils/string_utils.py +10 -0
  82. md_utils/url_utils.py +0 -2
  83. md_utils/write_html_image_list.py +9 -26
  84. md_visualization/plot_utils.py +12 -8
  85. md_visualization/visualization_utils.py +106 -7
  86. md_visualization/visualize_db.py +16 -8
  87. md_visualization/visualize_detector_output.py +208 -97
  88. {megadetector-5.0.5.dist-info → megadetector-5.0.7.dist-info}/METADATA +3 -6
  89. {megadetector-5.0.5.dist-info → megadetector-5.0.7.dist-info}/RECORD +98 -121
  90. {megadetector-5.0.5.dist-info → megadetector-5.0.7.dist-info}/WHEEL +1 -1
  91. taxonomy_mapping/map_lila_taxonomy_to_wi_taxonomy.py +1 -1
  92. taxonomy_mapping/map_new_lila_datasets.py +43 -39
  93. taxonomy_mapping/prepare_lila_taxonomy_release.py +5 -2
  94. taxonomy_mapping/preview_lila_taxonomy.py +27 -27
  95. taxonomy_mapping/species_lookup.py +33 -13
  96. taxonomy_mapping/taxonomy_csv_checker.py +7 -5
  97. api/synchronous/api_core/yolov5/detect.py +0 -252
  98. api/synchronous/api_core/yolov5/export.py +0 -607
  99. api/synchronous/api_core/yolov5/hubconf.py +0 -146
  100. api/synchronous/api_core/yolov5/models/__init__.py +0 -0
  101. api/synchronous/api_core/yolov5/models/common.py +0 -738
  102. api/synchronous/api_core/yolov5/models/experimental.py +0 -104
  103. api/synchronous/api_core/yolov5/models/tf.py +0 -574
  104. api/synchronous/api_core/yolov5/models/yolo.py +0 -338
  105. api/synchronous/api_core/yolov5/train.py +0 -670
  106. api/synchronous/api_core/yolov5/utils/__init__.py +0 -36
  107. api/synchronous/api_core/yolov5/utils/activations.py +0 -103
  108. api/synchronous/api_core/yolov5/utils/augmentations.py +0 -284
  109. api/synchronous/api_core/yolov5/utils/autoanchor.py +0 -170
  110. api/synchronous/api_core/yolov5/utils/autobatch.py +0 -66
  111. api/synchronous/api_core/yolov5/utils/aws/__init__.py +0 -0
  112. api/synchronous/api_core/yolov5/utils/aws/resume.py +0 -40
  113. api/synchronous/api_core/yolov5/utils/benchmarks.py +0 -148
  114. api/synchronous/api_core/yolov5/utils/callbacks.py +0 -71
  115. api/synchronous/api_core/yolov5/utils/dataloaders.py +0 -1087
  116. api/synchronous/api_core/yolov5/utils/downloads.py +0 -178
  117. api/synchronous/api_core/yolov5/utils/flask_rest_api/example_request.py +0 -19
  118. api/synchronous/api_core/yolov5/utils/flask_rest_api/restapi.py +0 -46
  119. api/synchronous/api_core/yolov5/utils/general.py +0 -1018
  120. api/synchronous/api_core/yolov5/utils/loggers/__init__.py +0 -187
  121. api/synchronous/api_core/yolov5/utils/loggers/wandb/__init__.py +0 -0
  122. api/synchronous/api_core/yolov5/utils/loggers/wandb/log_dataset.py +0 -27
  123. api/synchronous/api_core/yolov5/utils/loggers/wandb/sweep.py +0 -41
  124. api/synchronous/api_core/yolov5/utils/loggers/wandb/wandb_utils.py +0 -577
  125. api/synchronous/api_core/yolov5/utils/loss.py +0 -234
  126. api/synchronous/api_core/yolov5/utils/metrics.py +0 -355
  127. api/synchronous/api_core/yolov5/utils/plots.py +0 -489
  128. api/synchronous/api_core/yolov5/utils/torch_utils.py +0 -314
  129. api/synchronous/api_core/yolov5/val.py +0 -394
  130. md_utils/matlab_porting_tools.py +0 -97
  131. {megadetector-5.0.5.dist-info → megadetector-5.0.7.dist-info}/LICENSE +0 -0
  132. {megadetector-5.0.5.dist-info → megadetector-5.0.7.dist-info}/top_level.txt +0 -0
@@ -1,1087 +0,0 @@
1
- # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2
- """
3
- Dataloaders and dataset utils
4
- """
5
-
6
- import glob
7
- import hashlib
8
- import json
9
- import math
10
- import os
11
- import random
12
- import shutil
13
- import time
14
- from itertools import repeat
15
- from multiprocessing.pool import Pool, ThreadPool
16
- from pathlib import Path
17
- from threading import Thread
18
- from urllib.parse import urlparse
19
- from zipfile import ZipFile
20
-
21
- import numpy as np
22
- import torch
23
- import torch.nn.functional as F
24
- import yaml
25
- from PIL import ExifTags, Image, ImageOps
26
- from torch.utils.data import DataLoader, Dataset, dataloader, distributed
27
- from tqdm import tqdm
28
-
29
- from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective
30
- from utils.general import (DATASETS_DIR, LOGGER, NUM_THREADS, check_dataset, check_requirements, check_yaml, clean_str,
31
- cv2, segments2boxes, xyn2xy, xywh2xyxy, xywhn2xyxy, xyxy2xywhn)
32
- from utils.torch_utils import torch_distributed_zero_first
33
-
34
- # Parameters
35
- HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
36
- IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp' # include image suffixes
37
- VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' # include video suffixes
38
- BAR_FORMAT = '{l_bar}{bar:10}{r_bar}{bar:-10b}' # tqdm bar format
39
- LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
40
-
41
- # Get orientation exif tag
42
- for orientation in ExifTags.TAGS.keys():
43
- if ExifTags.TAGS[orientation] == 'Orientation':
44
- break
45
-
46
-
47
- def get_hash(paths):
48
- # Returns a single hash value of a list of paths (files or dirs)
49
- size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes
50
- h = hashlib.md5(str(size).encode()) # hash sizes
51
- h.update(''.join(paths).encode()) # hash paths
52
- return h.hexdigest() # return hash
53
-
54
-
55
- def exif_size(img):
56
- # Returns exif-corrected PIL size
57
- s = img.size # (width, height)
58
- try:
59
- rotation = dict(img._getexif().items())[orientation]
60
- if rotation in [6, 8]: # rotation 270 or 90
61
- s = (s[1], s[0])
62
- except Exception:
63
- pass
64
-
65
- return s
66
-
67
-
68
- def exif_transpose(image):
69
- """
70
- Transpose a PIL image accordingly if it has an EXIF Orientation tag.
71
- Inplace version of https://github.com/python-pillow/Pillow/blob/master/src/PIL/ImageOps.py exif_transpose()
72
-
73
- :param image: The image to transpose.
74
- :return: An image.
75
- """
76
- exif = image.getexif()
77
- orientation = exif.get(0x0112, 1) # default 1
78
- if orientation > 1:
79
- method = {
80
- 2: Image.FLIP_LEFT_RIGHT,
81
- 3: Image.ROTATE_180,
82
- 4: Image.FLIP_TOP_BOTTOM,
83
- 5: Image.TRANSPOSE,
84
- 6: Image.ROTATE_270,
85
- 7: Image.TRANSVERSE,
86
- 8: Image.ROTATE_90,}.get(orientation)
87
- if method is not None:
88
- image = image.transpose(method)
89
- del exif[0x0112]
90
- image.info["exif"] = exif.tobytes()
91
- return image
92
-
93
-
94
- def create_dataloader(path,
95
- imgsz,
96
- batch_size,
97
- stride,
98
- single_cls=False,
99
- hyp=None,
100
- augment=False,
101
- cache=False,
102
- pad=0.0,
103
- rect=False,
104
- rank=-1,
105
- workers=8,
106
- image_weights=False,
107
- quad=False,
108
- prefix='',
109
- shuffle=False):
110
- if rect and shuffle:
111
- LOGGER.warning('WARNING: --rect is incompatible with DataLoader shuffle, setting shuffle=False')
112
- shuffle = False
113
- with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
114
- dataset = LoadImagesAndLabels(
115
- path,
116
- imgsz,
117
- batch_size,
118
- augment=augment, # augmentation
119
- hyp=hyp, # hyperparameters
120
- rect=rect, # rectangular batches
121
- cache_images=cache,
122
- single_cls=single_cls,
123
- stride=int(stride),
124
- pad=pad,
125
- image_weights=image_weights,
126
- prefix=prefix)
127
-
128
- batch_size = min(batch_size, len(dataset))
129
- nd = torch.cuda.device_count() # number of CUDA devices
130
- nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers
131
- sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
132
- loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates
133
- return loader(dataset,
134
- batch_size=batch_size,
135
- shuffle=shuffle and sampler is None,
136
- num_workers=nw,
137
- sampler=sampler,
138
- pin_memory=True,
139
- collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn), dataset
140
-
141
-
142
- class InfiniteDataLoader(dataloader.DataLoader):
143
- """ Dataloader that reuses workers
144
-
145
- Uses same syntax as vanilla DataLoader
146
- """
147
-
148
- def __init__(self, *args, **kwargs):
149
- super().__init__(*args, **kwargs)
150
- object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
151
- self.iterator = super().__iter__()
152
-
153
- def __len__(self):
154
- return len(self.batch_sampler.sampler)
155
-
156
- def __iter__(self):
157
- for _ in range(len(self)):
158
- yield next(self.iterator)
159
-
160
-
161
- class _RepeatSampler:
162
- """ Sampler that repeats forever
163
-
164
- Args:
165
- sampler (Sampler)
166
- """
167
-
168
- def __init__(self, sampler):
169
- self.sampler = sampler
170
-
171
- def __iter__(self):
172
- while True:
173
- yield from iter(self.sampler)
174
-
175
-
176
- class LoadImages:
177
- # YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`
178
- def __init__(self, path, img_size=640, stride=32, auto=True):
179
- p = str(Path(path).resolve()) # os-agnostic absolute path
180
- if '*' in p:
181
- files = sorted(glob.glob(p, recursive=True)) # glob
182
- elif os.path.isdir(p):
183
- files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
184
- elif os.path.isfile(p):
185
- files = [p] # files
186
- else:
187
- raise Exception(f'ERROR: {p} does not exist')
188
-
189
- images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
190
- videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
191
- ni, nv = len(images), len(videos)
192
-
193
- self.img_size = img_size
194
- self.stride = stride
195
- self.files = images + videos
196
- self.nf = ni + nv # number of files
197
- self.video_flag = [False] * ni + [True] * nv
198
- self.mode = 'image'
199
- self.auto = auto
200
- if any(videos):
201
- self.new_video(videos[0]) # new video
202
- else:
203
- self.cap = None
204
- assert self.nf > 0, f'No images or videos found in {p}. ' \
205
- f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}'
206
-
207
- def __iter__(self):
208
- self.count = 0
209
- return self
210
-
211
- def __next__(self):
212
- if self.count == self.nf:
213
- raise StopIteration
214
- path = self.files[self.count]
215
-
216
- if self.video_flag[self.count]:
217
- # Read video
218
- self.mode = 'video'
219
- ret_val, img0 = self.cap.read()
220
- while not ret_val:
221
- self.count += 1
222
- self.cap.release()
223
- if self.count == self.nf: # last video
224
- raise StopIteration
225
- path = self.files[self.count]
226
- self.new_video(path)
227
- ret_val, img0 = self.cap.read()
228
-
229
- self.frame += 1
230
- s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '
231
-
232
- else:
233
- # Read image
234
- self.count += 1
235
- img0 = cv2.imread(path) # BGR
236
- assert img0 is not None, f'Image Not Found {path}'
237
- s = f'image {self.count}/{self.nf} {path}: '
238
-
239
- # Padded resize
240
- img = letterbox(img0, self.img_size, stride=self.stride, auto=self.auto)[0]
241
-
242
- # Convert
243
- img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
244
- img = np.ascontiguousarray(img)
245
-
246
- return path, img, img0, self.cap, s
247
-
248
- def new_video(self, path):
249
- self.frame = 0
250
- self.cap = cv2.VideoCapture(path)
251
- self.frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
252
-
253
- def __len__(self):
254
- return self.nf # number of files
255
-
256
-
257
- class LoadWebcam: # for inference
258
- # YOLOv5 local webcam dataloader, i.e. `python detect.py --source 0`
259
- def __init__(self, pipe='0', img_size=640, stride=32):
260
- self.img_size = img_size
261
- self.stride = stride
262
- self.pipe = eval(pipe) if pipe.isnumeric() else pipe
263
- self.cap = cv2.VideoCapture(self.pipe) # video capture object
264
- self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
265
-
266
- def __iter__(self):
267
- self.count = -1
268
- return self
269
-
270
- def __next__(self):
271
- self.count += 1
272
- if cv2.waitKey(1) == ord('q'): # q to quit
273
- self.cap.release()
274
- cv2.destroyAllWindows()
275
- raise StopIteration
276
-
277
- # Read frame
278
- ret_val, img0 = self.cap.read()
279
- img0 = cv2.flip(img0, 1) # flip left-right
280
-
281
- # Print
282
- assert ret_val, f'Camera Error {self.pipe}'
283
- img_path = 'webcam.jpg'
284
- s = f'webcam {self.count}: '
285
-
286
- # Padded resize
287
- img = letterbox(img0, self.img_size, stride=self.stride)[0]
288
-
289
- # Convert
290
- img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
291
- img = np.ascontiguousarray(img)
292
-
293
- return img_path, img, img0, None, s
294
-
295
- def __len__(self):
296
- return 0
297
-
298
-
299
- class LoadStreams:
300
- # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams`
301
- def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True):
302
- self.mode = 'stream'
303
- self.img_size = img_size
304
- self.stride = stride
305
-
306
- if os.path.isfile(sources):
307
- with open(sources) as f:
308
- sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]
309
- else:
310
- sources = [sources]
311
-
312
- n = len(sources)
313
- self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
314
- self.sources = [clean_str(x) for x in sources] # clean source names for later
315
- self.auto = auto
316
- for i, s in enumerate(sources): # index, source
317
- # Start thread to read frames from video stream
318
- st = f'{i + 1}/{n}: {s}... '
319
- if urlparse(s).hostname in ('www.youtube.com', 'youtube.com', 'youtu.be'): # if source is YouTube video
320
- check_requirements(('pafy', 'youtube_dl==2020.12.2'))
321
- import pafy
322
- s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL
323
- s = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam
324
- cap = cv2.VideoCapture(s)
325
- assert cap.isOpened(), f'{st}Failed to open {s}'
326
- w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
327
- h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
328
- fps = cap.get(cv2.CAP_PROP_FPS) # warning: may return 0 or nan
329
- self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback
330
- self.fps[i] = max((fps if math.isfinite(fps) else 0) % 100, 0) or 30 # 30 FPS fallback
331
-
332
- _, self.imgs[i] = cap.read() # guarantee first frame
333
- self.threads[i] = Thread(target=self.update, args=([i, cap, s]), daemon=True)
334
- LOGGER.info(f"{st} Success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)")
335
- self.threads[i].start()
336
- LOGGER.info('') # newline
337
-
338
- # check for common shapes
339
- s = np.stack([letterbox(x, self.img_size, stride=self.stride, auto=self.auto)[0].shape for x in self.imgs])
340
- self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal
341
- if not self.rect:
342
- LOGGER.warning('WARNING: Stream shapes differ. For optimal performance supply similarly-shaped streams.')
343
-
344
- def update(self, i, cap, stream):
345
- # Read stream `i` frames in daemon thread
346
- n, f, read = 0, self.frames[i], 1 # frame number, frame array, inference every 'read' frame
347
- while cap.isOpened() and n < f:
348
- n += 1
349
- # _, self.imgs[index] = cap.read()
350
- cap.grab()
351
- if n % read == 0:
352
- success, im = cap.retrieve()
353
- if success:
354
- self.imgs[i] = im
355
- else:
356
- LOGGER.warning('WARNING: Video stream unresponsive, please check your IP camera connection.')
357
- self.imgs[i] = np.zeros_like(self.imgs[i])
358
- cap.open(stream) # re-open stream if signal was lost
359
- time.sleep(1 / self.fps[i]) # wait time
360
-
361
- def __iter__(self):
362
- self.count = -1
363
- return self
364
-
365
- def __next__(self):
366
- self.count += 1
367
- if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit
368
- cv2.destroyAllWindows()
369
- raise StopIteration
370
-
371
- # Letterbox
372
- img0 = self.imgs.copy()
373
- img = [letterbox(x, self.img_size, stride=self.stride, auto=self.rect and self.auto)[0] for x in img0]
374
-
375
- # Stack
376
- img = np.stack(img, 0)
377
-
378
- # Convert
379
- img = img[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW
380
- img = np.ascontiguousarray(img)
381
-
382
- return self.sources, img, img0, None, ''
383
-
384
- def __len__(self):
385
- return len(self.sources) # 1E12 frames = 32 streams at 30 FPS for 30 years
386
-
387
-
388
- def img2label_paths(img_paths):
389
- # Define label paths as a function of image paths
390
- sa, sb = f'{os.sep}images{os.sep}', f'{os.sep}labels{os.sep}' # /images/, /labels/ substrings
391
- return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths]
392
-
393
-
394
- class LoadImagesAndLabels(Dataset):
395
- # YOLOv5 train_loader/val_loader, loads images and labels for training and validation
396
- cache_version = 0.6 # dataset labels *.cache version
397
- rand_interp_methods = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4]
398
-
399
- def __init__(self,
400
- path,
401
- img_size=640,
402
- batch_size=16,
403
- augment=False,
404
- hyp=None,
405
- rect=False,
406
- image_weights=False,
407
- cache_images=False,
408
- single_cls=False,
409
- stride=32,
410
- pad=0.0,
411
- prefix=''):
412
- self.img_size = img_size
413
- self.augment = augment
414
- self.hyp = hyp
415
- self.image_weights = image_weights
416
- self.rect = False if image_weights else rect
417
- self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
418
- self.mosaic_border = [-img_size // 2, -img_size // 2]
419
- self.stride = stride
420
- self.path = path
421
- self.albumentations = Albumentations() if augment else None
422
-
423
- try:
424
- f = [] # image files
425
- for p in path if isinstance(path, list) else [path]:
426
- p = Path(p) # os-agnostic
427
- if p.is_dir(): # dir
428
- f += glob.glob(str(p / '**' / '*.*'), recursive=True)
429
- # f = list(p.rglob('*.*')) # pathlib
430
- elif p.is_file(): # file
431
- with open(p) as t:
432
- t = t.read().strip().splitlines()
433
- parent = str(p.parent) + os.sep
434
- f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
435
- # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
436
- else:
437
- raise Exception(f'{prefix}{p} does not exist')
438
- self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
439
- # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib
440
- assert self.im_files, f'{prefix}No images found'
441
- except Exception as e:
442
- raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')
443
-
444
- # Check cache
445
- self.label_files = img2label_paths(self.im_files) # labels
446
- cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')
447
- try:
448
- cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict
449
- assert cache['version'] == self.cache_version # same version
450
- assert cache['hash'] == get_hash(self.label_files + self.im_files) # same hash
451
- except Exception:
452
- cache, exists = self.cache_labels(cache_path, prefix), False # cache
453
-
454
- # Display cache
455
- nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupt, total
456
- if exists and LOCAL_RANK in {-1, 0}:
457
- d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupt"
458
- tqdm(None, desc=prefix + d, total=n, initial=n, bar_format=BAR_FORMAT) # display cache results
459
- if cache['msgs']:
460
- LOGGER.info('\n'.join(cache['msgs'])) # display warnings
461
- assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}'
462
-
463
- # Read cache
464
- [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items
465
- labels, shapes, self.segments = zip(*cache.values())
466
- self.labels = list(labels)
467
- self.shapes = np.array(shapes, dtype=np.float64)
468
- self.im_files = list(cache.keys()) # update
469
- self.label_files = img2label_paths(cache.keys()) # update
470
- n = len(shapes) # number of images
471
- bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
472
- nb = bi[-1] + 1 # number of batches
473
- self.batch = bi # batch index of image
474
- self.n = n
475
- self.indices = range(n)
476
-
477
- # Update labels
478
- include_class = [] # filter labels to include only these classes (optional)
479
- include_class_array = np.array(include_class).reshape(1, -1)
480
- for i, (label, segment) in enumerate(zip(self.labels, self.segments)):
481
- if include_class:
482
- j = (label[:, 0:1] == include_class_array).any(1)
483
- self.labels[i] = label[j]
484
- if segment:
485
- self.segments[i] = segment[j]
486
- if single_cls: # single-class training, merge all classes into 0
487
- self.labels[i][:, 0] = 0
488
- if segment:
489
- self.segments[i][:, 0] = 0
490
-
491
- # Rectangular Training
492
- if self.rect:
493
- # Sort by aspect ratio
494
- s = self.shapes # wh
495
- ar = s[:, 1] / s[:, 0] # aspect ratio
496
- irect = ar.argsort()
497
- self.im_files = [self.im_files[i] for i in irect]
498
- self.label_files = [self.label_files[i] for i in irect]
499
- self.labels = [self.labels[i] for i in irect]
500
- self.shapes = s[irect] # wh
501
- ar = ar[irect]
502
-
503
- # Set training image shapes
504
- shapes = [[1, 1]] * nb
505
- for i in range(nb):
506
- ari = ar[bi == i]
507
- mini, maxi = ari.min(), ari.max()
508
- if maxi < 1:
509
- shapes[i] = [maxi, 1]
510
- elif mini > 1:
511
- shapes[i] = [1, 1 / mini]
512
-
513
- self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
514
-
515
- # Cache images into RAM/disk for faster training (WARNING: large datasets may exceed system resources)
516
- self.ims = [None] * n
517
- self.npy_files = [Path(f).with_suffix('.npy') for f in self.im_files]
518
- if cache_images:
519
- gb = 0 # Gigabytes of cached images
520
- self.im_hw0, self.im_hw = [None] * n, [None] * n
521
- fcn = self.cache_images_to_disk if cache_images == 'disk' else self.load_image
522
- results = ThreadPool(NUM_THREADS).imap(fcn, range(n))
523
- pbar = tqdm(enumerate(results), total=n, bar_format=BAR_FORMAT, disable=LOCAL_RANK > 0)
524
- for i, x in pbar:
525
- if cache_images == 'disk':
526
- gb += self.npy_files[i].stat().st_size
527
- else: # 'ram'
528
- self.ims[i], self.im_hw0[i], self.im_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i)
529
- gb += self.ims[i].nbytes
530
- pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})'
531
- pbar.close()
532
-
533
- def cache_labels(self, path=Path('./labels.cache'), prefix=''):
534
- # Cache dataset labels, check images and read shapes
535
- x = {} # dict
536
- nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages
537
- desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..."
538
- with Pool(NUM_THREADS) as pool:
539
- pbar = tqdm(pool.imap(verify_image_label, zip(self.im_files, self.label_files, repeat(prefix))),
540
- desc=desc,
541
- total=len(self.im_files),
542
- bar_format=BAR_FORMAT)
543
- for im_file, lb, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar:
544
- nm += nm_f
545
- nf += nf_f
546
- ne += ne_f
547
- nc += nc_f
548
- if im_file:
549
- x[im_file] = [lb, shape, segments]
550
- if msg:
551
- msgs.append(msg)
552
- pbar.desc = f"{desc}{nf} found, {nm} missing, {ne} empty, {nc} corrupt"
553
-
554
- pbar.close()
555
- if msgs:
556
- LOGGER.info('\n'.join(msgs))
557
- if nf == 0:
558
- LOGGER.warning(f'{prefix}WARNING: No labels found in {path}. See {HELP_URL}')
559
- x['hash'] = get_hash(self.label_files + self.im_files)
560
- x['results'] = nf, nm, ne, nc, len(self.im_files)
561
- x['msgs'] = msgs # warnings
562
- x['version'] = self.cache_version # cache version
563
- try:
564
- np.save(path, x) # save cache for next time
565
- path.with_suffix('.cache.npy').rename(path) # remove .npy suffix
566
- LOGGER.info(f'{prefix}New cache created: {path}')
567
- except Exception as e:
568
- LOGGER.warning(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # not writeable
569
- return x
570
-
571
- def __len__(self):
572
- return len(self.im_files)
573
-
574
- # def __iter__(self):
575
- # self.count = -1
576
- # print('ran dataset iter')
577
- # #self.shuffled_vector = np.random.permutation(self.nF) if self.augment else np.arange(self.nF)
578
- # return self
579
-
580
- def __getitem__(self, index):
581
- index = self.indices[index] # linear, shuffled, or image_weights
582
-
583
- hyp = self.hyp
584
- mosaic = self.mosaic and random.random() < hyp['mosaic']
585
- if mosaic:
586
- # Load mosaic
587
- img, labels = self.load_mosaic(index)
588
- shapes = None
589
-
590
- # MixUp augmentation
591
- if random.random() < hyp['mixup']:
592
- img, labels = mixup(img, labels, *self.load_mosaic(random.randint(0, self.n - 1)))
593
-
594
- else:
595
- # Load image
596
- img, (h0, w0), (h, w) = self.load_image(index)
597
-
598
- # Letterbox
599
- shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
600
- img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
601
- shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
602
-
603
- labels = self.labels[index].copy()
604
- if labels.size: # normalized xywh to pixel xyxy format
605
- labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
606
-
607
- if self.augment:
608
- img, labels = random_perspective(img,
609
- labels,
610
- degrees=hyp['degrees'],
611
- translate=hyp['translate'],
612
- scale=hyp['scale'],
613
- shear=hyp['shear'],
614
- perspective=hyp['perspective'])
615
-
616
- nl = len(labels) # number of labels
617
- if nl:
618
- labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)
619
-
620
- if self.augment:
621
- # Albumentations
622
- img, labels = self.albumentations(img, labels)
623
- nl = len(labels) # update after albumentations
624
-
625
- # HSV color-space
626
- augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
627
-
628
- # Flip up-down
629
- if random.random() < hyp['flipud']:
630
- img = np.flipud(img)
631
- if nl:
632
- labels[:, 2] = 1 - labels[:, 2]
633
-
634
- # Flip left-right
635
- if random.random() < hyp['fliplr']:
636
- img = np.fliplr(img)
637
- if nl:
638
- labels[:, 1] = 1 - labels[:, 1]
639
-
640
- # Cutouts
641
- # labels = cutout(img, labels, p=0.5)
642
- # nl = len(labels) # update after cutout
643
-
644
- labels_out = torch.zeros((nl, 6))
645
- if nl:
646
- labels_out[:, 1:] = torch.from_numpy(labels)
647
-
648
- # Convert
649
- img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
650
- img = np.ascontiguousarray(img)
651
-
652
- return torch.from_numpy(img), labels_out, self.im_files[index], shapes
653
-
654
- def load_image(self, i):
655
- # Loads 1 image from dataset index 'i', returns (im, original hw, resized hw)
656
- im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i],
657
- if im is None: # not cached in RAM
658
- if fn.exists(): # load npy
659
- im = np.load(fn)
660
- else: # read image
661
- im = cv2.imread(f) # BGR
662
- assert im is not None, f'Image Not Found {f}'
663
- h0, w0 = im.shape[:2] # orig hw
664
- r = self.img_size / max(h0, w0) # ratio
665
- if r != 1: # if sizes are not equal
666
- interp = cv2.INTER_LINEAR if (self.augment or r > 1) else cv2.INTER_AREA
667
- im = cv2.resize(im, (int(w0 * r), int(h0 * r)), interpolation=interp)
668
- return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized
669
- else:
670
- return self.ims[i], self.im_hw0[i], self.im_hw[i] # im, hw_original, hw_resized
671
-
672
- def cache_images_to_disk(self, i):
673
- # Saves an image as an *.npy file for faster loading
674
- f = self.npy_files[i]
675
- if not f.exists():
676
- np.save(f.as_posix(), cv2.imread(self.im_files[i]))
677
-
678
- def load_mosaic(self, index):
679
- # YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
680
- labels4, segments4 = [], []
681
- s = self.img_size
682
- yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border) # mosaic center x, y
683
- indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
684
- random.shuffle(indices)
685
- for i, index in enumerate(indices):
686
- # Load image
687
- img, _, (h, w) = self.load_image(index)
688
-
689
- # place img in img4
690
- if i == 0: # top left
691
- img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
692
- x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
693
- x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
694
- elif i == 1: # top right
695
- x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
696
- x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
697
- elif i == 2: # bottom left
698
- x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
699
- x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
700
- elif i == 3: # bottom right
701
- x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
702
- x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
703
-
704
- img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
705
- padw = x1a - x1b
706
- padh = y1a - y1b
707
-
708
- # Labels
709
- labels, segments = self.labels[index].copy(), self.segments[index].copy()
710
- if labels.size:
711
- labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
712
- segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
713
- labels4.append(labels)
714
- segments4.extend(segments)
715
-
716
- # Concat/clip labels
717
- labels4 = np.concatenate(labels4, 0)
718
- for x in (labels4[:, 1:], *segments4):
719
- np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
720
- # img4, labels4 = replicate(img4, labels4) # replicate
721
-
722
- # Augment
723
- img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])
724
- img4, labels4 = random_perspective(img4,
725
- labels4,
726
- segments4,
727
- degrees=self.hyp['degrees'],
728
- translate=self.hyp['translate'],
729
- scale=self.hyp['scale'],
730
- shear=self.hyp['shear'],
731
- perspective=self.hyp['perspective'],
732
- border=self.mosaic_border) # border to remove
733
-
734
- return img4, labels4
735
-
736
- def load_mosaic9(self, index):
737
- # YOLOv5 9-mosaic loader. Loads 1 image + 8 random images into a 9-image mosaic
738
- labels9, segments9 = [], []
739
- s = self.img_size
740
- indices = [index] + random.choices(self.indices, k=8) # 8 additional image indices
741
- random.shuffle(indices)
742
- hp, wp = -1, -1 # height, width previous
743
- for i, index in enumerate(indices):
744
- # Load image
745
- img, _, (h, w) = self.load_image(index)
746
-
747
- # place img in img9
748
- if i == 0: # center
749
- img9 = np.full((s * 3, s * 3, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
750
- h0, w0 = h, w
751
- c = s, s, s + w, s + h # xmin, ymin, xmax, ymax (base) coordinates
752
- elif i == 1: # top
753
- c = s, s - h, s + w, s
754
- elif i == 2: # top right
755
- c = s + wp, s - h, s + wp + w, s
756
- elif i == 3: # right
757
- c = s + w0, s, s + w0 + w, s + h
758
- elif i == 4: # bottom right
759
- c = s + w0, s + hp, s + w0 + w, s + hp + h
760
- elif i == 5: # bottom
761
- c = s + w0 - w, s + h0, s + w0, s + h0 + h
762
- elif i == 6: # bottom left
763
- c = s + w0 - wp - w, s + h0, s + w0 - wp, s + h0 + h
764
- elif i == 7: # left
765
- c = s - w, s + h0 - h, s, s + h0
766
- elif i == 8: # top left
767
- c = s - w, s + h0 - hp - h, s, s + h0 - hp
768
-
769
- padx, pady = c[:2]
770
- x1, y1, x2, y2 = (max(x, 0) for x in c) # allocate coords
771
-
772
- # Labels
773
- labels, segments = self.labels[index].copy(), self.segments[index].copy()
774
- if labels.size:
775
- labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padx, pady) # normalized xywh to pixel xyxy format
776
- segments = [xyn2xy(x, w, h, padx, pady) for x in segments]
777
- labels9.append(labels)
778
- segments9.extend(segments)
779
-
780
- # Image
781
- img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax]
782
- hp, wp = h, w # height, width previous
783
-
784
- # Offset
785
- yc, xc = (int(random.uniform(0, s)) for _ in self.mosaic_border) # mosaic center x, y
786
- img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s]
787
-
788
- # Concat/clip labels
789
- labels9 = np.concatenate(labels9, 0)
790
- labels9[:, [1, 3]] -= xc
791
- labels9[:, [2, 4]] -= yc
792
- c = np.array([xc, yc]) # centers
793
- segments9 = [x - c for x in segments9]
794
-
795
- for x in (labels9[:, 1:], *segments9):
796
- np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
797
- # img9, labels9 = replicate(img9, labels9) # replicate
798
-
799
- # Augment
800
- img9, labels9 = random_perspective(img9,
801
- labels9,
802
- segments9,
803
- degrees=self.hyp['degrees'],
804
- translate=self.hyp['translate'],
805
- scale=self.hyp['scale'],
806
- shear=self.hyp['shear'],
807
- perspective=self.hyp['perspective'],
808
- border=self.mosaic_border) # border to remove
809
-
810
- return img9, labels9
811
-
812
- @staticmethod
813
- def collate_fn(batch):
814
- im, label, path, shapes = zip(*batch) # transposed
815
- for i, lb in enumerate(label):
816
- lb[:, 0] = i # add target image index for build_targets()
817
- return torch.stack(im, 0), torch.cat(label, 0), path, shapes
818
-
819
- @staticmethod
820
- def collate_fn4(batch):
821
- img, label, path, shapes = zip(*batch) # transposed
822
- n = len(shapes) // 4
823
- im4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]
824
-
825
- ho = torch.tensor([[0.0, 0, 0, 1, 0, 0]])
826
- wo = torch.tensor([[0.0, 0, 1, 0, 0, 0]])
827
- s = torch.tensor([[1, 1, 0.5, 0.5, 0.5, 0.5]]) # scale
828
- for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW
829
- i *= 4
830
- if random.random() < 0.5:
831
- im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear',
832
- align_corners=False)[0].type(img[i].type())
833
- lb = label[i]
834
- else:
835
- im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
836
- lb = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
837
- im4.append(im)
838
- label4.append(lb)
839
-
840
- for i, lb in enumerate(label4):
841
- lb[:, 0] = i # add target image index for build_targets()
842
-
843
- return torch.stack(im4, 0), torch.cat(label4, 0), path4, shapes4
844
-
845
-
846
- # Ancillary functions --------------------------------------------------------------------------------------------------
847
- def create_folder(path='./new'):
848
- # Create folder
849
- if os.path.exists(path):
850
- shutil.rmtree(path) # delete output folder
851
- os.makedirs(path) # make new output folder
852
-
853
-
854
- def flatten_recursive(path=DATASETS_DIR / 'coco128'):
855
- # Flatten a recursive directory by bringing all files to top level
856
- new_path = Path(str(path) + '_flat')
857
- create_folder(new_path)
858
- for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)):
859
- shutil.copyfile(file, new_path / Path(file).name)
860
-
861
-
862
- def extract_boxes(path=DATASETS_DIR / 'coco128'): # from utils.dataloaders import *; extract_boxes()
863
- # Convert detection dataset into classification dataset, with one directory per class
864
- path = Path(path) # images dir
865
- shutil.rmtree(path / 'classifier') if (path / 'classifier').is_dir() else None # remove existing
866
- files = list(path.rglob('*.*'))
867
- n = len(files) # number of files
868
- for im_file in tqdm(files, total=n):
869
- if im_file.suffix[1:] in IMG_FORMATS:
870
- # image
871
- im = cv2.imread(str(im_file))[..., ::-1] # BGR to RGB
872
- h, w = im.shape[:2]
873
-
874
- # labels
875
- lb_file = Path(img2label_paths([str(im_file)])[0])
876
- if Path(lb_file).exists():
877
- with open(lb_file) as f:
878
- lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32) # labels
879
-
880
- for j, x in enumerate(lb):
881
- c = int(x[0]) # class
882
- f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename
883
- if not f.parent.is_dir():
884
- f.parent.mkdir(parents=True)
885
-
886
- b = x[1:] * [w, h, w, h] # box
887
- # b[2:] = b[2:].max() # rectangle to square
888
- b[2:] = b[2:] * 1.2 + 3 # pad
889
- b = xywh2xyxy(b.reshape(-1, 4)).ravel().astype(np.int)
890
-
891
- b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image
892
- b[[1, 3]] = np.clip(b[[1, 3]], 0, h)
893
- assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}'
894
-
895
-
896
- def autosplit(path=DATASETS_DIR / 'coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False):
897
- """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
898
- Usage: from utils.dataloaders import *; autosplit()
899
- Arguments
900
- path: Path to images directory
901
- weights: Train, val, test weights (list, tuple)
902
- annotated_only: Only use images with an annotated txt file
903
- """
904
- path = Path(path) # images dir
905
- files = sorted(x for x in path.rglob('*.*') if x.suffix[1:].lower() in IMG_FORMATS) # image files only
906
- n = len(files) # number of files
907
- random.seed(0) # for reproducibility
908
- indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
909
-
910
- txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files
911
- [(path.parent / x).unlink(missing_ok=True) for x in txt] # remove existing
912
-
913
- print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only)
914
- for i, img in tqdm(zip(indices, files), total=n):
915
- if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label
916
- with open(path.parent / txt[i], 'a') as f:
917
- f.write('./' + img.relative_to(path.parent).as_posix() + '\n') # add image to txt file
918
-
919
-
920
- def verify_image_label(args):
921
- # Verify one image-label pair
922
- im_file, lb_file, prefix = args
923
- nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments
924
- try:
925
- # verify images
926
- im = Image.open(im_file)
927
- im.verify() # PIL verify
928
- shape = exif_size(im) # image size
929
- assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels'
930
- assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}'
931
- if im.format.lower() in ('jpg', 'jpeg'):
932
- with open(im_file, 'rb') as f:
933
- f.seek(-2, 2)
934
- if f.read() != b'\xff\xd9': # corrupt JPEG
935
- ImageOps.exif_transpose(Image.open(im_file)).save(im_file, 'JPEG', subsampling=0, quality=100)
936
- msg = f'{prefix}WARNING: {im_file}: corrupt JPEG restored and saved'
937
-
938
- # verify labels
939
- if os.path.isfile(lb_file):
940
- nf = 1 # label found
941
- with open(lb_file) as f:
942
- lb = [x.split() for x in f.read().strip().splitlines() if len(x)]
943
- if any(len(x) > 6 for x in lb): # is segment
944
- classes = np.array([x[0] for x in lb], dtype=np.float32)
945
- segments = [np.array(x[1:], dtype=np.float32).reshape(-1, 2) for x in lb] # (cls, xy1...)
946
- lb = np.concatenate((classes.reshape(-1, 1), segments2boxes(segments)), 1) # (cls, xywh)
947
- lb = np.array(lb, dtype=np.float32)
948
- nl = len(lb)
949
- if nl:
950
- assert lb.shape[1] == 5, f'labels require 5 columns, {lb.shape[1]} columns detected'
951
- assert (lb >= 0).all(), f'negative label values {lb[lb < 0]}'
952
- assert (lb[:, 1:] <= 1).all(), f'non-normalized or out of bounds coordinates {lb[:, 1:][lb[:, 1:] > 1]}'
953
- _, i = np.unique(lb, axis=0, return_index=True)
954
- if len(i) < nl: # duplicate row check
955
- lb = lb[i] # remove duplicates
956
- if segments:
957
- segments = segments[i]
958
- msg = f'{prefix}WARNING: {im_file}: {nl - len(i)} duplicate labels removed'
959
- else:
960
- ne = 1 # label empty
961
- lb = np.zeros((0, 5), dtype=np.float32)
962
- else:
963
- nm = 1 # label missing
964
- lb = np.zeros((0, 5), dtype=np.float32)
965
- return im_file, lb, shape, segments, nm, nf, ne, nc, msg
966
- except Exception as e:
967
- nc = 1
968
- msg = f'{prefix}WARNING: {im_file}: ignoring corrupt image/label: {e}'
969
- return [None, None, None, None, nm, nf, ne, nc, msg]
970
-
971
-
972
- def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profile=False, hub=False):
973
- """ Return dataset statistics dictionary with images and instances counts per split per class
974
- To run in parent directory: export PYTHONPATH="$PWD/yolov5"
975
- Usage1: from utils.dataloaders import *; dataset_stats('coco128.yaml', autodownload=True)
976
- Usage2: from utils.dataloaders import *; dataset_stats('path/to/coco128_with_yaml.zip')
977
- Arguments
978
- path: Path to data.yaml or data.zip (with data.yaml inside data.zip)
979
- autodownload: Attempt to download dataset if not found locally
980
- verbose: Print stats dictionary
981
- """
982
-
983
- def _round_labels(labels):
984
- # Update labels to integer class and 6 decimal place floats
985
- return [[int(c), *(round(x, 4) for x in points)] for c, *points in labels]
986
-
987
- def _find_yaml(dir):
988
- # Return data.yaml file
989
- files = list(dir.glob('*.yaml')) or list(dir.rglob('*.yaml')) # try root level first and then recursive
990
- assert files, f'No *.yaml file found in {dir}'
991
- if len(files) > 1:
992
- files = [f for f in files if f.stem == dir.stem] # prefer *.yaml files that match dir name
993
- assert files, f'Multiple *.yaml files found in {dir}, only 1 *.yaml file allowed'
994
- assert len(files) == 1, f'Multiple *.yaml files found: {files}, only 1 *.yaml file allowed in {dir}'
995
- return files[0]
996
-
997
- def _unzip(path):
998
- # Unzip data.zip
999
- if str(path).endswith('.zip'): # path is data.zip
1000
- assert Path(path).is_file(), f'Error unzipping {path}, file not found'
1001
- ZipFile(path).extractall(path=path.parent) # unzip
1002
- dir = path.with_suffix('') # dataset directory == zip name
1003
- assert dir.is_dir(), f'Error unzipping {path}, {dir} not found. path/to/abc.zip MUST unzip to path/to/abc/'
1004
- return True, str(dir), _find_yaml(dir) # zipped, data_dir, yaml_path
1005
- else: # path is data.yaml
1006
- return False, None, path
1007
-
1008
- def _hub_ops(f, max_dim=1920):
1009
- # HUB ops for 1 image 'f': resize and save at reduced quality in /dataset-hub for web/app viewing
1010
- f_new = im_dir / Path(f).name # dataset-hub image filename
1011
- try: # use PIL
1012
- im = Image.open(f)
1013
- r = max_dim / max(im.height, im.width) # ratio
1014
- if r < 1.0: # image too large
1015
- im = im.resize((int(im.width * r), int(im.height * r)))
1016
- im.save(f_new, 'JPEG', quality=75, optimize=True) # save
1017
- except Exception as e: # use OpenCV
1018
- print(f'WARNING: HUB ops PIL failure {f}: {e}')
1019
- im = cv2.imread(f)
1020
- im_height, im_width = im.shape[:2]
1021
- r = max_dim / max(im_height, im_width) # ratio
1022
- if r < 1.0: # image too large
1023
- im = cv2.resize(im, (int(im_width * r), int(im_height * r)), interpolation=cv2.INTER_AREA)
1024
- cv2.imwrite(str(f_new), im)
1025
-
1026
- zipped, data_dir, yaml_path = _unzip(Path(path))
1027
- with open(check_yaml(yaml_path), errors='ignore') as f:
1028
- data = yaml.safe_load(f) # data dict
1029
- if zipped:
1030
- data['path'] = data_dir # TODO: should this be dir.resolve()?
1031
- check_dataset(data, autodownload) # download dataset if missing
1032
- hub_dir = Path(data['path'] + ('-hub' if hub else ''))
1033
- stats = {'nc': data['nc'], 'names': data['names']} # statistics dictionary
1034
- for split in 'train', 'val', 'test':
1035
- if data.get(split) is None:
1036
- stats[split] = None # i.e. no test set
1037
- continue
1038
- x = []
1039
- dataset = LoadImagesAndLabels(data[split]) # load dataset
1040
- for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'):
1041
- x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc']))
1042
- x = np.array(x) # shape(128x80)
1043
- stats[split] = {
1044
- 'instance_stats': {
1045
- 'total': int(x.sum()),
1046
- 'per_class': x.sum(0).tolist()},
1047
- 'image_stats': {
1048
- 'total': dataset.n,
1049
- 'unlabelled': int(np.all(x == 0, 1).sum()),
1050
- 'per_class': (x > 0).sum(0).tolist()},
1051
- 'labels': [{
1052
- str(Path(k).name): _round_labels(v.tolist())} for k, v in zip(dataset.im_files, dataset.labels)]}
1053
-
1054
- if hub:
1055
- im_dir = hub_dir / 'images'
1056
- im_dir.mkdir(parents=True, exist_ok=True)
1057
- for _ in tqdm(ThreadPool(NUM_THREADS).imap(_hub_ops, dataset.im_files), total=dataset.n, desc='HUB Ops'):
1058
- pass
1059
-
1060
- # Profile
1061
- stats_path = hub_dir / 'stats.json'
1062
- if profile:
1063
- for _ in range(1):
1064
- file = stats_path.with_suffix('.npy')
1065
- t1 = time.time()
1066
- np.save(file, stats)
1067
- t2 = time.time()
1068
- x = np.load(file, allow_pickle=True)
1069
- print(f'stats.npy times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
1070
-
1071
- file = stats_path.with_suffix('.json')
1072
- t1 = time.time()
1073
- with open(file, 'w') as f:
1074
- json.dump(stats, f) # save stats *.json
1075
- t2 = time.time()
1076
- with open(file) as f:
1077
- x = json.load(f) # load hyps dict
1078
- print(f'stats.json times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write')
1079
-
1080
- # Save, print and return
1081
- if hub:
1082
- print(f'Saving {stats_path.resolve()}...')
1083
- with open(stats_path, 'w') as f:
1084
- json.dump(stats, f) # save stats.json
1085
- if verbose:
1086
- print(json.dumps(stats, indent=2, sort_keys=False))
1087
- return stats