From 9dc048dd584f33de05da740ed8ee08301921c40d Mon Sep 17 00:00:00 2001
From: "junjie.jiang" <junjie.jiang@zilliz.com>
Date: Fri, 12 Jul 2024 11:01:11 +0800
Subject: [PATCH] Add timer

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
---
 image_decode_cv2.py | 42 +++++++++++++++++++++++++++++++++++-------
 1 file changed, 35 insertions(+), 7 deletions(-)

diff --git a/image_decode_cv2.py b/image_decode_cv2.py
index 5e07424..619cda9 100644
--- a/image_decode_cv2.py
+++ b/image_decode_cv2.py
@@ -1,32 +1,60 @@
 # encoding=utf-8
 
-from collections import namedtuple
-
+import time
 import cv2
 import requests
 import numpy as np
 
 from towhee._types import Image
 from towhee.operator import PyOperator, SharedType
-from towhee import register
-from towhee.utils import ndarray_utils
 from towhee.utils.log import engine_log
 
 
+class Timer:
+    """
+    Timer
+    """
+    def __init__(self, text: str = 'Elapsed time: {:0.4f} seconds'):
+        self._text = text
+        self._start_time = None
+
+    def start(self):
+        if self._start_time is not None:
+            raise RuntimeError('Timer is running. Use .stop() to stop it')
+        self._start_time = time.perf_counter()
+
+    def stop(self):
+        if self._start_time is None:
+            raise RuntimeError('Timer is not running. Use .start() to start it')
+        elapsed_time = time.perf_counter() - self._start_time
+        self._start_time = None
+        engine_log.error(self._text.format(elapsed_time))
+        return elapsed_time
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, *exc_info):
+        self.stop()
+
+
 class ImageDecoder(PyOperator):
     def __init__(self):
         super().__init__()
 
     def _load_from_remote(image_url: str) -> np.ndarray:
         try:
-            r = requests.get(image_url, timeout=(20, 20))
+            with Timer("Image download time: {:0.4f} seconds"):
+                r = requests.get(image_url, timeout=(20, 20))
             if r.status_code // 100 != 2:
                 engine_log.error('Download image from %s failed, error msg: %s, request code: %s ' % (image_url,
                                                                                                       r.text,
                                                                                                       r.status_code))
                 return None
-            arr = np.asarray(bytearray(r.content), dtype=np.uint8)
-            return cv2.imdecode(arr, -1)
+            with Timer("Image decode time: {:0.4f} seconds"):
+                arr = np.asarray(bytearray(r.content), dtype=np.uint8)
+                return cv2.imdecode(arr, -1)
         except Exception as e:
             engine_log.error('Download image from %s failed, error msg: %s' % (image_url, str(e)))
             return False