Browse Source
Fix image mode issue
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
2 additions and
1 deletions
-
benchmark/run.py
|
|
|
@ -49,7 +49,7 @@ query_data = query_data[:query_size] if query_size else query_data |
|
|
|
# Warm up |
|
|
|
print('Warming up...') |
|
|
|
op = ops.image_embedding.timm(model_name=model_name, device=device).get_op() |
|
|
|
dummy_input = numpy.random.randn(1, 3, op.config['input_size'][-2], op.config['input_size'][-1]) |
|
|
|
dummy_input = numpy.random.randn(op.config['input_size'][0]) |
|
|
|
dim = op(dummy_input).shape[0] |
|
|
|
print(f'output dim: {dim}') |
|
|
|
|
|
|
|
@ -120,6 +120,7 @@ elif args.format == 'onnx': |
|
|
|
|
|
|
|
@towhee.register |
|
|
|
def run_onnx(img): |
|
|
|
img = img.convert('RGB') |
|
|
|
img = op.tfms(img).cpu().detach().numpy() |
|
|
|
features = sess.run(output_names=['output_0'], input_feed={'input_0': img})[0] |
|
|
|
if len(features.shape) == 4: |
|
|
|
|
