from image_prediction.estimator.estimator import Estimator


class EstimatorMock(Estimator):

    def __init__(self, estimator=lambda x: x):
        super().__init__(estimator=estimator)

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