diff --git a/MANIFEST.in b/MANIFEST.in index 97a3faf..c720c17 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,2 @@ -include cvu/detector/yolov5/backends/weights/*.json +include cvu/detector/weights.json include cvu/utils/backend/*.json diff --git a/cvu/detector/interface/__init__.py b/cvu/detector/interface/__init__.py new file mode 100644 index 0000000..9782df8 --- /dev/null +++ b/cvu/detector/interface/__init__.py @@ -0,0 +1,2 @@ +from .core import IDetector +from .model import IDetectorModel diff --git a/cvu/detector/interface/core.py b/cvu/detector/interface/core.py new file mode 100644 index 0000000..0b7b373 --- /dev/null +++ b/cvu/detector/interface/core.py @@ -0,0 +1,61 @@ +"""Includes interface for CVU Detectors' core. +""" +import abc +from typing import List, Union + +import numpy as np + +from cvu.interface.core import ICore +from cvu.detector.predictions import Predictions + + +class IDetector(ICore, metaclass=abc.ABCMeta): + """Interface which will be implemented for every CVU Detector. + A core defines one complete method/solution for certain use cases. + For example, YoloV5 is a detector core of Object Detection use cases. + """ + + @abc.abstractmethod + def __init__(self, classes: Union[str, List[str]], backend: str, + weight: str, device: str, *args, **kwargs) -> None: + """Initiate Core. + + Args: + classes (Union[str, List[str]]): name of classes to be detected. + It can be set to individual classes like 'coco', 'person', 'cat' etc. + Alternatively, it can be a list of classes such as ['person', 'cat']. + For default weights, 'classes' is used to filter out objects + according to provided argument from coco class (unless specified otherwise). + + backend (str): name of the backend to be used for inference purposes. + + weight (str): path to weight files (according to selected backend). + + device (str): name of the device to be used. Valid + devices can be "cpu", "gpu", "tpu", "auto". + + auto_install (bool): auto install missing requirements for the selected + backend. + """ + ... + + @abc.abstractmethod + def __call__(self, inputs: np.ndarray, **kwargs) -> Predictions: + """Run object detection on given image + + Args: + inputs (np.ndarray): image in BGR format. + + Returns: + Predictions: detected objects. + """ + ... + + @abc.abstractmethod + def __repr__(self) -> str: + """Returns Backend and Model Information + + Returns: + str: formatted string with method and config info. + """ + ... diff --git a/cvu/detector/interface/model.py b/cvu/detector/interface/model.py new file mode 100644 index 0000000..f116e2a --- /dev/null +++ b/cvu/detector/interface/model.py @@ -0,0 +1,39 @@ +"""This module contains interface definition of the detector model +""" +import abc + +import numpy as np + + +class IDetectorModel(metaclass=abc.ABCMeta): + """Interface of the detector model + + A model performs inference, using a certain backend runtime, + on a numpy array, and returns result after performing NMS. + + Inputs are expected to be normalized in channels-first order + with/without batch axis. + """ + + @abc.abstractmethod + def __call__(self, inputs: np.ndarray) -> np.ndarray: + """Performs model inference on given inputs, and returns + inference's output after NMS. + + Args: + inputs (np.ndarray): normalized in channels-first format, + with batch axis. + + Returns: + np.ndarray: inference's output after NMS + """ + ... + + @abc.abstractmethod + def __repr__(self) -> str: + """Represents model with method and configuration informations. + + Returns: + str: formatted string with method and config info. + """ + ... diff --git a/cvu/detector/weights.json b/cvu/detector/weights.json new file mode 100644 index 0000000..e4002c4 --- /dev/null +++ b/cvu/detector/weights.json @@ -0,0 +1,12 @@ +{ + "yolov5":{ + "yolov5s":{ + "onnx": "1piC3ZGuc4D8MMJQQRK3dgaCa66-4Ucxi", + "tensorflow":"1SA7hT4jUx9szqJePZ8NssMPSCelRvYpo", + "tensorrt":"1piC3ZGuc4D8MMJQQRK3dgaCa66-4Ucxi", + "tflite": "1oeKZm81Gz5OejmcgDA0biemkPWCY6qNu", + "torch":"1cWeCusb2IB-x-h_IZs7t3Y_rs9GkzQsI", + "torch.cuda":"16BFVGsEYAXsupFoCXxlOhSBovoMH6juK" + } + } +} diff --git a/cvu/detector/yolo/__init__.py b/cvu/detector/yolo/__init__.py new file mode 100644 index 0000000..931816d --- /dev/null +++ b/cvu/detector/yolo/__init__.py @@ -0,0 +1,4 @@ +"""This module contains core Yolo Implementation with various +functional backends. +""" +from .core import Yolo diff --git a/cvu/detector/yolov5/backends/__init__.py b/cvu/detector/yolo/backends/__init__.py similarity index 84% rename from cvu/detector/yolov5/backends/__init__.py rename to cvu/detector/yolo/backends/__init__.py index c0b786d..c61bb87 100644 --- a/cvu/detector/yolov5/backends/__init__.py +++ b/cvu/detector/yolo/backends/__init__.py @@ -1,4 +1,4 @@ -"""This module contains implementation of Yolov5 model +"""This module contains implementation of Yolo model for various backends. A model (aka backend) basically performs inference on a given input numpy array, diff --git a/cvu/detector/yolov5/backends/yolov5_onnx.py b/cvu/detector/yolo/backends/yolo_onnx.py similarity index 78% rename from cvu/detector/yolov5/backends/yolov5_onnx.py rename to cvu/detector/yolo/backends/yolo_onnx.py index 449a88d..a32e525 100644 --- a/cvu/detector/yolov5/backends/yolov5_onnx.py +++ b/cvu/detector/yolo/backends/yolo_onnx.py @@ -1,4 +1,4 @@ -"""This file contains Yolov5's IModel implementation in ONNX. +"""This file contains Yolo's IDetectorModel implementation in ONNX. This model (onnx-backend) performs inference using ONNXRUNTIME, on a given input numpy array, and returns result after performing nms and other backend specific postprocessings. @@ -13,14 +13,12 @@ import numpy as np import onnxruntime -from cvu.interface.model import IModel -from cvu.utils.general import get_path -from cvu.detector.yolov5.backends.common import download_weights +from cvu.detector.interface import IDetectorModel from cvu.postprocess.nms.yolov5 import non_max_suppression_np -class Yolov5(IModel): - """Implements IModel for Yolov5 using ONNX. +class Yolo(IDetectorModel): + """Implements IDetectorModel for Yolo using ONNX. This model (onnx-backend) performs inference, using ONNXRUNTIME, on a numpy array, and returns result after performing NMS. @@ -29,13 +27,11 @@ class Yolov5(IModel): with/without batch axis. """ - def __init__(self, weight: str = "yolov5s", device: str = 'auto') -> None: + def __init__(self, weight: str, device: str = 'auto') -> None: """Initiate Model Args: - weight (str, optional): path to onnx weight file. Alternatively, - it also accepts identifiers (such as yolvo5s, yolov5m, etc.) to load - pretrained models. Defaults to "yolov5s". + weight (str): path to onnx weight file. device (str, optional): name of the device to be used. Valid devices can be "cpu", "gpu", "auto". Defaults to "auto" which tries to use the device @@ -50,26 +46,18 @@ def _load_model(self, weight: str) -> None: """Internally loads ONNX Args: - weight (str): path to ONNX weight file or predefined-identifiers - (such as yolvo5s, yolov5m, etc.) + weight (str): path to ONNX weight file """ + if not os.path.exists(weight): + raise FileNotFoundError(f"Unable to locate model weights {weight}") + if weight.endswith("torchscript"): # convert to onnx convert_to_onnx = importlib.import_module( - ".convert_to_onnx","cvu.utils.backend_onnx") + ".convert_to_onnx", "cvu.utils.backend_onnx") weight = convert_to_onnx.onnx_from_torchscript( weight, save_path=weight.replace("torchscript", "onnx")) - # attempt to load predefined weights - elif not os.path.exists(weight): - - # get path to pretrained weights - weight += '.onnx' - weight = get_path(__file__, "weights", weight) - - # download weights if not already downloaded - download_weights(weight, "onnx") - # load model if self._device == "cpu": # load model on cpu @@ -132,7 +120,7 @@ def __repr__(self) -> str: Returns: str: information string """ - return f"Yolov5s ONNX-{self._device}" + return f"Yolo ONNX-{self._device}" @staticmethod def _postprocess(outputs: np.ndarray) -> List[np.ndarray]: diff --git a/cvu/detector/yolov5/backends/yolov5_tensorflow.py b/cvu/detector/yolo/backends/yolo_tensorflow.py similarity index 80% rename from cvu/detector/yolov5/backends/yolov5_tensorflow.py rename to cvu/detector/yolo/backends/yolo_tensorflow.py index e2b502d..d947c54 100644 --- a/cvu/detector/yolov5/backends/yolov5_tensorflow.py +++ b/cvu/detector/yolo/backends/yolo_tensorflow.py @@ -1,4 +1,4 @@ -"""This file contains Yolov5's IModel implementation in Tensorflow. +"""This file contains Yolo's IDetectorModel implementation in Tensorflow. This model (tensorflow-backend) performs inference using SavedModel, on a given input numpy array, and returns result after performing nms and other backend specific postprocessings. @@ -14,15 +14,13 @@ import tensorflow as tf from tensorflow.keras import mixed_precision -from cvu.interface.model import IModel -from cvu.utils.general import get_path -from cvu.detector.yolov5.backends.common import download_weights +from cvu.detector.interface import IDetectorModel from cvu.postprocess.bbox import denormalize from cvu.postprocess.backend_tf.nms.yolov5 import non_max_suppression_tf -class Yolov5(IModel): - """Implements IModel for Yolov5 using Tensorflow. +class Yolo(IDetectorModel): + """Implements IDetectorModel for Yolo using Tensorflow. This model (tensorflow-backend) performs inference, using SavedModel, on a numpy array, and returns result after performing NMS. @@ -30,13 +28,11 @@ class Yolov5(IModel): Inputs are expected to be normalized in channels-last order with batch axis. """ - def __init__(self, weight: str = "yolov5s", device='auto') -> None: + def __init__(self, weight: str, device='auto') -> None: """Initiate Model Args: - weight (str, optional): path to SavedModel weight files. Alternatively, - it also accepts identifiers (such as yolvo5s, yolov5m, etc.) to load - pretrained models. Defaults to "yolov5s". + weight (str): path to SavedModel weight files. device (str, optional): name of the device to be used. Valid devices can be "cpu", "gpu", "auto", "tpu". Defaults to "auto" which tries to use the @@ -92,18 +88,10 @@ def _load_model(self, weight: str) -> None: """Internally loads SavedModel Args: - weight (str): path to SavedModel weight files or predefined-identifiers - (such as yolvo5s, yolov5m, etc.) + weight (str): path to SavedModel weight files """ - # attempt to load predefined weights if not os.path.exists(weight): - weight += '_tensorflow' - - # get path to pretrained weights - weight = get_path(__file__, "weights", weight) - - # download weights if not already downloaded - download_weights(weight, "tensorflow", unzip=True) + raise FileNotFoundError(f"Unable to locate model weights {weight}") # set load_options needed for TPU (if needed) load_options = None @@ -143,7 +131,7 @@ def __repr__(self) -> str: Returns: str: information string """ - return f"Yolov5-Tensorflow: {self._device}" + return f"Yolo Tensorflow: {self._device}" @staticmethod def _postprocess(outputs: np.ndarray) -> List[np.ndarray]: diff --git a/cvu/detector/yolov5/backends/yolov5_tensorrt.py b/cvu/detector/yolo/backends/yolo_tensorrt.py similarity index 87% rename from cvu/detector/yolov5/backends/yolov5_tensorrt.py rename to cvu/detector/yolo/backends/yolo_tensorrt.py index f15fdfe..9321190 100644 --- a/cvu/detector/yolov5/backends/yolov5_tensorrt.py +++ b/cvu/detector/yolo/backends/yolo_tensorrt.py @@ -1,4 +1,4 @@ -"""This file contains Yolov5's IModel implementation in TensorRT. +"""This file contains Yolo's IDetectorModel implementation in TensorRT. This model (tensorRT-backend) performs inference using TensorRT, on a given input numpy array, and returns result after performing nms and other backend specific postprocessings. @@ -14,19 +14,17 @@ import pycuda.autoinit # noqa # pylint: disable=unused-import import pycuda.driver as cuda -from cvu.interface.model import IModel +from cvu.detector.interface import IDetectorModel from cvu.preprocess.image.letterbox import letterbox from cvu.preprocess.image.general import (basic_preprocess, bgr_to_rgb, hwc_to_chw) -from cvu.utils.general import get_path -from cvu.detector.yolov5.backends.common import download_weights from cvu.postprocess.nms.yolov5 import non_max_suppression_np from cvu.utils.backend_tensorrt.int8_calibrator import Int8EntropyCalibrator2 -class Yolov5(IModel): +class Yolo(IDetectorModel): # noqa # pylint: disable=too-many-instance-attributes - """Implements IModel for Yolov5 using TensorRT. + """Implements IDetectorModel for Yolo using TensorRT. This model (tensorrt-backend) performs inference, using TensorRT, on a numpy array, and returns result after performing NMS. @@ -39,6 +37,7 @@ class Yolov5(IModel): Inputs are expected to be normalized in channels-first order with/without batch axis. """ + def __init__(self, weight: str = None, num_classes: int = 80, @@ -65,7 +64,8 @@ def __init__(self, self._dtype = dtype if self._dtype == "int8": if calib_images_dir is None: - raise Exception("[CVU-Error] calib_images_dir is None with dtype int8.") + raise Exception( + "[CVU-Error] calib_images_dir is None with dtype int8.") self._calib_images_dir = calib_images_dir # initiate model specific class attributes @@ -104,31 +104,10 @@ def _load_model(self, weight: str) -> None: """Internally loads TensorRT cuda engine and creates execution context. Args: - weight (str): path to ONNX weight file, TensorRT Engine file or - predefined-identifiers (such as yolvo5s, yolov5m, etc.) + weight (str): path to ONNX weight file or TensorRT Engine file """ - # load default models using predefined-identifiers - if "." not in weight: - height, width = self._input_shape[:2] - - # get path to pretrained weights - engine_path = get_path(__file__, "weights", - f"{weight}_{height}_{width}_{self._dtype}_trt.engine") - - onnx_weight = get_path(__file__, "weights", f"{weight}_trt.onnx") - - # download onnx weights if needed, and/or generate engine file - if not os.path.exists(engine_path): - - # download weights if not already downloaded - download_weights(onnx_weight, "tensorrt") - - # build engine with current configs and load it - self._engine = self._build_engine(onnx_weight, engine_path, - self._input_shape) - else: - # deserialize and load engine - self._engine = self._deserialize_engine(engine_path) + if not os.path.exists(weight): + raise FileNotFoundError(f"Unable to locate model weights {weight}") # use custom models else: @@ -152,7 +131,8 @@ def _load_model(self, weight: str) -> None: self._context = self._engine.create_execution_context() if not self._context: raise Exception( - "[CVU-Error] Couldn't create execution context from engine successfully !") + "[CVU-Error] Couldn't create execution context from engine successfully !" + ) @staticmethod def get_supported_dtypes(builder) -> List[str]: @@ -173,7 +153,6 @@ def get_supported_dtypes(builder) -> List[str]: supported_dtypes.append("int8") return supported_dtypes - def _build_engine(self, onnx_weight: str, trt_engine_path: str, input_shape: Tuple[int]) -> trt.tensorrt.ICudaEngine: """Builds and serializes TensorRT engine by parsing the onnx model. @@ -212,7 +191,7 @@ def _build_engine(self, onnx_weight: str, trt_engine_path: str, trt.OnnxParser(network, self._logger) as parser: # get supported dtypes on this platform - supported_dtypes = Yolov5.get_supported_dtypes(builder) + supported_dtypes = Yolo.get_supported_dtypes(builder) if self._dtype not in supported_dtypes: raise Exception(f"[CVU-Error] Invalid dtype '{self._dtype}'. "\ f"Please choose from {str(supported_dtypes)}") @@ -239,13 +218,9 @@ def _build_engine(self, onnx_weight: str, trt_engine_path: str, input_w=input_shape[1], img_dir=self._calib_images_dir, preprocess=[ - letterbox, - bgr_to_rgb, - hwc_to_chw, - np.ascontiguousarray, - basic_preprocess - ] - ) + letterbox, bgr_to_rgb, hwc_to_chw, + np.ascontiguousarray, basic_preprocess + ]) # parse onnx model with open(onnx_weight, 'rb') as onnx_file: @@ -375,7 +350,7 @@ def __repr__(self) -> str: Returns: str: information string """ - return f"Yolov5s TensorRT-Cuda-{self._input_shape}" + return f"Yolo TensorRT-Cuda-{self._input_shape}" def __del__(self): """Clean up execution context stack. diff --git a/cvu/detector/yolov5/backends/yolov5_tflite.py b/cvu/detector/yolo/backends/yolo_tflite.py similarity index 79% rename from cvu/detector/yolov5/backends/yolov5_tflite.py rename to cvu/detector/yolo/backends/yolo_tflite.py index 4b0bec4..ecc0805 100644 --- a/cvu/detector/yolov5/backends/yolov5_tflite.py +++ b/cvu/detector/yolo/backends/yolo_tflite.py @@ -1,4 +1,4 @@ -"""This file contains Yolov5's IModel implementation in TFLite. +"""This file contains Yolo's IDetectorModel implementation in TFLite. This model (tflite-backend) performs inference using TFLite, on a given input numpy array, and returns result after performing nms and other backend specific postprocessings. @@ -12,27 +12,24 @@ import numpy as np import tensorflow.lite as tflite -from cvu.interface.model import IModel -from cvu.utils.general import get_path -from cvu.detector.yolov5.backends.common import download_weights +from cvu.detector.interface import IDetectorModel from cvu.postprocess.backend_tf.nms.yolov5 import non_max_suppression_tf -class Yolov5(IModel): - """Implements IModel for Yolov5 using TFLite. +class Yolo(IDetectorModel): + """Implements IDetectorModel for Yolo using TFLite. This model (tflite-backend) performs inference, using TFLite, on a numpy array, and returns result after performing NMS. Inputs are expected to be normalized in channels-last order with batch axis. """ - def __init__(self, weight: str = "yolov5s", device='auto') -> None: + + def __init__(self, weight: str, device='auto') -> None: """Initiate Model Args: - weight (str, optional): path to SavedModel weight files. Alternatively, - it also accepts identifiers (such as yolvo5s, yolov5m, etc.) to load - pretrained models. Defaults to "yolov5s". + weight (str, optional): path to TFLite weight files. device (str, optional): name of the device to be used. Valid devices can be "cpu", "auto". Defaults to "auto" which tries to use the @@ -71,18 +68,10 @@ def _load_model(self, weight: str) -> None: """Internally loads TFLite Model Args: - weight (str): path to TFLite weight files or predefined-identifiers - (such as yolvo5s, yolov5m, etc.) + weight (str): path to TFLite weight files """ - # attempt to load predefined weights if not os.path.exists(weight): - weight += '.tflite' - - # get path to pretrained weights - weight = get_path(__file__, "weights", weight) - - # download weights if not already downloaded - download_weights(weight, "tflite") + raise FileNotFoundError(f"Unable to locate model weights {weight}") # load model self._model = tflite.Interpreter(model_path=weight) # pylint: disable=maybe-no-member @@ -129,7 +118,7 @@ def __repr__(self) -> str: Returns: str: information string """ - return f"Yolov5: {self._device}" + return f"Yolo TFLite-{self._device}" @staticmethod def _postprocess(outputs: np.ndarray) -> List[np.ndarray]: diff --git a/cvu/detector/yolov5/backends/yolov5_torch.py b/cvu/detector/yolo/backends/yolo_torch.py similarity index 80% rename from cvu/detector/yolov5/backends/yolov5_torch.py rename to cvu/detector/yolo/backends/yolo_torch.py index 5d9366e..44d9a32 100644 --- a/cvu/detector/yolov5/backends/yolov5_torch.py +++ b/cvu/detector/yolo/backends/yolo_torch.py @@ -1,4 +1,4 @@ -"""This file contains Yolov5's IModel implementation in PyTorch. +"""This file contains Yolo's IDetectorModel implementation in PyTorch. This model (torch-backend) performs inference using Torch-script, on a given input numpy array, and returns result after performing nms and other backend specific postprocessings. @@ -12,14 +12,12 @@ import numpy as np import torch -from cvu.interface.model import IModel -from cvu.utils.general import get_path -from cvu.detector.yolov5.backends.common import download_weights +from cvu.detector.interface import IDetectorModel from cvu.postprocess.backend_torch.nms.yolov5 import non_max_suppression_torch -class Yolov5(IModel): - """Implements IModel for Yolov5 using PyTorch. +class Yolo(IDetectorModel): + """Implements IDetectorModel for Yolo using PyTorch. This model (torch-backend) performs inference, using Torch-script, on a numpy array, and returns result after performing NMS. @@ -28,13 +26,11 @@ class Yolov5(IModel): (with/without batch axis). """ - def __init__(self, weight: str = "yolov5s", device='auto') -> None: + def __init__(self, weight: str, device='auto') -> None: """Initiate Model Args: - weight (str, optional): path to jit-script .pt weight files. Alternatively, - it also accepts identifiers (such as yolvo5s, yolov5m, etc.) to load - pretrained models. Defaults to "yolov5s". + weight (str, optional): path to jit-script .pt weight files. device (str, optional): name of the device to be used. Valid devices can be "cpu", "gpu", "auto" or specific cuda devices such as @@ -67,19 +63,10 @@ def _load_model(self, weight: str) -> None: """Internally loads JIT-Model Args: - weight (str): path to jit-script .pt weight files or predefined-identifiers - (such as yolvo5s, yolov5m, etc.) + weight (str): path to jit-script .pt weight files """ - # attempt to load predefined weights if not os.path.exists(weight): - if self._device.type != 'cpu': - weight += '.cuda' - - # get path to pretrained weights - weight = get_path(__file__, "weights", f"{weight}.torchscript.pt") - - # download weights if not already downloaded - download_weights(weight, "torch") + raise FileNotFoundError(f"Unable to locate model weights {weight}") # load model self._model = torch.jit.load(weight, map_location=self._device) @@ -113,7 +100,7 @@ def __repr__(self) -> str: Returns: str: information string """ - return f"Yolov5: {self._device.type}" + return f"Yolo Torch-{self._device.type}" def _preprocess(self, inputs: np.ndarray) -> torch.Tensor: """Process inputs for model inference. diff --git a/cvu/detector/yolo/common.py b/cvu/detector/yolo/common.py new file mode 100644 index 0000000..59906bb --- /dev/null +++ b/cvu/detector/yolo/common.py @@ -0,0 +1,65 @@ +"""This file contains common functions used between different backends. +""" +import os + +from cvu.utils.google_utils import gdrive_download +from cvu.utils.general import (load_json, get_path) + + +def get_weights_filename(weights: str, backend: str) -> str: + save_path = "" + if backend == "torch": + save_path = f"{weights}_torch.torchscript.pt.pt" + elif backend == "onnx": + save_path = f"{weights}_onnx.onnx" + elif backend == "tensorrt": + save_path = f"{weights}_tensorrt.onnx" + elif backend == "tensorflow": + save_path = f"{weights}_tensorflow" + elif backend == "tflite": + save_path = f"{weights}_tflite.tflite" + return save_path + + +def download_weights( + yolo_version: str, + weights: str, + backend: str, + unzip: bool = False, +) -> str: + """Download weight if not downloaded already. + + Args: + yolo_version (str): name of yolo version (i.e. yolov5, yolov7 etc.) + weights (str): name of the weights (i.e. yolov5s, yolov7s etc.) + backend (str): name of the backend + unzip (bool): unzip downloaded file + + Raises: + FileNotFoundError: raised if weight is not a valid pretrained + weight name. + + Returns: + save_path (str): path where file is saved + """ + # already downloaded + save_path = get_weights_filename(weights, backend) + if os.path.exists(save_path): + return save_path + + # get dict of all available pretrained weights + weights_json = get_path(__file__, "..", "weights.json") + available_weights = load_json(weights_json) + + weight_key = available_weights.get(yolo_version.lower(), + {}).get(weights, {}).get(backend, None) + + # check if a valid weight is requested + if not weight_key: + raise FileNotFoundError( + f"Invalid default weights {weights} for model {yolo_version.title()}-{backend}" + ) + + # download weights + gdrive_download(weight_key, save_path, unzip=unzip) + return save_path diff --git a/cvu/detector/yolo/core.py b/cvu/detector/yolo/core.py new file mode 100644 index 0000000..8dcca87 --- /dev/null +++ b/cvu/detector/yolo/core.py @@ -0,0 +1,224 @@ +"""This file contains Yolo's core IDetector implementation. + +Yolo Core represents a common interface over which users can perform +Yolo powered object detection, without having to worry about the details of +different backends and their specific requirements and implementations. Core +will internally resolve and handle all internal details. +""" +from typing import Any, Callable, Tuple, Union, List +from importlib import import_module + +import numpy as np + +from cvu.detector.interface import IDetector +from cvu.detector.predictions import Predictions +from cvu.detector.configs import COCO_CLASSES +from cvu.preprocess.image.letterbox import letterbox +from cvu.preprocess.image.general import (basic_preprocess, bgr_to_rgb, + hwc_to_chw) +from cvu.postprocess.bbox import scale_coords +from cvu.utils.backend import setup_backend + + +class Yolo(IDetector): + """Implements ICore for Yolo + + Yolo Core represents a common interface to perform + Yolo (v5 & v7) powered object detection. + """ + _BACKEND_PKG = "cvu.detector.yolo.backends" + + def __init__(self, + classes: Union[str, List[str]], + backend: str, + weight: str, + device: str = "auto", + auto_install: bool = False, + **kwargs: Any) -> None: + """Initiate Yolo Object Detector + + Args: + classes (Union[str, List[str]]): name of classes to be detected. + It can be set to individual classes like 'coco', 'person', 'cat' etc. + Alternatively, it also accepts list of classes such as ['person', 'cat']. + For default models/weights, 'classes' is used to filter out objects + according to provided argument from coco class. For custom models, all + classes should be provided in original order as list of strings. + + backend (str): name of the backend to be used for inference purposes. + + weight (str): path to weight files (according to selected backend). + + device (str, optional): name of the device to be used. Valid + devices can be "cpu", "gpu", "tpu", "auto". Defaults to "auto" which tries + to use the device best suited for selected backend and the hardware avaibility. + + auto_install (bool, optional): auto install missing requirements for the selected + backend. + """ + # initiate class attributes + if kwargs.get("input_shape", None) is not None: + self._preprocess = [ + lambda image: letterbox( + image, kwargs['input_shape'], auto=False), bgr_to_rgb + ] + else: + self._preprocess = [letterbox, bgr_to_rgb] + + self._postprocess = [] + self._classes = {} + self._model = None + + # setup backend and load model + if auto_install: + setup_backend(backend, device) + self._load_classes(classes) + self._load_model(backend, weight, device, **kwargs) + + def __repr__(self) -> str: + """Returns Backend and Model Information + + Returns: + str: information string + """ + return str(self._model) + + def __call__(self, inputs: np.ndarray, **kwargs) -> Predictions: + """Performs Yolo Object Detection on given inputs. + Returns detected objects as Predictions object. + + Args: + inputs (np.ndarray): image in BGR format. + + Returns: + Predictions: detected objects. + """ + # preprocess + processed_inputs = self._apply(inputs, self._preprocess) + + # inference on backend + outputs = self._model(processed_inputs) + + # postprocess + outputs = self._apply(outputs, self._postprocess) + + # scale up + outputs = self._scale(inputs.shape, processed_inputs.shape, outputs) + + # convert to preds + return self._to_preds(outputs) + + @staticmethod + def _apply( + value: np.ndarray, + functions: List[Callable[[np.ndarray], np.ndarray]]) -> np.ndarray: + """Recursively applies list of callable functions to given value + + Args: + value (np.ndarray): input to be processed + + functions (List[Callable[[np.ndarray], np.ndarray]]): list of + callable functions. + + Returns: + np.ndarray: value resulting from applying all functions + """ + for func in functions: + value = func(value) + return value + + @staticmethod + def _scale(original_shape: Tuple[int], process_shape: Tuple[int], + outputs: np.ndarray) -> np.ndarray: + """Scale outputs based on process_shape to original_shape + + Args: + original_shape (Tuple[int]): shape of original inputs. + process_shape (Tuple[int]): shape of processed inputs. + outputs (np.ndarray): outputs from yolo model + + Returns: + np.ndarray: scaled outputs + """ + # channels first, pick widht-height accordingly + if len(process_shape) > 3: + process_shape = process_shape[1:] + + if process_shape[2] != 3: + process_shape = process_shape[1:] + + # scale bounding box + outputs[:, :4] = scale_coords(process_shape[:2], outputs[:, :4], + original_shape).round() + return outputs + + def _load_model(self, backend_name: str, weight: str, device: str, + **kwargs: Any) -> None: + """Internally loads Model (backend) + + Args: + backend_name (str): name of the backend + weight (str): path to weight file or default identifiers + device (str): name of target device (auto, cpu, gpu, tpu) + """ + # load model + backend = import_module(f".yolo_{backend_name}", self._BACKEND_PKG) + + if backend_name != 'tensorrt': + self._model = backend.Yolo(weight, device) + else: + self._model = backend.Yolo(weight, + num_classes=len(self._classes), + **kwargs) + + # add preprocess + if backend_name in ['torch', 'onnx', 'tensorrt']: + self._preprocess.append(hwc_to_chw) + + # contigousarray + self._preprocess.append(np.ascontiguousarray) + + # add common preprocess + if backend_name in ['onnx', 'tensorflow', 'tflite', 'tensorrt']: + self._preprocess.append(basic_preprocess) + + def _load_classes(self, classes: Union[str, List[str]]) -> None: + """Internally loads target classes + + Args: + classes (Union[str, List[str]]): name or list of classes to be detected. + """ + if classes == 'coco': + classes = COCO_CLASSES + + elif isinstance(classes, str): + classes = [classes] + + if set(classes).issubset(COCO_CLASSES): + for i, name in enumerate(COCO_CLASSES): + if name in classes: + self._classes[i] = name + else: + self._classes = dict(enumerate(classes)) + + def _to_preds(self, outputs: np.ndarray) -> Predictions: + """Convert Yolo's numpy inputs to Predictions object. + + Args: + outputs (np.ndarray): basic outputs from yolo inference. + + Returns: + Predictions: detected objects + """ + # create container + preds = Predictions() + + # add detection + for *xyxy, conf, class_id in outputs: + # filter class + if class_id in self._classes: + preds.create_and_append(xyxy, + conf, + class_id, + class_name=self._classes[class_id]) + return preds diff --git a/cvu/detector/yolov5/backends/common.py b/cvu/detector/yolov5/backends/common.py deleted file mode 100644 index 6a937a8..0000000 --- a/cvu/detector/yolov5/backends/common.py +++ /dev/null @@ -1,37 +0,0 @@ -"""This file contains common functions used between different backends. -""" -import os - -from cvu.utils.google_utils import gdrive_download -from cvu.utils.general import (load_json, get_path) - - -def download_weights(weight: str, backend: str, unzip=False) -> None: - """Download weight if not downloaded already. - - Args: - weight (str): path where weights should be downloaded - backend (str): name of the backend - unzip (bool, optional): unzip downloaded file. Defaults to False. - - Raises: - FileNotFoundError: raised if weight is not a valid pretrained - weight name. - """ - # already downloaded - if os.path.exists(weight): - return - - # get weight's identifier key - weight_key = os.path.split(weight)[-1] - - # get dict of all available pretrained weights - weights_json = get_path(__file__, "weights", f"{backend}_weights.json") - available_weights = load_json(weights_json) - - # check if a valid weight is requested - if weight_key not in available_weights: - raise FileNotFoundError - - # download weights - gdrive_download(available_weights[weight_key], weight, unzip) diff --git a/cvu/detector/yolov5/backends/weights/onnx_weights.json b/cvu/detector/yolov5/backends/weights/onnx_weights.json deleted file mode 100644 index df032e5..0000000 --- a/cvu/detector/yolov5/backends/weights/onnx_weights.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "yolov5s.onnx": "1piC3ZGuc4D8MMJQQRK3dgaCa66-4Ucxi" -} \ No newline at end of file diff --git a/cvu/detector/yolov5/backends/weights/tensorflow_weights.json b/cvu/detector/yolov5/backends/weights/tensorflow_weights.json deleted file mode 100644 index 2578a78..0000000 --- a/cvu/detector/yolov5/backends/weights/tensorflow_weights.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "yolov5s_tensorflow":"1SA7hT4jUx9szqJePZ8NssMPSCelRvYpo" -} diff --git a/cvu/detector/yolov5/backends/weights/tensorrt_weights.json b/cvu/detector/yolov5/backends/weights/tensorrt_weights.json deleted file mode 100644 index 1c46041..0000000 --- a/cvu/detector/yolov5/backends/weights/tensorrt_weights.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "yolov5s_trt.onnx":"1piC3ZGuc4D8MMJQQRK3dgaCa66-4Ucxi" -} diff --git a/cvu/detector/yolov5/backends/weights/tflite_weights.json b/cvu/detector/yolov5/backends/weights/tflite_weights.json deleted file mode 100644 index 8ec108a..0000000 --- a/cvu/detector/yolov5/backends/weights/tflite_weights.json +++ /dev/null @@ -1,3 +0,0 @@ -{ - "yolov5s.tflite": "1oeKZm81Gz5OejmcgDA0biemkPWCY6qNu" -} \ No newline at end of file diff --git a/cvu/detector/yolov5/backends/weights/torch_weights.json b/cvu/detector/yolov5/backends/weights/torch_weights.json deleted file mode 100644 index 6f21448..0000000 --- a/cvu/detector/yolov5/backends/weights/torch_weights.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "yolov5s.cuda.torchscript.pt":"16BFVGsEYAXsupFoCXxlOhSBovoMH6juK", - "yolov5s.torchscript.pt":"1cWeCusb2IB-x-h_IZs7t3Y_rs9GkzQsI" -} \ No newline at end of file diff --git a/cvu/detector/yolov5/core.py b/cvu/detector/yolov5/core.py index a411bda..88bacd6 100644 --- a/cvu/detector/yolov5/core.py +++ b/cvu/detector/yolov5/core.py @@ -8,28 +8,19 @@ Find more about Yolov5 here from their official repository https://github.com/ultralytics/yolov5 """ -from typing import Any, Callable, Tuple, Union, List -from importlib import import_module +import os +from typing import Any, Union, List -import numpy as np +from cvu.detector.yolo import Yolo +from cvu.detector.yolo.common import download_weights -from cvu.interface.core import ICore -from cvu.detector.predictions import Predictions -from cvu.detector.configs import COCO_CLASSES -from cvu.preprocess.image.letterbox import letterbox -from cvu.preprocess.image.general import (basic_preprocess, bgr_to_rgb, - hwc_to_chw) -from cvu.postprocess.bbox import scale_coords -from cvu.utils.backend import setup_backend - -class Yolov5(ICore): +class Yolov5(Yolo): """Implements ICore for Yolov5 Yolov5 Core represents a common interface to perform Yolov5 powered object detection. """ - _BACKEND_PKG = "cvu.detector.yolov5.backends" def __init__(self, classes: Union[str, List[str]], @@ -62,168 +53,10 @@ def __init__(self, auto_install (bool, optional): auto install missing requirements for the selected backend. """ - # ICore - super().__init__(classes, backend) - - # initiate class attributes - if kwargs.get("input_shape", None) is not None: - self._preprocess = [lambda image: letterbox(image, kwargs['input_shape'], auto=False), - bgr_to_rgb] - else: - self._preprocess = [letterbox, bgr_to_rgb] - self._postprocess = [] - self._classes = {} - self._model = None - - # setup backend and load model - if auto_install: - setup_backend(backend, device) - self._load_classes(classes) - self._load_model(backend, weight, device, **kwargs) - - def __repr__(self) -> str: - """Returns Backend and Model Information - - Returns: - str: information string - """ - return str(self._model) - - def __call__(self, inputs: np.ndarray, **kwargs) -> Predictions: - """Performs Yolov5 Object Detection on given inputs. - Returns detected objects as Predictions object. - - Args: - inputs (np.ndarray): image in BGR format. - - Returns: - Predictions: detected objects. - """ - # preprocess - processed_inputs = self._apply(inputs, self._preprocess) - - # inference on backend - outputs = self._model(processed_inputs) - - # postprocess - outputs = self._apply(outputs, self._postprocess) - - # scale up - outputs = self._scale(inputs.shape, processed_inputs.shape, outputs) - - # convert to preds - return self._to_preds(outputs) - - @staticmethod - def _apply( - value: np.ndarray, - functions: List[Callable[[np.ndarray], np.ndarray]]) -> np.ndarray: - """Recursively applies list of callable functions to given value - - Args: - value (np.ndarray): input to be processed - - functions (List[Callable[[np.ndarray], np.ndarray]]): list of - callable functions. - - Returns: - np.ndarray: value resulting from applying all functions - """ - for func in functions: - value = func(value) - return value - - @staticmethod - def _scale(original_shape: Tuple[int], process_shape: Tuple[int], - outputs: np.ndarray) -> np.ndarray: - """Scale outputs based on process_shape to original_shape - - Args: - original_shape (Tuple[int]): shape of original inputs. - process_shape (Tuple[int]): shape of processed inputs. - outputs (np.ndarray): outputs from yolov5 model - - Returns: - np.ndarray: scaled outputs - """ - # channels first, pick widht-height accordingly - if len(process_shape) > 3: - process_shape = process_shape[1:] - - if process_shape[2] != 3: - process_shape = process_shape[1:] - - # scale bounding box - outputs[:, :4] = scale_coords(process_shape[:2], outputs[:, :4], - original_shape).round() - return outputs - - def _load_model(self, backend_name: str, weight: str, device: str, **kwargs: Any) -> None: - """Internally loads Model (backend) - - Args: - backend_name (str): name of the backend - weight (str): path to weight file or default identifiers - device (str): name of target device (auto, cpu, gpu, tpu) - """ - # load model - backend = import_module(f".yolov5_{backend_name}", self._BACKEND_PKG) - - if backend_name != 'tensorrt': - self._model = backend.Yolov5(weight, device) - else: - self._model = backend.Yolov5(weight, - num_classes=len(self._classes), - **kwargs) - - # add preprocess - if backend_name in ['torch', 'onnx', 'tensorrt']: - self._preprocess.append(hwc_to_chw) - - # contigousarray - self._preprocess.append(np.ascontiguousarray) - - # add common preprocess - if backend_name in ['onnx', 'tensorflow', 'tflite', 'tensorrt']: - self._preprocess.append(basic_preprocess) - - def _load_classes(self, classes: Union[str, List[str]]) -> None: - """Internally loads target classes - - Args: - classes (Union[str, List[str]]): name or list of classes to be detected. - """ - if classes == 'coco': - classes = COCO_CLASSES - - elif isinstance(classes, str): - classes = [classes] - - if set(classes).issubset(COCO_CLASSES): - for i, name in enumerate(COCO_CLASSES): - if name in classes: - self._classes[i] = name - else: - self._classes = dict(enumerate(classes)) - - def _to_preds(self, outputs: np.ndarray) -> Predictions: - """Convert Yolov5's numpy inputs to Predictions object. - - Args: - outputs (np.ndarray): basic outputs from yolov5 inference. - - Returns: - Predictions: detected objects - """ - # create container - preds = Predictions() + if not os.path.exists(weight): + # attemp weight download + weight = download_weights("yolov5", weight, backend, + backend == "tensorflow") - # add detection - for *xyxy, conf, class_id in outputs: - # filter class - if class_id in self._classes: - preds.create_and_append(xyxy, - conf, - class_id, - class_name=self._classes[class_id]) - return preds + super().__init__(classes, backend, weight, device, auto_install, + **kwargs) diff --git a/cvu/interface/core.py b/cvu/interface/core.py index 4818363..a65e3e7 100644 --- a/cvu/interface/core.py +++ b/cvu/interface/core.py @@ -3,7 +3,6 @@ For example, YoloV5 is a core of Object Detection use cases. """ import abc -from typing import Union, List import numpy as np @@ -15,13 +14,12 @@ class ICore(metaclass=abc.ABCMeta): A core defines one complete method/solution for certain use cases. For example, YoloV5 is a core of Object Detection use cases. """ + @abc.abstractmethod - def __init__(self, classes: Union[str, List[str]], backend: str, *args, - **kwargs) -> None: + def __init__(self, backend: str, *args, **kwargs) -> None: """Initiate Core. Args: - classes (Union[str, List[str]]): single object class name or list of classes backend (str): name of the backend to run core on. """ ... diff --git a/cvu/interface/model.py b/cvu/interface/model.py deleted file mode 100644 index 1954b0b..0000000 --- a/cvu/interface/model.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Defines interface for Model that represents CVU-Backends. -A Model combines the process for individual model inference for certain backend. -For example, YoloV5Torch can be a model of Yolov5 a core. -""" -import abc - -import numpy as np - - -class IModel(metaclass=abc.ABCMeta): - """Model Interface which will be implemented for every CVU-Backend. - A Model combines the process for individual model inference - for certain backend. - """ - @abc.abstractmethod - def __call__(self, inputs: np.ndarray) -> np.ndarray: - """Execute core on inputs - - Args: - inputs (np.ndarray): inputs to be exectued core on - - Returns: - Predictions: results of executation - """ - ... - - @abc.abstractmethod - def __repr__(self) -> str: - """Represents model with method and configuration informations. - - Returns: - str: formatted string with method and config info. - """ - ...