Added logger to detect target scripts

UWARG · Nov 7, 2024 · f2bcff8 · f2bcff8
1 parent fd3cea9
commit f2bcff8
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 1 deletion.
diff --git a/modules/detect_target/detect_target_factory.py b/modules/detect_target/detect_target_factory.py
@@ -6,6 +6,7 @@
 
 from . import base_detect_target
 from . import detect_target_ultralytics
+from ..common.modules.logger import logger
 
 
 class DetectTargetOption(enum.Enum):
@@ -21,6 +22,7 @@ def create_detect_target(
     device: "str | int",
     model_path: str,
     override_full: bool,
+    local_logger: logger.Logger,
     show_annotations: bool,
     save_name: str,
 ) -> tuple[bool, base_detect_target.BaseDetectTarget | None]:
@@ -33,6 +35,7 @@ def create_detect_target(
                 device,
                 model_path,
                 override_full,
+                local_logger,
                 show_annotations,
                 save_name,
             )

diff --git a/modules/detect_target/detect_target_ultralytics.py b/modules/detect_target/detect_target_ultralytics.py
@@ -10,6 +10,7 @@
 from . import base_detect_target
 from .. import image_and_time
 from .. import detections_and_time
+from ..common.modules.logger import logger
 
 
 class DetectTargetUltralytics(base_detect_target.BaseDetectTarget):
@@ -22,6 +23,7 @@ def __init__(
         device: "str | int",
         model_path: str,
         override_full: bool,
+        local_logger: logger.Logger,
         show_annotations: bool = False,
         save_name: str = "",
     ) -> None:
@@ -36,6 +38,7 @@ def __init__(
         self.__model = ultralytics.YOLO(model_path)
         self.__counter = 0
         self.__enable_half_precision = not self.__device == "cpu"
+        self.__local_logger = local_logger
         self.__show_annotations = show_annotations
         if override_full:
             self.__enable_half_precision = False
@@ -54,6 +57,7 @@ def run(
         Return: Success and the detections.
         """
         image = data.image
+        start_time = time.time()
         predictions = self.__model.predict(
             source=image,
             half=self.__enable_half_precision,
@@ -91,6 +95,10 @@ def run(
 
                 detections.append(detection)
 
+        self.__local_logger.info(
+            f"{time.time()} Object detection took {time.time() - start_time} seconds", True
+        )
+
         # Logging
         if self.__filename_prefix != "":
             filename = self.__filename_prefix + str(self.__counter)

diff --git a/modules/detect_target/detect_target_worker.py b/modules/detect_target/detect_target_worker.py
@@ -47,6 +47,7 @@ def detect_target_worker(
         device,
         model_path,
         override_full,
+        local_logger,
         show_annotations,
         save_name,
     )

diff --git a/tests/unit/test_detect_target_ultralytics.py b/tests/unit/test_detect_target_ultralytics.py
@@ -13,6 +13,7 @@
 from modules.detect_target import detect_target_ultralytics
 from modules import image_and_time
 from modules import detections_and_time
+from modules.common.modules.logger import logger
 
 
 TEST_PATH = pathlib.Path("tests", "model_example")
@@ -108,8 +109,11 @@ def detector() -> detect_target_ultralytics.DetectTargetUltralytics:  # type: ig
     """
     Construct DetectTargetUltralytics.
     """
+
+    _, test_logger = logger.Logger.create("test_logger", False)
+
     detection = detect_target_ultralytics.DetectTargetUltralytics(
-        DEVICE, str(MODEL_PATH), OVERRIDE_FULL
+        DEVICE, str(MODEL_PATH), OVERRIDE_FULL, test_logger
     )
     yield detection  # type: ignore