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