import tensorflow as tf


class ObjectDetector:
    def __init__(self, model_path='./model', label_file='./model/label.names',
                 num_classes=2, score_threshold=0.5, image_sz=(416, 416, 3)):
        self._model_path = model_path
        self._label_file = label_file
        self._num_classes = num_classes
        self._score_threshold = score_threshold
        self._image_sz = image_sz[0:2]

        self._config = ConfigProto()
        self._config.gpu_options.allow_growth = True

        self._graph = tf.Graph()

        with self._graph.as_default():
            self._sess = tf.Session(config=self._config)

            tf.saved_model.load(
                self._sess, [tag_constants.SERVING], self._model_path)

            self._image_tensor = self._sess.graph.get_tensor_by_name(
                'serving_default_input_1:0')
            self._output_tensor = self._sess.graph.get_tensor_by_name(
                'StatefulPartitionedCall:0')

