from edge_engine.ai.model.modelwraper import ModelWrapper
import cv2
import base64


class LoopBackModel(ModelWrapper):

    def __init__(self, pubs, path=None, ):
        super().__init__(path)
        self.mqtt = pubs.mqtt_pub

    def _pre_process(self, x):
        return x

    def _post_process(self, x):
        image = cv2.imencode('.jpg', x['frame'])[1].tostring()
        image = 'data:image/jpeg;base64,' + base64.b64encode(image).decode("utf-8")
        x['frame'] = image
        self.mqtt.publish(x)
        return x

    def _predict(self, x):
        return x

    def predict(self, x):
        return super().predict(x)
