Spaces:
Runtime error
Runtime error
File size: 2,117 Bytes
213de71 10b0510 213de71 4dc7058 213de71 |
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 64 65 66 67 68 69 70 71 72 73 74 |
import onnxruntime
import random
import numpy as np
from pathlib import Path
from numpy.typing import NDArray
from typing import Any, List
from utils import count_gpus, get_memory_free_MiB
from abc import ABC, abstractclassmethod
__dir__ = Path(__file__).parent
class ONNXBaseTask(ABC):
num_gpus: int = 0
def __init__(self, weight: str) -> None:
self.session = self.initialize_session(weight)
self.input_metadata = self.session.get_inputs()[0]
self.prepare_input = self.setup_prepare_input()
# warmup model
input_height, input_width = self.input_metadata.shape[-2:]
temp = np.zeros((1, 3, int(input_height) if int(input_height) > 0 else 320, int(input_width) if int(input_width) > 0 else 320), dtype=np.float32)
self.run_session(temp)
@abstractclassmethod
def process_output(self, raw_outputs: List[NDArray], **kwargs) -> Any:
pass
@abstractclassmethod
def setup_prepare_input(self):
pass
def call(self, image) -> Any:
input_height, input_width = self.input_metadata.shape[-2:]
# predict
input_value = self.prepare_input(image, height=input_height, width=input_width)
raw_outputs = self.run_session(input_value)
return self.process_output(raw_outputs)
def run_session(self, input_value: NDArray) -> List[NDArray]:
input_dict = {self.input_metadata.name : input_value}
return self.session.run(None, input_dict)
def initialize_session(self, weight: str):
# get avaiable runtime
providers=[]
if self.num_gpus == 0:
providers += [("CPUExecutionProvider", {})]
else:
providers += [(
"CUDAExecutionProvider",
{
"device_id": random.choice([i for i in range(self.num_gpus) if get_memory_free_MiB(i) >= 1000])
}
)]
# init session
return onnxruntime.InferenceSession(
str(__dir__.parent.parent.parent/weight),
providers=providers
)
|