app.py

import gradio as gr
from transformers import pipeline
from PIL import Image
import requests
import torch

# Verifique se uma GPU está disponível e configure o dispositivo
# Se você tiver um Hugging Face Space com GPU, use "cuda". Caso contrário, use "cpu".
device = "cuda" if torch.cuda.is_available() else "cpu"

# Inicializa o pipeline do modelo MedGemma
# O modelo google/medgemma-4b-it é uma variante de 4 bilhões de parâmetros
# O torch_dtype=torch.bfloat16 é recomendado para desempenho em GPUs mais recentes
# Certifique-se de que a versão do transformers seja >= 4.50.0
try:
    medgemma_pipe = pipeline(
        "image-text-to-text",
        model="google/medgemma-4b-it",
        torch_dtype=torch.bfloat16,
        device=device,
    )
except Exception as e:
    print(f"Erro ao carregar o modelo MedGemma: {e}")
    print("Verifique se você aceitou os termos de uso na página do modelo Hugging Face e se suas dependências estão corretas.")
    # Se o modelo não carregar, use uma função dummy para que a interface ainda funcione
    def medgemma_pipe(messages, max_new_tokens):
        print("Modelo não carregado, usando função dummy.")
        return [{"generated_text": [{"content": "Desculpe, o modelo não pôde ser carregado. Verifique os logs."}]}]

def process_medical_query(image_input, text_input):
    """
    Processa uma consulta médica que pode incluir uma imagem e/ou texto.
    """
    messages = []

    # O MedGemma é otimizado para aplicações médicas com geração de texto [cite: 113]
    # Se uma imagem for fornecida, adicione-a ao conteúdo da mensagem do usuário
    if image_input:
        # O modelo espera imagens normalizadas para 896x896 pixels [cite: 232]
        # Embora o pipeline possa lidar com isso automaticamente, é uma boa prática estar ciente
        image = image_input.convert("RGB") # Garante que a imagem esteja em formato RGB

        # A role 'system' pode ser usada para instruir o modelo, como ser um 'radiologista' [cite: 141, 142]
        messages.append({"role": "system", "content": [{"type": "text", "text": "Você é um assistente médico especialista."}]})
        
        # A role 'user' contém a pergunta e a imagem [cite: 145, 146, 147, 148, 149]
        user_content = []
        if text_input:
            user_content.append({"type": "text", "text": text_input})
        user_content.append({"type": "image", "image": image})
        messages.append({"role": "user", "content": user_content})

    elif text_input:
        messages.append({"role": "system", "content": [{"type": "text", "text": "Você é um assistente médico especialista."}]})
        messages.append({"role": "user", "content": [{"type": "text", "text": text_input}]})
    else:
        return "Por favor, forneça uma imagem e/ou uma pergunta."

    try:
        # A geração de texto é o principal objetivo do MedGemma [cite: 113]
        # max_new_tokens define o comprimento máximo da resposta gerada 
        output = medgemma_pipe(text=messages, max_new_tokens=500)
        
        # O resultado é uma lista de dicionários; precisamos extrair o texto gerado
        generated_text = output[0]["generated_text"][-1]["content"]
        return generated_text
    except Exception as e:
        return f"Ocorreu um erro durante a inferência: {e}. Verifique a entrada ou os logs do Space."

# Configuração da interface Gradio
# O MedGemma suporta entrada de texto e visão, e saída de texto [cite: 210, 211]
iface = gr.Interface(
    fn=process_medical_query,
    inputs=[
        gr.Image(type="pil", label="Envie uma Imagem Médica (Opcional)"),
        gr.Textbox(lines=5, placeholder="Descreva a imagem ou faça uma pergunta médica...", label="Sua Pergunta/Descrição"),
    ],
    outputs="text",
    title="⚕️ Assistente Médico com MedGemma 4B ⚕️",
    description=(
        "Este é um demonstrador interativo utilizando o modelo Google MedGemma 4B, "
        "treinado para compreensão de texto e imagem médica[cite: 91]. "
        "Envie uma imagem médica (como um raio-X) e/ou faça uma pergunta para obter uma análise. "
        "Lembre-se: As saídas geradas por modelos de IA são **preliminares** e **não devem ser usadas para diagnóstico ou tratamento clínico**[cite: 463, 464, 469]."
    ),
    live=False,
    allow_flagging="auto",
    examples=[
        [None, "Quais são os principais riscos de diabetes não controlada?"],
        ["https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest.jpg", "Descreva este raio-X do tórax."],
        ["https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest.jpg", "Há sinais de pneumonia nesta imagem?"],
    ]
)

# Inicia a interface Gradio
if __name__ == "__main__":
    iface.launch()