import gradio as gr
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import io
from pyceph.pyceph import predict

def process_image(image):
    """
    Обрабатывает изображение и возвращает результат с отмеченными точками
    """
    if image is None:
        return None, "Пожалуйста, загрузите изображение"
    
    # Сохраняем временное изображение
    temp_path = "temp_input.jpg"
    image.save(temp_path)
    
    # Получаем предсказания
    try:
        results = predict(image_src=temp_path, device='cpu')
        
        # Создаем изображение с точками
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.imshow(image)
        
        # Рисуем точки и подписи
        landmarks_text = []
        for landmark_dict in results:
            for name, (x, y) in landmark_dict.items():
                ax.plot(x, y, 'ro', markersize=8)
                ax.text(x, y, name, fontsize=8, color='yellow', 
                       bbox=dict(boxstyle='round,pad=0.3', facecolor='black', alpha=0.5))
                landmarks_text.append(f"{name}: ({x}, {y})")
        
        ax.axis('off')
        plt.tight_layout()
        
        # Конвертируем в изображение
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', dpi=150)
        buf.seek(0)
        result_image = Image.open(buf)
        plt.close()
        
        # Форматируем текст с координатами
        landmarks_output = "\n".join(landmarks_text)
        
        return result_image, landmarks_output
    
    except Exception as e:
        return None, f"Ошибка обработки: {str(e)}"

# Создаем интерфейс Gradio
with gr.Blocks(title="Py-Ceph: Цефалометрический анализ") as demo:
    gr.Markdown("""
    # 🦷 Py-Ceph: Определение цефалометрических точек
    
    Загрузите боковой рентгеновский снимок черепа (цефалограмму), 
    и модель автоматически определит 19 анатомических точек.
    
    Основано на [py-ceph](https://github.com/zhangted/py-ceph)
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Загрузите цефалограмму")
            submit_btn = gr.Button("Анализировать", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Результат с точками")
            output_text = gr.Textbox(label="Координаты точек", lines=20)
    
    submit_btn.click(
        fn=process_image,
        inputs=input_image,
        outputs=[output_image, output_text]
    )
    
    gr.Markdown("""
    ### Определяемые точки:
    Sella, Nasion, Orbitale, Porion, Subspinale, Supramentale, Pogonion, 
    Menton, Gnathion, Gonion, Incision Inferius, Incision Superius, 
    Upper Lip, Lower Lip, Subnasale, Soft Tissue Pogonion, 
    Posterior Nasal Spine, Anterior Nasal Spine, Articulare
    """)

if __name__ == "__main__":
    demo.launch()