Update app.py

app.py

@@ -1,20 +1,14 @@
 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
 
-# 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'),
-}
+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):
@@ -23,7 +17,43 @@ def ordenar_e_preprocessar_imagem(img, modelo):
     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))
+    
+    # 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
+
+# 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, salvar_resultado=False):
     try:
         # Faça uma cópia da imagem original
@@ -32,6 +62,9 @@ def processar_imagem_de_entrada(img, modelo, pipe, salvar_resultado=False):
         # 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)
         
@@ -44,6 +77,10 @@ def processar_imagem_de_entrada(img, modelo, pipe, salvar_resultado=False):
                 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))
@@ -60,27 +97,32 @@ def processar_imagem_de_entrada(img, modelo, pipe, salvar_resultado=False):
             nome_arquivo = os.path.join(caminho_salvar_resultado, f'resultado_{int(time.time())}.png')
             cv2.imwrite(nome_arquivo, img_out)
         
-        return img_out
+        return spent, img_out
     
     except Exception as e:
         # Em caso de erro, retorne uma mensagem de erro
         return 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'),
+}
 
-# Resto do código permanece inalterado
+# Função para processar a imagem de entrada com os campos extras removidos
 def processar_imagem_de_entrada_wrapper(img, modelo, salvar_resultado=False):
-    resultado, img_out, _, _, _ = processar_imagem_de_entrada(img, modelo, models[modelo], salvar_resultado)
-    
-    # Sobrepor a máscara na imagem original
-    overlaid_img = cv2.addWeighted(img, 0.7, cv2.cvtColor(img_out, cv2.COLOR_GRAY2BGR), 0.3, 0)
-    
-    return resultado, overlaid_img
-
+    model = models[modelo]
+    resultado = processar_imagem_de_entrada(img, modelo, model, salvar_resultado)
+    return resultado
 
 # Inicializar a contagem de análises
 num_analises = 0
 
-# Criar a interface Gradio
+# Criar a interface Gradio sem os campos extras
 my_app = gr.Interface(
     fn=processar_imagem_de_entrada_wrapper,
     inputs=[
@@ -93,8 +135,9 @@ my_app = gr.Interface(
     description="Esta aplicação segmenta angiogramas coronarianos usando modelos de segmentação pré-treinados.",
     theme="default",
     layout="vertical",
-    allow_flagging=False,
+    allow_flagging="never",  # Altere para uma string "auto", "manual" ou "never"
 )
 
 # Iniciar a interface Gradio
-my_app.launch()
+my_app.launch()
+