YOLO Module
Yolo
Class for YOLO PyTorch model operations.
This class provides methods to load a YOLO model, get class names, export to various formats,
and run inference. Also, this class contains constants related to YOLO models, directories, colors, and utility functions
for checking model names, versions, dataset statuses, and dataset names.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
|
export_onnx(model)
staticmethod
Export the model to ONNX format.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model
|
YOLO
|
Loaded YOLO model. |
required |
Returns:
str: Path to the exported ONNX model file.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
62 63 64 65 66 67 68 69 70 71 72 |
|
export_tensor_rt(model, quantized=True)
staticmethod
Export the model to TensorRT format.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model
|
YOLO
|
Loaded YOLO model. |
required |
quantized
|
bool
|
Whether to quantize the model, default is True. |
True
|
Returns:
str: Path to the exported TensorRT engine file.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
49 50 51 52 53 54 55 56 57 58 59 60 |
|
export_tflite(model, quantized=True)
staticmethod
Export the model to TFLite format.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model
|
YOLO
|
Loaded YOLO model. |
required |
quantized
|
bool
|
Whether to quantize the model, default is True. |
True
|
Returns:
str: Path to the exported TFLite model file.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
74 75 76 77 78 79 80 81 82 83 84 85 |
|
get_class_names(model)
staticmethod
Get YOLO PyTorch model class names.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model
|
YOLO
|
Loaded YOLO model. |
required |
Returns:
Dict[int, str]: Dictionary mapping class indices to class names.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
37 38 39 40 41 42 43 44 45 46 47 |
|
load(model_path, task='detect')
staticmethod
Load YOLO PyTorch model.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_path
|
str | PathLike[str]
|
Path to the YOLO model file. |
required |
task
|
str
|
Task type, default is 'detect'. |
'detect'
|
Returns:
YOLO: Loaded YOLO model.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
|
run_inference(model, preprocessed_image)
staticmethod
Run inference from PyTorch model.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model
|
YOLO
|
Loaded YOLO model. |
required |
preprocessed_image
|
Tensor
|
Preprocessed image tensor. |
required |
Returns:
tuple(List, float): Inferences and elapsed time in seconds.
Source code in devices\raspberry_pi_5\src\yolo\__init__.py
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
|
args
Args
Bases: Args
Class to handle command line arguments.
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 |
|
add_yolo_classes_argument(required=True)
Add YOLO classes argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
required
|
bool
|
Whether the argument is required or not. Defaults to True. |
True
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
|
add_yolo_device_argument(required=True)
Add YOLO device argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
required
|
bool
|
Whether the argument is required or not. Defaults to True. |
True
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 |
|
add_yolo_epochs_argument(required=True)
Add YOLO epochs argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
required
|
bool
|
Whether the argument is required or not. Defaults to True. |
True
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
|
add_yolo_format_argument(required=False)
Add YOLO format argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
required
|
bool
|
Whether the argument is required or not. Defaults to False. |
False
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
|
add_yolo_ignore_classes_argument(required=True)
Add YOLO ignore classes argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
required
|
bool
|
Whether the argument is required or not. Defaults to True. |
True
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
|
add_yolo_image_size_argument(required=True)
Add YOLO image size argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
required
|
bool
|
Whether the argument is required or not. Defaults to True. |
True
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 |
|
add_yolo_input_model_pt_argument()
Add YOLO input PyTorch model argument to the parser.
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
14 15 16 17 18 19 20 21 22 23 |
|
add_yolo_output_model_argument()
Add YOLO output model argument to the parser.
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
34 35 36 37 38 39 40 41 42 43 |
|
add_yolo_quantized_argument(default=False)
Add YOLO quantized argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
default
|
bool
|
Default value for the quantized argument. Defaults to False. |
False
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
82 83 84 85 86 87 88 89 90 91 92 |
|
add_yolo_retraining_argument(default=False)
Add YOLO retraining argument to the parser.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
default
|
bool
|
Default value for the retraining argument. Defaults to False. |
False
|
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
103 104 105 106 107 108 109 110 111 112 113 |
|
get_yolo_classes()
Get the YOLO classes argument.
Returns:
Type | Description |
---|---|
List[str]
|
List[str]: The YOLO classes. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
142 143 144 145 146 147 148 149 |
|
get_yolo_device()
Get the YOLO device argument.
Returns:
Name | Type | Description |
---|---|---|
str |
str
|
The YOLO device. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
224 225 226 227 228 229 230 231 |
|
get_yolo_epochs()
Get the YOLO epochs argument.
Returns:
Name | Type | Description |
---|---|---|
int |
int
|
The number of YOLO epochs. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
196 197 198 199 200 201 202 203 |
|
get_yolo_format()
Get the YOLO format argument.
Returns:
Name | Type | Description |
---|---|---|
str |
str
|
The YOLO format. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
73 74 75 76 77 78 79 80 |
|
get_yolo_ignore_classes()
Get the YOLO ignore classes argument.
Returns:
Type | Description |
---|---|
List[str]
|
List[str]: The YOLO ignore classes. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
169 170 171 172 173 174 175 176 |
|
get_yolo_image_size()
Get the YOLO image size argument.
Returns:
Name | Type | Description |
---|---|---|
int |
int
|
The YOLO image size. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
251 252 253 254 255 256 257 258 |
|
get_yolo_input_model_pt()
Get the YOLO input PyTorch model argument.
Returns:
Name | Type | Description |
---|---|---|
str |
str
|
The YOLO input PyTorch model. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
25 26 27 28 29 30 31 32 |
|
get_yolo_output_model()
Get the YOLO output model argument.
Returns:
Name | Type | Description |
---|---|---|
str |
str
|
The YOLO output model. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
45 46 47 48 49 50 51 52 |
|
get_yolo_quantized()
Get the YOLO quantized argument.
Returns:
Name | Type | Description |
---|---|---|
bool |
bool
|
The YOLO quantized flag. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
94 95 96 97 98 99 100 101 |
|
get_yolo_retraining()
Get the YOLO retraining argument.
Returns:
Name | Type | Description |
---|---|---|
bool |
bool
|
The YOLO retraining flag. |
Source code in devices\raspberry_pi_5\src\yolo\args\__init__.py
115 116 117 118 119 120 121 122 |
|
enums
Flag
Bases: Enum
Enum to represent command line flags.
Source code in devices\raspberry_pi_5\src\yolo\args\enums.py
6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
|
parsed_name
property
Get the flag name with the prefix.
Returns:
Name | Type | Description |
---|---|---|
str |
str
|
The flag name with the prefix. |
files
Files
Bases: Files
Files utility class.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 |
|
check_dataset_name(dataset_name)
classmethod
Check the validity of dataset name.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_name
|
str
|
The name of the dataset to check. |
required |
Raises:
ValueError: If the dataset name is invalid.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
|
check_dataset_status(dataset_status)
staticmethod
Check the validity of dataset status.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_status
|
str | None
|
The dataset status to check. |
required |
Raises:
ValueError: If the dataset status is invalid.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
|
check_model_dataset_name(dataset_name, model_name)
classmethod
Check the validity of the model dataset name.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_name
|
str
|
The name of the dataset. |
required |
model_name
|
str | None
|
The name of the model. |
required |
Raises:
ValueError: If the dataset name is invalid for the given model name.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
|
check_model_dataset_status(dataset_name, dataset_status)
staticmethod
Check the validity of model dataset status.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_name
|
str
|
The name of the dataset. |
required |
dataset_status
|
str | None
|
The dataset status to check. |
required |
Raises:
ValueError: If the dataset name is invalid for the given dataset status.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
|
get_dataset_model_dir_path(dataset_name, dataset_status, model_name)
classmethod
Get the dataset model directory path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_name
|
str
|
The name of the dataset. |
required |
dataset_status
|
str | None
|
The status of the dataset (e.g., 'to_process', 'processed'). |
required |
model_name
|
str | None
|
The name of the YOLO model. |
required |
Returns:
str | PathLike[str]: The path to the dataset model directory.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
|
get_model_best_onnx_path(model_name, yolo_version)
classmethod
Get the model best ONNX path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_name
|
str
|
Name of the YOLO model. |
required |
yolo_version
|
str
|
Version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the model best ONNX file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 |
|
get_model_best_pt_path(model_name, yolo_version)
classmethod
Get the model best PyTorch path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_name
|
str
|
Name of the YOLO model. |
required |
yolo_version
|
str
|
Version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the model best PyTorch file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 |
|
get_model_colab_data_path(model_name)
classmethod
Get the model Colab data path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_name
|
str
|
Name of the YOLO model. |
required |
Returns:
str | PathLike: The path to the model Colab data file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 |
|
get_model_data_name(model_name)
classmethod
Get the model data name.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_name
|
str
|
Name of the YOLO model. |
required |
Returns:
str: The name of the model data file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
305 306 307 308 309 310 311 312 313 314 315 316 317 318 |
|
get_model_local_data_path(model_name)
classmethod
Get the model local data path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_name
|
str
|
Name of the YOLO model. |
required |
Returns:
str | PathLike: The path to the model local data file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 |
|
get_tf_record_path(model_name, yolo_version)
classmethod
Get the TensorFlow Record path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model_name
|
str
|
Name of the YOLO model. |
required |
yolo_version
|
str
|
Version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the TensorFlow Record file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 |
|
get_yolo_colab_data_dir_path()
classmethod
Get the YOLO colab data folder path.
Returns:
Type | Description |
---|---|
str | PathLike
|
str | PathLike: The path to the YOLO colab data folder. |
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
279 280 281 282 283 284 285 286 287 288 289 290 |
|
get_yolo_data_dir_path()
classmethod
Get the YOLO data folder path.
Returns:
Type | Description |
---|---|
str | PathLike
|
str | PathLike: The path to the YOLO data folder. |
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
269 270 271 272 273 274 275 276 277 |
|
get_yolo_dataset_classes_file_path(dataset_name, dataset_status, model_name)
classmethod
Get the YOLO dataset classes file path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_name
|
str
|
The name of the dataset. |
required |
dataset_status
|
str | None
|
The status of the dataset (e.g., 'to_process', 'processed'). |
required |
model_name
|
str | None
|
The name of the YOLO model. |
required |
Returns:
str | PathLike: The path to the YOLO dataset classes file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 |
|
get_yolo_dataset_notes_file_path(dataset_name, dataset_status, model_name)
classmethod
Get the YOLO dataset notes file path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
dataset_name
|
str
|
The name of the dataset. |
required |
dataset_status
|
str | None
|
The status of the dataset (e.g., 'to_process', 'processed'). |
required |
model_name
|
str | None
|
The name of the YOLO model. |
required |
Returns:
str | PathLike: The path to the YOLO dataset notes file.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 |
|
get_yolo_local_data_dir_path()
classmethod
Get the YOLO local data folder path.
Returns:
Type | Description |
---|---|
str | PathLike
|
str | PathLike: The path to the YOLO local data folder. |
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
292 293 294 295 296 297 298 299 300 301 302 303 |
|
get_yolo_notebooks_dir_path(yolo_version)
classmethod
Get the YOLO notebooks folder path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
yolo_version
|
str
|
The version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the YOLO notebooks folder.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
356 357 358 359 360 361 362 363 364 365 366 367 368 369 |
|
get_yolo_old_runs_dir_path(yolo_version)
classmethod
Get the YOLO old runs folder path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
yolo_version
|
str
|
The version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the YOLO old runs folder.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
195 196 197 198 199 200 201 202 203 204 205 206 207 208 |
|
get_yolo_runs_new_name_dir_path(yolo_version)
classmethod
Get the YOLO runs folder path with the new name.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
yolo_version
|
str
|
The version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the YOLO runs folder with a new name.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
|
get_yolo_zip_dir_path(yolo_version)
classmethod
Get the YOLO zip folder path.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
yolo_version
|
str
|
The version of the YOLO model. |
required |
Returns:
str | PathLike: The path to the YOLO zip folder.
Source code in devices\raspberry_pi_5\src\yolo\files\__init__.py
254 255 256 257 258 259 260 261 262 263 264 265 266 267 |
|