diff --git a/timm_image.py b/timm_image.py
index 3d2c84d..32ed98d 100644
--- a/timm_image.py
+++ b/timm_image.py
@@ -15,8 +15,8 @@
 import logging
 import numpy
 
+import towhee
 from towhee.operator.base import NNOperator, OperatorFlag
-from towhee.types import Image as towheeImage
 from towhee.types.arg import arg, to_image_color
 from towhee import register
 
@@ -59,7 +59,7 @@ class TimmImage(NNOperator):
         self.skip_tfms = skip_preprocess
 
     @arg(1, to_image_color('RGB'))
-    def __call__(self, img: towheeImage) -> numpy.ndarray:
+    def __call__(self, img: towhee._types.Image) -> numpy.ndarray:
         img = PILImage.fromarray(img.astype('uint8'), 'RGB')
         if not self.skip_tfms:
             img = self.tfms(img).unsqueeze(0)