from fastai.vision.all import *
import gradio as gr

# Carregar o modelo
learn = load_learner('export.pkl')

# Categorias em português
categorias = ('cortes_e_feridas', 'fratura', 'erupção_cutânea', 'fragmento')

def classificar_imagem(img):
    """Função para classificar a imagem e retornar as probabilidades"""
    pred, idx, probs = learn.predict(img)
    return {cat: float(prob) for cat, prob in zip(categorias, probs)}

# Exemplos de imagens para teste
exemplos = [
    'Cuts_for_nn.jpeg',
    'Fracture_examp.jpeg',
    'Rash.jpeg',
    'Splinter_examp.jpeg',
    'Splinter_download.jpeg',
    'Cut_download.jpeg',
    'Rash_download.jpeg',
    'Fracture_download.jpeg'
]

# Criando a interface moderna do Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏥 Classificador de Imagens Médicas
    
    Faça upload de uma imagem para classificar entre as seguintes categorias:
    - Cortes e Feridas
    - Fraturas
    - Erupções Cutâneas
    - Fragmentos
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Faça upload ou arraste uma imagem aqui",
                type="pil",
                height=300
            )
            upload_button = gr.Button("Classificar Imagem", variant="primary")
        
        with gr.Column():
            label_output = gr.Label(
                label="Resultados da Classificação",
                num_top_classes=4
            )
    
    # Galeria de exemplos
    gr.Examples(
        examples=exemplos,
        inputs=input_image,
        outputs=label_output,
        fn=classificar_imagem,
        cache_examples=True,
        label="Imagens de Exemplo"
    )
    
    # Configurar o evento de clique do botão
    upload_button.click(
        fn=classificar_imagem,
        inputs=input_image,
        outputs=label_output
    )
    
    gr.Markdown("""
    ### Como usar:
    1. Faça upload de uma imagem ou use um dos exemplos abaixo
    2. Clique no botão 'Classificar Imagem'
    3. Veja os resultados da classificação com as probabilidades
    
    ⚠️ Nota: Este é apenas um sistema de demonstração e não deve ser usado para diagnóstico médico.
    """)

# Iniciar a aplicação
demo.launch(share=True)