Update app.py

app.py

@@ -1,13 +1,78 @@
 import os
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
 import gradio as gr
 import torch
 import cv2
 import numpy as np
 from preprocess import unsharp_masking
+import time
+from sklearn.cluster import KMeans
+import os
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Função para ordenar e pré-processar a imagem de entrada
+def ordenar_e_preprocessar_imagem(img, modelo):
+    ori = img.copy()
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    h, w = img.shape
+    img_out = preprocessar_imagem(img, modelo)
+    return img_out, h, w, img, ori
+
+# Função para pré-processar a imagem com base no modelo selecionado
+def preprocessar_imagem(img, modelo='SE-RegUNet 4GF'):
+    # Redimensionar a imagem para 512x512
+    img = cv2.resize(img, (512, 512))
+
+    # Normalizar a imagem
+    img = img.astype(np.float32)
+    img = (img - img.min()) / (img.max() - img.min() + 1e-6)
+
+    # Aplicar o CLAHE
+    if modelo == 'SE-RegUNet 4GF':
+        img = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(img)
+
+    return img
+
+
+import os
 
 # Caminho absoluto para a pasta de salvamento
 caminho_salvar_resultado = "/Segmento_de_Angio_Coronariana_v5/Salvar Resultado"
 
+# Função para processar a imagem de entrada
+def processar_imagem_de_entrada(img, modelo, pipe):
+    """
+    Processa uma imagem de entrada usando um modelo de segmentação pré-treinado.
+
+    Args:
+        img: Imagem de entrada (numpy array).
+        modelo: Modelo de segmentação.
+        pipe: Pipe de inferência do modelo.
+
+    Returns:
+        Tempo decorrido e imagem de saída.
+    """
+
+    # Pré-processe a imagem
+    img, h, w = ordenar_e_preprocessar_imagem(img, modelo)
+
+    # Execute a inferência do modelo
+    with torch.no_grad():
+        logit = pipe(img.unsqueeze(0).to(device))[0].argmax(dim=1).cpu().numpy()
+
+    # Redimensione o resultado, se necessário
+    if h != 512 or w != 512:
+        logit = cv2.resize(logit, (h, w))
+
+    # Crie a imagem de saída
+    img_out = img.copy()
+    img_out[logit] = 255
+
+    return f"{time.time() - start:.3f} segundos", img_out
+
+
 # Carregar modelos pré-treinados
 models = {
     'SE-RegUNet 4GF': torch.jit.load('./model/SERegUNet4GF.pt'),
@@ -18,120 +83,54 @@ models = {
     'UNet3+': torch.jit.load('./model/UNet3plus.pt'),
 }
 
-def ordenar_e_preprocessar_imagem(img, modelo):
-    ori = img.copy()
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    h, w = img.shape
-    img_out = preprocessar_imagem(img, modelo)
-    return img_out, h, w, img, ori
+from scipy.spatial import distance
+from scipy.ndimage import label
+import numpy as np
+import numpy as np
+import cv2
+from sklearn.cluster import KMeans
 
-def preprocessar_imagem(img, modelo='SE-RegUNet 4GF'):
-    # Redimensionar a imagem para 512x512
-    img = cv2.resize(img, (512, 512))
-    
-    # Aplicar a máscara de nitidez à imagem
-    img = unsharp_masking(img).astype(np.uint8)
-    
-    # Função auxiliar para normalizar a imagem
-    def normalizar_imagem(img):
-        return np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
-    
-    if modelo == 'AngioNet' or modelo == 'UNet3+':
-        img = normalizar_imagem(img)
-        img_out = np.expand_dims(img, axis=0)
-    elif modelo == 'SE-RegUNet 4GF':
-        clahe1 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-        clahe2 = cv2.createCLAHE(clipLimit=8.0, tileGridSize=(8, 8))
-        image1 = clahe1.apply(img)
-        image2 = clahe2.apply(img)
-        img = normalizar_imagem(img)
-        image1 = normalizar_imagem(image1)
-        image2 = normalizar_imagem(image2)
-        img_out = np.stack((img, image1, image2), axis=0)
-    else:
-        clahe1 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-        image1 = clahe1.apply(img)
-
-def processar_imagem_de_entrada(img, modelo, pipe, salvar_resultado=False):
-    try:
-        # Faça uma cópia da imagem original
-        img = img.copy()
-        
-        # Coloque o modelo na GPU (se disponível) e configure-o para modo de avaliação
-        pipe = pipe.to(device).eval()
-        
-        # Registre o tempo de início
-        start = time.time()
-        
-        # Pré-processe a imagem e obtenha informações de dimensão
-        img, h, w, ori_gray, ori = ordenar_e_preprocessar_imagem(img, modelo)
-        
-        # Converta a imagem para o formato esperado pelo modelo e coloque-a na GPU
-        img = torch.FloatTensor(img).unsqueeze(0).to(device)
-        
-        # Realize a inferência do modelo sem gradientes
-        with torch.no_grad():
-            if modelo == 'AngioNet':
-                img = torch.cat([img, img], dim=0)
-            logit = np.round(torch.softmax(pipe.forward(img), dim=1).detach().cpu().numpy()[0, 0]).astype(np.uint8)
-        
-        # Calcule o tempo decorrido
-        spent = time.time() - start
-        spent = f"{spent:.3f} segundos"
-
-        # Redimensione o resultado, se necessário
-        if h != 512 or w != 512:
-            logit = cv2.resize(logit, (h, w))
-
-        # Converta o resultado para um formato booleano
-        logit = logit.astype(bool)
-        
-        # Crie uma cópia da imagem original para saída e aplique a máscara
-        img_out = ori.copy()
-        img_out[logit, 0] = 255
-
-        # Salve o resultado em um arquivo se a opção estiver ativada
-        if salvar_resultado:
-            nome_arquivo = os.path.join(caminho_salvar_resultado, f'resultado_{int(time.time())}.png')
-            cv2.imwrite(nome_arquivo, img_out)
-        
-        return spent, img_out
-    
-    except Exception as e:
-        # Em caso de erro, retorne uma mensagem de erro
-        return str(e), None
-
-def main():
-    num_analises = 0
-
-    # Criar interface Gradio
-    my_app = gr.Interface(
-        fn=processar_imagem_de_entrada_wrapper,
-        inputs=[
-            gr.inputs.Image(label="Angiograma:", shape=(512, 512)),
-            gr.inputs.Dropdown(['SE-RegUNet 4GF','SE-RegUNet 16GF', 'AngioNet', 'EffUNet++ B5', 'Reg-SA-UNet++', 'UNet3+'], label='Modelo', default='SE-RegUNet 4GF'),
-            gr.inputs.Checkbox(label="Salvar Resultado"),
-        ],
-        outputs=[
-            gr.outputs.Label(label="Tempo decorrido"),
-            gr.outputs.Image(type="numpy", label="Imagem de Saída"),
-            gr.outputs.Label(label="Possui Doença?"),
-            gr.outputs.Label(label="Explicação"),
-            gr.outputs.Label(label="Análises Realizadas"),
-        ],
-        title="Segmentação de Angiograma Coronariano",
-        description="Esta aplicação segmenta angiogramas coronarianos usando modelos de segmentação pré-treinados.",
-        theme="default",
-        layout="vertical",
-        allow_flagging=False,
-    )
-
-    # Iniciar interface Gradio
-    my_app.launch() 
-    
-    # Criar uma pasta para salvar os resultados
-    if not os.path.exists(caminho_salvar_resultado):
-        os.makedirs(caminho_salvar_resultado)
-
-if __name__ == "__main__":
-    main()
+def processar_imagem_de_entrada_wrapper(img, modelo, salvar_resultado=False):
+    model = models[modelo]
+    resultado, img_out = processar_imagem_de_entrada(img, model)
+
+    # Extração de características dos clusters
+    kmeans = KMeans(n_clusters=2, random_state=0)
+    kmeans.fit(img_out[:, :, 0].reshape((-1, 1)))
+    labels = kmeans.labels_
+    cluster_centers = kmeans.cluster_centers_
+
+    # Detecção de doença
+    has_disease = any(cluster_centers[i][0] < 0.3 for i in range(2))
+
+    # Resultado
+    status_doenca = "Sim" if has_disease else "Não"
+    explanation = "A máquina detectou uma possível doença nos vasos sanguíneos." if has_disease else "A máquina não detectou nenhuma doença nos vasos sanguíneos."
+
+    return resultado, img_out, status_doenca, explanation
+
+# Inicializar a contagem de análises
+num_analises = 0
+
+# Criar a interface Gradio
+# Criar a interface Gradio
+my_app = gr.Interface(
+    fn=processar_imagem_de_entrada_wrapper,
+    inputs=[
+        gr.inputs.Image(label="Angiograma:", shape=(512, 512)),
+        gr.inputs.Dropdown(['SE-RegUNet 4GF','SE-RegUNet 16GF', 'AngioNet', 'EffUNet++ B5', 'Reg-SA-UNet++', 'UNet3+'], label='Modelo', default='SE-RegUNet 4GF'),
+    ],
+    outputs=[
+        gr.outputs.Label(label="Tempo decorrido"),
+        gr.outputs.Image(type="numpy", label="Imagem de Saída"),
+        gr.outputs.Label(label="Possui Doença ou não"),
+    ],
+    title="Segmentação de Angiograma Coronariano",
+    description="Esta aplicação segmenta angiogramas coronarianos usando modelos de segmentação pré-treinados.",
+    theme="default",
+    layout="vertical",
+    allow_flagging=False,
+)
+
+# Iniciar a interface Gradio
+my_app.launch()