mirror of
https://github.com/azaion/annotations.git
synced 2026-04-23 02:16:29 +00:00
b21f8e320f
add tensorrt engine
29 lines
716 B
Cython
29 lines
716 B
Cython
from typing import List, Tuple
|
|
import numpy as np
|
|
|
|
|
|
cdef class InferenceEngine:
|
|
cdef public int batch_size
|
|
cdef tuple get_input_shape(self)
|
|
cdef int get_batch_size(self)
|
|
cdef get_class_names(self)
|
|
cpdef run(self, input_data)
|
|
|
|
cdef class OnnxEngine(InferenceEngine):
|
|
cdef object session
|
|
cdef list model_inputs
|
|
cdef str input_name
|
|
cdef object input_shape
|
|
cdef object class_names
|
|
|
|
cdef class TensorRTEngine(InferenceEngine):
|
|
cdef object stream
|
|
cdef object context
|
|
cdef str input_name
|
|
cdef str output_name
|
|
cdef object d_input
|
|
cdef object d_output
|
|
cdef object input_shape
|
|
cdef object output_shape
|
|
cdef object h_output
|
|
cdef object class_names |