from fastapi import APIRouter

from fastapi import File
import cv2
import numpy as np
import io
from scripts.core.handler.Text_Extraction_Handler import load_ocr_model,load_yolo_model, predict
from scripts.logging.application_logging import logger
from scripts.config.app_constants import TEXT_EXTRACT

ocr = load_ocr_model()
yolo_model = load_yolo_model()
router = APIRouter()


@router.post(TEXT_EXTRACT.text)
def extract_license_number(file: bytes = File(...)):
    try:
        stream = io.BytesIO(file)
        image = np.asarray(bytearray(stream.read()), dtype="uint8")
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)
        cv2.imwrite("test.jpg", image)
        image_path = "test.jpg"

        text = predict(image_path, ocr, yolo_model)
        logger.info("text detected")

        return text
    except Exception as e:
        logger.error("Error is :{}".format(e))

