Update app.py

app.py

@@ -1,6 +1,4 @@
 import os
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-
 import gradio as gr
 import torch
 import cv2
@@ -8,109 +6,68 @@ 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
+def preprocess_image(img, model='SE-RegUNet 4GF'):
     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):
+
+    def normalize_image(img):
         return np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
-    
-    if modelo == 'AngioNet' or modelo == 'UNet3+':
-        img = normalizar_imagem(img)
+
+    if model in ('AngioNet', 'UNet3+'):
+        img = normalize_image(img)
         img_out = np.expand_dims(img, axis=0)
-    elif modelo == 'SE-RegUNet 4GF':
+    elif model == '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 = normalize_image(img)
+        image1 = normalize_image(image1)
+        image2 = normalize_image(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)
+        image1 = normalize_image(image1)
         img_out = np.stack((image1,) * 3, axis=0)
-    
-    return img_out
 
+    return img_out
 
-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, salvar_resultado=False):
+def process_input_image(img, model, save_result=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
+        pipe = models[model].to(device).eval()
         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, h, w, ori_gray, ori = preprocess_image(img, model)
         img = torch.FloatTensor(img).unsqueeze(0).to(device)
-        
-        # Realize a inferência do modelo sem gradientes
+
         with torch.no_grad():
-            if modelo == 'AngioNet':
+            if model == '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
+        spent = f"{time.time() - start:.3f} segundos"
+
         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)
-        
+        if save_result:
+            file_name = os.path.join(caminho_salvar_resultado, f'resultado_{int(time.time())}.png')
+            cv2.imwrite(file_name, img_out)
+
         return spent, img_out
-    
+
     except Exception as e:
-        # Em caso de erro, retorne uma mensagem de erro
         return str(e), None
 
-
-
-# Carregar modelos pré-treinados
 models = {
     'SE-RegUNet 4GF': torch.jit.load('./model/SERegUNet4GF.pt'),
     'SE-RegUNet 16GF': torch.jit.load('./model/SERegUNet16GF.pt'),
@@ -120,72 +77,45 @@ models = {
     'UNet3+': torch.jit.load('./model/UNet3plus.pt'),
 }
 
-from scipy.spatial import distance
-from scipy.ndimage import label
-import numpy as np
+def process_input_image_wrapper(img, model, save_result=False):
+    resultado, img_out = process_input_image(img, model, save_result)
 
-# Adicionar a opção de salvar o resultado em um arquivo
-def processar_imagem_de_entrada_wrapper(img, modelo, salvar_resultado=False):
-    model = models[modelo]
-    resultado, img_out = processar_imagem_de_entrada(img, modelo, model, salvar_resultado)
-    
-    # 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_
-    
-    # Detecção de doenças usando K-Means
-    kmeans = KMeans(n_clusters=2, random_state=0)
-    flattened_img = img_out[:, :, 0].reshape((-1, 1))  # Use o canal de intensidade
-    kmeans.fit(flattened_img)
-    labels = kmeans.labels_
-    cluster_centers = kmeans.cluster_centers_
-    
-    # Resto do código permanece inalterado
-    
-    # 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
+        cluster_mask = labels == i
         area = np.sum(cluster_mask)
-    
-        if area == 0:  # Skip empty clusters
+
+        if area == 0:
             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."
-    
-    # Resto do código permanece inalterado
-    
+
+    status_doenca = "Sim" if has_disease_flag else "Não"
+    explanation = "A máquina detectou uma possível doença nos vasos sanguíneos." if has_disease_flag else "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
+caminho_salvar_resultado = "/Segmento_de_Angio_Coronariana_v5/Salvar Resultado"
 num_analises = 0
 
-# Criar a interface Gradio
 my_app = gr.Interface(
-    fn=processar_imagem_de_entrada_wrapper,
+    fn=process_input_image_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'),
@@ -205,5 +135,4 @@ my_app = gr.Interface(
     allow_flagging=False,
 )
 
-# Iniciar a interface Gradio
-my_app.launch()
+my_app.launch()