import os
import logging
from PIL import Image
import torch
import gradio as gr
from surya.foundation import FoundationPredictor
from surya.recognition import RecognitionPredictor
from surya.detection import DetectionPredictor

# Настройка логирования
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Настройка переменных окружения для ускорения
logger.info("Настройка переменных окружения для ускорения обработки")
os.environ["RECOGNITION_BATCH_SIZE"] = "512"
os.environ["DETECTOR_BATCH_SIZE"] = "36"
os.environ["RECOGNITION_STATIC_CACHE"] = "true"

# Включение TorchDynamo
torch._dynamo.config.capture_scalar_outputs = True

# Инициализация предикторов
logger.info("Инициализация предикторов")
foundation_predictor = FoundationPredictor()
recognition_predictor = RecognitionPredictor(foundation_predictor)
detection_predictor = DetectionPredictor()

# Компиляция модели распознавания для ускорения
try:
    logger.info("Компиляция модели распознавания")
    recognition_predictor.model.decoder.model = torch.compile(recognition_predictor.model.decoder.model)
except Exception as e:
    logger.warning(f"Не удалось скомпилировать модель распознавания: {e}. Продолжаем без компиляции.")

def recognize_text(image: Image.Image) -> str:
    """
    Функция распознавания текста с изображения.
    Использует детектор для поиска текстовых областей и предиктор распознавания для OCR.
    """
    try:
        # Принудительно конвертируем изображение в RGB, если нужно
        if image.mode != "RGB":
            image = image.convert("RGB")
        predictions = recognition_predictor([image], det_predictor=detection_predictor)
        result = "\n".join(line.text for prediction in predictions for line in prediction.text_lines)
        return result.strip()
    except Exception as e:
        logger.error(f"Ошибка при распознавании текста: {e}")
        return f"Ошибка: {e}"

iface = gr.Interface(
    fn=recognize_text,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Распознанный текст"),
    title="Surya OCR",
    description="Загрузите изображение, чтобы распознать текст с помощью Surya."
)

if __name__ == "__main__":
    logger.info("Запуск приложения Gradio")
    iface.launch()