Update app.py

app.py

@@ -6,15 +6,11 @@ import torch
 import cv2
 import numpy as np
 from preprocess import unsharp_masking
-import glob
 import time
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print(
-    "torch: ", torch.__version__,
-)
-
+# Função para ordenar e pré-processar a imagem de entrada
 def ordenar_arquivos(img, modelo):
     ori = img.copy()
     img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
@@ -22,6 +18,7 @@ def ordenar_arquivos(img, modelo):
     img_out = preprocessamento(img, modelo)
     return img_out, h, w, img, ori
 
+# Função para pré-processar a imagem com base no modelo selecionado
 def preprocessamento(img, modelo='SE-RegUNet 4GF'):
     img = cv2.resize(img, (512, 512))
     img = unsharp_masking(img).astype(np.uint8)
@@ -44,6 +41,7 @@ def preprocessamento(img, modelo='SE-RegUNet 4GF'):
         img_out = np.stack((image1,) * 3, axis=0)
     return img_out
 
+# Função para processar a imagem de entrada
 def processar_imagem_de_entrada(img, modelo, pipe):
     img = img.copy()
     pipe = pipe.to(device).eval()
@@ -65,7 +63,7 @@ def processar_imagem_de_entrada(img, modelo, pipe):
     img_out[logit, 0] = 255
     return spent, img_out
 
-# Load the models outside the function
+# Carregar modelos pré-treinados
 models = {
     'SE-RegUNet 4GF': torch.jit.load('./model/SERegUNet4GF.pt'),
     'SE-RegUNet 16GF': torch.jit.load('./model/SERegUNet16GF.pt'),
@@ -75,46 +73,48 @@ models = {
     'UNet3+': torch.jit.load('./model/UNet3plus.pt'),
 }
 
-def processar_imagem_de_entrada_wrapper(img, modelo):
-    model = models[modelo]
-    spent, img_out = processar_imagem_de_entrada(img, modelo, model)
+from sklearn.cluster import KMeans
 
-    def has_disease(img, model, min_area=500):
-      """Checks if the angiogram has disease.
-    
-      Args:
-        img: The input angiogram.
-        model: The segmentation model.
-        min_area: The minimum area for an anomaly to be considered a disease.
-    
-      Returns:
-        True if the angiogram has disease, False otherwise.
-      """
+def has_disease(img, min_area=500):
+    # Aplicar K-Means na imagem segmentada ou máscara
+    # Certifique-se de fornecer a imagem ou máscara correta
+    # A imagem ou máscara deve ser preprocessada e segmentada antes desta etapa
     
-      # Segment the angiogram.
-      mask = model.predict(img)
+    kmeans = KMeans(n_clusters=2, random_state=0)
+    flattened_img = img.reshape((-1, 1))  # Transformar a imagem em uma matriz unidimensional
+    kmeans.fit(flattened_img)
     
-      # Detect anomalies in the mask.
-      anomalies = detect_anomalies(img, mask)
+    # Rotular as regiões da imagem
+    labels = kmeans.labels_
     
-      # Check if any of the anomalies are large enough to be considered a disease.
-      for anomaly in anomalies:
-        if anomaly[1] >= min_area:
-          return True
+    # Calcular a área das regiões
+    area_0 = np.sum(labels == 0)
+    area_1 = np.sum(labels == 1)
     
-      return False
+    # Verificar se a área da região com doença é maior que o limite mínimo
+    if area_1 >= min_area:
+        return True
+    else:
+        return False
 
-    has_disease = has_disease(img_out)
+
+# Função que encapsula o processamento e a verificação de doenças
+def processar_imagem_de_entrada_wrapper(img, modelo):
+    model = models[modelo]
+    spent, img_out = processar_imagem_de_entrada(img, modelo, model)
+
+    # Chame a função has_disease para verificar se há doença
+    has_disease_flag = has_disease(img_out)
 
     # Adicione a explicação à interface
-    if has_disease:
-        explanation = "A máquina detectou uma possível doença nos vasos sanguíneos. Ela olha para a parte destacada da imagem e calcula se há mais vasos do que o normal. Se for o caso, ela diz que há uma doença. Caso contrário, ela diz que não há doença."
+    if has_disease_flag:
+        explanation = "A máquina detectou uma possível doença nos vasos sanguíneos."
     else:
         explanation = "A máquina não detectou nenhuma doença nos vasos sanguíneos."
     
-    return spent, img_out, has_disease, explanation
-
+    return spent, img_out, has_disease_flag, explanation
 
+# Criar a interface Gradio
 my_app = gr.Interface(
     fn=processar_imagem_de_entrada_wrapper,
     inputs=[
@@ -128,10 +128,11 @@ my_app = gr.Interface(
         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. Faça o upload de uma imagem de angiograma e selecione um modelo para visualizar o resultado da segmentação.\n\nSelecione uma imagem de angiograma coronariano e um modelo de segmentação no painel à esquerda.",
+    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()