logo
Browse Source

Update

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
test
Jael Gu 3 years ago
parent
commit
7e51ef22d1
  1. 5
      timm_image.py

5
timm_image.py

@ -57,7 +57,7 @@ class TimmImage(NNOperator):
def __call__(self, img: numpy.ndarray) -> numpy.ndarray: def __call__(self, img: numpy.ndarray) -> numpy.ndarray:
if hasattr(img, 'mode'): if hasattr(img, 'mode'):
if img.mode not in ['RGB', 'BGR']: if img.mode not in ['RGB', 'BGR']:
log.error(f'Invalid image mode: expect "RGB" but receive "{img.mode}".')
log.error(f'Invalid image mode: expect "RGB" or "BGR" but receive "{img.mode}".')
raise AssertionError(f'Invalid image mode "{img.mode}".') raise AssertionError(f'Invalid image mode "{img.mode}".')
elif img.mode == 'BGR': elif img.mode == 'BGR':
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
@ -79,14 +79,13 @@ class TimmImage(NNOperator):
# if __name__ == '__main__': # if __name__ == '__main__':
# import cv2
# from towhee._types import Image # from towhee._types import Image
# #
# #
# path = '/path/to/image' # path = '/path/to/image'
# img = cv2.imread(path) # img = cv2.imread(path)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# img = Image(img, 'RGB')
# img = Image(img)
# #
# op = TimmImage('resnet50') # op = TimmImage('resnet50')
# out = op(img) # out = op(img)

Loading…
Cancel
Save