Update app.py

app.py

@@ -1,4 +1,6 @@
 import os
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
 import gradio as gr
 import torch
 import cv2
@@ -6,6 +8,17 @@ 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'):
@@ -39,20 +52,14 @@ def preprocessar_imagem(img, modelo='SE-RegUNet 4GF'):
     
     return img_out
 
-# # Função para realizar a inferência do modelo
-# def inferir_modelo(img, pipe, device):
-#     # Faça a inferência do modelo aqui
-#     # Certifique-se de pré-processar a imagem, colocá-la na GPU e retornar o resultado
-#     # ...
 
-# # Função para analisar os resultados
-# def analisar_resultados(img_out):
-#     # Realize a análise dos resultados aqui
-#     # Determine se há doença com base nas características dos clusters, etc.
-#     # ...
+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, device, salvar_resultado=False, nome_arquivo=None):
+def processar_imagem_de_entrada(img, modelo, pipe, salvar_resultado=False):
     try:
         # Faça uma cópia da imagem original
         img = img.copy()
@@ -64,10 +71,7 @@ def processar_imagem_de_entrada(img, modelo, pipe, device, salvar_resultado=Fals
         start = time.time()
         
         # Pré-processe a imagem e obtenha informações de dimensão
-        img = img.copy()
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-        h, w = img.shape
-        img = preprocessar_imagem(img, modelo)
+        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)
@@ -90,92 +94,115 @@ def processar_imagem_de_entrada(img, modelo, pipe, device, salvar_resultado=Fals
         logit = logit.astype(bool)
         
         # Crie uma cópia da imagem original para saída e aplique a máscara
-        img_out = img.copy()
+        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)
         
-        # Resto do código permanece inalterado
-        kmeans = KMeans(n_clusters=2, random_state=0)
-        flattened_img = img_out[:, :, 0].reshape((-1, 1))
-        kmeans.fit(flattened_img)
-        labels = kmeans.labels_
-        cluster_centers = kmeans.cluster_centers_
-            
-        # Extração de características dos clusters
-        num_clusters = len(cluster_centers)
-        cluster_features = []
-        for i in range(num_clusters):
-            cluster_mask = labels == i  # Create a boolean mask for the cluster
-        
-            # Calcular área do cluster
-            area = np.sum(cluster_mask)
-        
-            if area == 0:  # Skip empty clusters
-                continue
-        
-            # Calcular forma do cluster usando a relação entre área e perímetro
-            contours, _ = cv2.findContours(np.uint8(cluster_mask), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            if len(contours) > 0:
-                perimeter = cv2.arcLength(contours[0], True)
-                compactness = 4 * np.pi * area / (perimeter ** 2)
-        
-                cluster_features.append({'area': area, 'compactness': compactness})
-        
-        # Decidir se há doença com base nas características dos clusters
-        has_disease_flag = any(feature['area'] >= 200 and feature['compactness'] < 0.3 for feature in cluster_features)
-        
-        # Formatar o indicador de doença como uma string
-        if has_doenca_flag:
-            status_doenca = "Sim"
-            explanation = "A máquina detectou uma possível doença nos vasos sanguíneos."
-        else:
-            status_doenca = "Não"
-            explanation = "A máquina não detectou nenhuma doença nos vasos sanguíneos."
-        
-        return spent, img_out, status_doenca, explanation
+        return spent, img_out
     
     except Exception as e:
         # Em caso de erro, retorne uma mensagem de erro
         return str(e), None
 
-# Função para carregar um modelo pré-treinado com base no nome do modelo
-def carregar_modelo(modelo_nome):
-    return torch.jit.load(f'./model/{modelo_nome}.pt')
 
-# Função principal para iniciar o aplicativo
-def processar_imagem():
-    try:
-        # Configuração inicial
-        modelo_inicial = 'SE-RegUNet 4GF'
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        pipe = carregar_modelo(modelo_inicial)
-
-        # Criar a interface Gradio
-        my_app = gr.Interface(
-            fn=processar_imagem_de_entrada,
-            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=modelo_inicial),
-                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"),
-            ],
-            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()
 
-    except Exception as e:
-        print(str(e))
+# Carregar modelos pré-treinados
+models = {
+    'SE-RegUNet 4GF': torch.jit.load('./model/SERegUNet4GF.pt'),
+    'SE-RegUNet 16GF': torch.jit.load('./model/SERegUNet16GF.pt'),
+    'AngioNet': torch.jit.load('./model/AngioNet.pt'),
+    'EffUNet++ B5': torch.jit.load('./model/EffUNetppb5.pt'),
+    'Reg-SA-UNet++': torch.jit.load('./model/RegSAUnetpp.pt'),
+    'UNet3+': torch.jit.load('./model/UNet3plus.pt'),
+}
+
+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 processar_imagem_de_entrada_wrapper(img, modelo, salvar_resultado=False, nome_arquivo=None):
+    model = models[modelo]
+    resultado, img_out = processar_imagem_de_entrada(img, modelo, model)
+    
+    # Resto do código permanece inalterado
+    kmeans = KMeans(n_clusters=2, random_state=0)
+    flattened_img = img_out[:, :, 0].reshape((-1, 1))
+    kmeans.fit(flattened_img)
+    labels = kmeans.labels_
+    cluster_centers = kmeans.cluster_centers_
+        
+    # Extração de características dos clusters
+    num_clusters = len(cluster_centers)
+    cluster_features = []
+    for i in range(num_clusters):
+        cluster_mask = labels == i  # Create a boolean mask for the cluster
+    
+        # Calcular área do cluster
+        area = np.sum(cluster_mask)
+    
+        if area == 0:  # Skip empty clusters
+            continue
+    
+        # Calcular forma do cluster usando a relação entre área e perímetro
+        contours, _ = cv2.findContours(np.uint8(cluster_mask), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        if len(contours) > 0:
+            perimeter = cv2.arcLength(contours[0], True)
+            compactness = 4 * np.pi * area / (perimeter ** 2)
+    
+            cluster_features.append({'area': area, 'compactness': compactness})
+    
+    # Decidir se há doença com base nas características dos clusters
+    has_disease_flag = any(feature['area'] >= 200 and feature['compactness'] < 0.3 for feature in cluster_features)
+    
+    # Formatar o indicador de doença como uma string
+    if has_disease_flag:
+        status_doenca = "Sim"
+        explanation = "A máquina detectou uma possível doença nos vasos sanguíneos."
+    else:
+        status_doenca = "Não"
+        explanation = "A máquina não detectou nenhuma doença nos vasos sanguíneos."
+    
+    return resultado, img_out, status_doenca, explanation, f"{num_analises} análises realizadas"
+
+
+# Inicializar a contagem de análises
+num_analises = 0
+
+# 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'),
+        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 a interface Gradio
+my_app.launch() otimize
 
-# Chamar a função principal para iniciar o aplicativo
-processar_imagem()
+# Criar uma pasta para salvar os resultados
+if not os.path.exists(caminho_salvar_resultado):
+    os.makedirs(caminho_salvar_resultado)
 
+# Iniciar o servidor
+my_app.launch()