src/services/ocr/replicate_ocr/replicate_ocr.py

"""
project @ convers-ai
created @ 2024-11-06
author  @ github.com/ishworrsubedii
"""
import base64

import cv2
import replicate


def replicate_run(input):
    output = replicate.run(
        "abiruyt/text-extract-ocr:a524caeaa23495bc9edc805ab08ab5fe943afd3febed884a4f3747aa32e9cd61",
        input=input
    )
    return output


class ReplicateOCR:
    def __init__(self):
        pass

    def read_text(self, image_path):
        _, buffer = cv2.imencode('.png', image_path)
        base_64_image = base64.b64encode(buffer).decode('utf-8')

        input_image_uri = f"data:image/png;base64,{base_64_image}"
        input = {
            "image": input_image_uri,
        }
        output = replicate_run(input)
        output = output.replace("\n", " ").replace("\t", " ")

        return output