autolabelling-models/easy_ocr.py

import dagshub
import mlflow

conda_env = {'channels': ['defaults'],
             'dependencies': [
                 'python~=3.12',
                 'pip',
                 {'pip': ['mlflow',
                          'easyocr']}],
             'name': 'env'}



class EasyOCR(mlflow.pyfunc.PythonModel):
    def __init__(self):
        super().__init__()
        import easyocr

        self.easyocr = easyocr

    def load_context(self, context):
        self.model = self.easyocr.Reader(['mn','en'])

    def predict(self, context, model_input):
        if type(model_input) == list:
            res = []
            for filepath in model_input:
                res.append(self.model.readtext(filepath))
            return res
        else: return self.model.readtext(model_input)

if __name__ == '__main__':
    dagshub.init('autolabelling-models', 'jinensetpal')
    model = EasyOCR()

    with mlflow.start_run():
        mlflow.pyfunc.log_model(
            artifact_path='models',
            python_model=model,
            conda_env=conda_env,
            registered_model_name='easy_ocr')