File size: 1,868 Bytes
da77aaf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
"""
Copyright (c) 2024-present Naver Cloud Corp.
This source code is based on code from the Segment Anything Model (SAM)
(https://github.com/facebookresearch/segment-anything).
This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""
import torch
from typing import Any, Callable
import onnxruntime
def np2tensor(np_array, device):
return torch.from_numpy(np_array).to(device)
def tensor2np(torch_tensor):
return torch_tensor.detach().cpu().numpy()
class ZIM_Encoder():
def __init__(self, onnx_path, num_threads=16):
self.onnx_path = onnx_path
sessionOptions = onnxruntime.SessionOptions()
sessionOptions.intra_op_num_threads = num_threads
sessionOptions.inter_op_num_threads = num_threads
providers = ["CPUExecutionProvider"]
self.ort_session = onnxruntime.InferenceSession(
onnx_path, sess_options=sessionOptions, providers=providers
)
def cuda(self, device_id=0):
providers = [
(
"CUDAExecutionProvider",
{
"device_id": device_id,
},
),
]
self.ort_session.set_providers(providers)
def forward(
self,
image,
):
device = image.device
ort_inputs = {
"image": tensor2np(image),
}
image_embeddings, feat_D0, feat_D1, feat_D2 = self.ort_session.run(None, ort_inputs)
image_embeddings = np2tensor(image_embeddings, device)
feat_D0 = np2tensor(feat_D0, device)
feat_D1 = np2tensor(feat_D1, device)
feat_D2 = np2tensor(feat_D2, device)
return image_embeddings, (feat_D0, feat_D1, feat_D2)
__call__: Callable[..., Any] = forward
|