mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-23 21:44:22 +08:00
fe27db2f6e
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Burhan <62214284+Burhan-Q@users.noreply.github.com> Co-authored-by: Kayzwer <68285002+Kayzwer@users.noreply.github.com>
78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
"""Monkey patches to update/extend functionality of existing functions."""
|
|
|
|
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_save = torch.save # copy to avoid recursion errors
|
|
|
|
|
|
def torch_save(*args, **kwargs):
|
|
"""
|
|
Use dill (if exists) to serialize the lambda functions where pickle does not do this.
|
|
|
|
Args:
|
|
*args (tuple): Positional arguments to pass to torch.save.
|
|
**kwargs (dict): Keyword arguments to pass to torch.save.
|
|
"""
|
|
try:
|
|
import dill as pickle # noqa
|
|
except ImportError:
|
|
import pickle
|
|
|
|
if "pickle_module" not in kwargs:
|
|
kwargs["pickle_module"] = pickle # noqa
|
|
return _torch_save(*args, **kwargs)
|