Update app.py

app.py

@@ -24,16 +24,33 @@ def ordenar_e_preprocessar_imagem(img, modelo):
 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
+    
+    # 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)
+        image1 = normalizar_imagem(image1)
+        img_out = np.stack((image1,) * 3, axis=0)
+    
+    return img_out
 
 
 import os
@@ -42,35 +59,55 @@ import os
 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()
+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))
+        # 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
+        # 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
 
-    return f"{time.time() - start:.3f} segundos", img_out
 
 
 # Carregar modelos pré-treinados
@@ -90,12 +127,9 @@ import numpy as np
 import cv2
 from sklearn.cluster import KMeans
 
-# ...
-
-# Função para processar a imagem de entrada e calcular as características dos clusters
 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)  # Passe 'model' como 'pipe' aqui
+    resultado, img_out = processar_imagem_de_entrada(img, modelo, model)
     
     # Resto do código permanece inalterado
     kmeans = KMeans(n_clusters=2, random_state=0)
@@ -137,27 +171,24 @@ def processar_imagem_de_entrada_wrapper(img, modelo, salvar_resultado=False, nom
     
     return resultado, img_out, status_doenca, explanation, f"{num_analises} análises realizadas"
 
-# ...
-
-# Iniciar a interface Gradio
-my_app.launch()
-
 
 # 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'),
+        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 ou não"),
+        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.",
@@ -167,4 +198,4 @@ my_app = gr.Interface(
 )
 
 # Iniciar a interface Gradio
-my_app.launch()
+my_app.launch()