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)
    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 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

    @staticmethod
    cdef bytes convert_from_onnx(bytes onnx_model)

    @staticmethod
    cdef unsigned long long get_gpu_memory_bytes(device_id=?)

    @staticmethod
    cdef str get_engine_filename(device_id=?)