diff --git a/src/utils/benchmark.py b/src/utils/benchmark.py new file mode 100644 index 00000000..83099a2f --- /dev/null +++ b/src/utils/benchmark.py @@ -0,0 +1,26 @@ +import time +import logging + +logging.basicConfig(format='%(message)s', level=logging.INFO) + +class Benchmark: + time_units = { + 'seconds': 1e9, + 'milliseconds': 1e6, + 'microseconds': 1e3, + 'nanoseconds': 1, + } + + def __new__(cls, func): + def wrapper(*args, **kwargs): + start_time = time.perf_counter_ns() + result = func(*args, **kwargs) + end_time = time.perf_counter_ns() + + elapsed_time = (end_time - start_time) + for unit, factor in cls.time_units.items(): + if elapsed_time >= factor: + logging.info(f"Function {func.__name__} - {elapsed_time / factor:.2f} {unit}") + break + return result + return wrapper diff --git a/src/utils/image.py b/src/utils/image.py index ce7d204d..8b0966b6 100644 --- a/src/utils/image.py +++ b/src/utils/image.py @@ -5,7 +5,7 @@ from typing import Union from src.shared.typings import BBox, GrayImage from src.utils.core import hashit, locate - +from src.utils.benchmark import Benchmark # TODO: add types # TODO: add unit tests @@ -41,17 +41,12 @@ def inner(screenshot: GrayImage) -> Union[BBox, None]: return inner return decorator - # TODO: add unit tests -@njit(cache=True, fastmath=True) + def convertGraysToBlack(arr: np.ndarray) -> np.ndarray: - for i in range(len(arr)): - for j in range(len(arr[0])): - if arr[i, j] >= 50 and arr[i, j] <= 100: - arr[i, j] = 0 + arr[(arr >= 50) & (arr <= 100)] = 0 return arr - # TODO: add unit tests def RGBtoGray(image: np.ndarray) -> GrayImage: return cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)