from image_prediction.estimator.adapter.adapter import EstimatorAdapter
from image_prediction.estimator.estimator import Estimator


class EstimatorAdapterPatch(EstimatorAdapter):
    def __init__(self, estimator: Estimator):
        super().__init__(estimator=estimator)
        self.__output_batch = None

    @property
    def output_batch(self):
        return self.__output_batch

    @output_batch.setter
    def output_batch(self, output_batch):
        self.__output_batch = output_batch

    def predict(self, batch):
        self.estimator.predict(batch)
        return self.__output_batch