import gradio as gr
from transformers import pipeline, AutoModelForImageClassification, AutoImageProcessor
import torch
from PIL import Image
import numpy as np

WOUND_TYPES = {
    "pressure_ulcer": "Úlcera de Pressão",
    "venous_ulcer": "Úlcera Venosa",
    "diabetic_ulcer": "Úlcera Diabética",
    "surgical_wound": "Ferida Cirúrgica",
    "burn": "Queimadura",
    "healthy_skin": "Pele Saudável"
}

def load_models():
    # Carregando modelos específicos para feridas
    wound_classifier = pipeline(
        "image-classification",
        model="stevhliu/wound-classification",  # Modelo específico para classificação de feridas
        device=0 if torch.cuda.is_available() else -1
    )
    
    tissue_classifier = pipeline(
        "image-classification",
        model="viktorcikojevic/wound-tissue-type",  # Modelo para classificação do tipo de tecido
        device=0 if torch.cuda.is_available() else -1
    )
    
    return wound_classifier, tissue_classifier

def preprocess_image(image):
    # Normalização e pré-processamento da imagem
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert('RGB')
    return image

def classify_wound(image):
    if image is None:
        return None, None
    
    wound_classifier, tissue_classifier = load_models()
    processed_image = preprocess_image(image)
    
    # Classificação do tipo de ferida
    wound_results = wound_classifier(processed_image)
    
    # Classificação do tipo de tecido
    tissue_results = tissue_classifier(processed_image)
    
    # Formatando resultados da classificação de feridas
    wound_formatted = []
    for result in wound_results:
        label = WOUND_TYPES.get(result['label'], result['label'])
        score = result['score']
        wound_formatted.append((label, score))
    
    # Formatando resultados da classificação de tecidos
    tissue_formatted = []
    for result in tissue_results:
        label = result['label'].replace('_', ' ').title()
        score = result['score']
        tissue_formatted.append((label, score))
    
    return wound_formatted, tissue_formatted

# Criando a interface do Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏥 Classificador Especializado de Feridas
    
    Sistema de classificação utilizando modelos especializados para análise de feridas e tecidos.
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload da Imagem",
                type="pil",
                tool="select"
            )
            submit_btn = gr.Button("Analisar Ferida", variant="primary", size="lg")
        
        with gr.Column():
            wound_output = gr.Label(
                label="Tipo de Ferida",
                num_top_classes=3
            )
            tissue_output = gr.Label(
                label="Tipo de Tecido",
                num_top_classes=3
            )
    
    with gr.Row():
        with gr.Accordion("Informações Importantes", open=False):
            gr.Markdown("""
            ### Tipos de Feridas Detectáveis:
            - Úlcera de Pressão
            - Úlcera Venosa
            - Úlcera Diabética
            - Ferida Cirúrgica
            - Queimadura
            
            ### Tipos de Tecidos Analisados:
            - Tecido de Granulação
            - Tecido Necrótico
            - Tecido Fibrinoso
            - Tecido Epitelial
            
            ### Recomendações para Melhores Resultados:
            1. Use imagens bem iluminadas
            2. Capture a imagem em um ângulo perpendicular à ferida
            3. Mantenha um fundo neutro e limpo
            4. Evite sombras ou reflexos excessivos
            """)
    
    # Configurando eventos
    submit_btn.click(
        fn=classify_wound,
        inputs=input_image,
        outputs=[wound_output, tissue_output]
    )
    
    # Exemplos
    gr.Examples(
        examples=[
            ["exemplo_pressure_ulcer.jpg"],
            ["exemplo_diabetic_wound.jpg"],
        ],
        inputs=input_image,
        outputs=[wound_output, tissue_output],
        cache_examples=True
    )

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