File size: 3,374 Bytes
f6228f9 |
1 2 3 4 5 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 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 |
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""Monkey patches to update/extend functionality of existing functions."""
import time
from pathlib import Path
import cv2
import numpy as np
import torch
# OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------
_imshow = cv2.imshow # copy to avoid recursion errors
def imread(filename: str, flags: int = cv2.IMREAD_COLOR):
"""
Read an image from a file.
Args:
filename (str): Path to the file to read.
flags (int, optional): Flag that can take values of cv2.IMREAD_*. Defaults to cv2.IMREAD_COLOR.
Returns:
(np.ndarray): The read image.
"""
return cv2.imdecode(np.fromfile(filename, np.uint8), flags)
def imwrite(filename: str, img: np.ndarray, params=None):
"""
Write an image to a file.
Args:
filename (str): Path to the file to write.
img (np.ndarray): Image to write.
params (list of ints, optional): Additional parameters. See OpenCV documentation.
Returns:
(bool): True if the file was written, False otherwise.
"""
try:
cv2.imencode(Path(filename).suffix, img, params)[1].tofile(filename)
return True
except Exception:
return False
def imshow(winname: str, mat: np.ndarray):
"""
Displays an image in the specified window.
Args:
winname (str): Name of the window.
mat (np.ndarray): Image to be shown.
"""
_imshow(winname.encode("unicode_escape").decode(), mat)
# PyTorch functions ----------------------------------------------------------------------------------------------------
_torch_load = torch.load # copy to avoid recursion errors
_torch_save = torch.save
def torch_load(*args, **kwargs):
"""
Load a PyTorch model with updated arguments to avoid warnings.
This function wraps torch.load and adds the 'weights_only' argument for PyTorch 1.13.0+ to prevent warnings.
Args:
*args (Any): Variable length argument list to pass to torch.load.
**kwargs (Any): Arbitrary keyword arguments to pass to torch.load.
Returns:
(Any): The loaded PyTorch object.
Note:
For PyTorch versions 2.0 and above, this function automatically sets 'weights_only=False'
if the argument is not provided, to avoid deprecation warnings.
"""
from ultralytics.utils.torch_utils import TORCH_1_13
if TORCH_1_13 and "weights_only" not in kwargs:
kwargs["weights_only"] = False
return _torch_load(*args, **kwargs)
def torch_save(*args, **kwargs):
"""
Optionally use dill to serialize lambda functions where pickle does not, adding robustness with 3 retries and
exponential standoff in case of save failure.
Args:
*args (tuple): Positional arguments to pass to torch.save.
**kwargs (Any): Keyword arguments to pass to torch.save.
"""
for i in range(4): # 3 retries
try:
return _torch_save(*args, **kwargs)
except RuntimeError as e: # unable to save, possibly waiting for device to flush or antivirus scan
if i == 3:
raise e
time.sleep((2**i) / 2) # exponential standoff: 0.5s, 1.0s, 2.0s
|