Update app.py

app.py

@@ -1,131 +1,85 @@
 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 glob
 import time
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print(
-    "torch: ", torch.__version__,
-)
-
-def ordenar_arquivos(img, modelo):
-    ori = img.copy()
+def preprocess_image(img, model_name):
+    # Preprocess the input image based on the selected model
     img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    h, w = img.shape
-    img_out = preprocessamento(img, modelo)
-    return img_out, h, w, img, ori
-
-def preprocessamento(img, modelo='SE-RegUNet 4GF'):
     img = cv2.resize(img, (512, 512))
-    img = unsharp_masking(img).astype(np.uint8)
-    if modelo == 'AngioNet' or modelo == 'UNet3+':
-        img = np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
-        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 = np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
-        image1 = np.float32((image1 - image1.min()) / (image1.max() - image1.min() + 1e-6))
-        image2 = np.float32((image2 - image2.min()) / (image2.max() - image2.min() + 1e-6))
-        img_out = np.stack((img, image1, image2), axis=0)
-    else:
-        clahe1 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-        image1 = clahe1.apply(img)
-        image1 = np.float32((image1 - image1.min()) / (image1.max() - image1.min() + 1e-6))
-        img_out = np.stack((image1,) * 3, axis=0)
-    return img_out
-
-def processar_imagem_de_entrada(img, modelo, pipe):
-    img = img.copy()
-    pipe = pipe.to(device).eval()
-    start = time.time()
-    img, h, w, ori_gray, ori = ordenar_arquivos(img, modelo)
-    img = torch.FloatTensor(img).unsqueeze(0).to(device)
-    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)
-    spent = time.time() - start
-    spent = f"{spent:.3f} segundos"
-
-    if h != 512 or w != 512:
-        logit = cv2.resize(logit, (h, w))
-
-    logit = logit.astype(bool)
-    img_out = ori.copy()
-    img_out[logit, 0] = 255
-    return spent, img_out
+    img = cv2.GaussianBlur(img, (0, 0), 1.0)
+    img = img.astype(np.float32) / 255.0
+    return img
 
-# Load the models outside the function
-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'),
-}
+def load_model(model_path):
+    model = torch.jit.load(model_path)
+    return model
 
-def processar_imagem_de_entrada_wrapper(img, modelo):
-    model = models[modelo]
-    spent, img_out = processar_imagem_de_entrada(img, modelo, model)
-
-    # Define the function `has_disease`
-    def has_disease(img_out):
-        """
-        Checks if the angiogram has disease based on the segmentation.
-        Args:
-            img_out: The segmented angiogram.
-        Returns:
-            True if the angiogram has disease, False otherwise.
-        """
-        percentage_of_vessels = np.sum(img_out) / (img_out.shape[0] * img_out.shape[1])
-        if percentage_of_vessels > 0.5:
-            return True
-        else:
-            return False
+def process_image(img, model):
+    img = np.expand_dims(img, axis=0)
+    img_tensor = torch.FloatTensor(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logit = torch.softmax(model.forward(img_tensor), dim=1).detach().cpu().numpy()[0, 0]
+    return (logit > 0.5).astype(np.uint8)
 
-    has_disease = has_disease(img_out)
+def has_disease(segmented_img):
+    percentage_of_vessels = np.sum(segmented_img) / (segmented_img.shape[0] * segmented_img.shape[1])
+    return percentage_of_vessels > 0.5
 
-    # Add the explanation to the interface
-    if has_disease:
-        explanation = (
-            f"* **True:** O angiograma tem doença. Isso é determinado pela função `has_disease`, que calcula o percentual de vasos no angiograma segmentado. Se o percentual de vasos for maior que 50%, a função retorna \"true\", indicando que o angiograma tem doença. Caso contrário, a função retorna \"false\", indicando que o angiograma não tem doença."
+def explanation(has_disease_flag):
+    if has_disease_flag:
+        explanation_text = (
+            "O angiograma tem doença. Isso é determinado pela função `has_disease`, que calcula o percentual de vasos no angiograma segmentado. Se o percentual de vasos for maior que 50%, a função retorna \"true\", indicando que o angiograma tem doença. Caso contrário, a função retorna \"false\", indicando que o angiograma não tem doença."
         )
     else:
-        explanation = (
-            f"* **False:** O angiograma não tem doença."
-        )
-
-    return spent, img_out, has_disease, explanation
-
-
-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'),
-    ],
-    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. 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.",
-    theme="default",
-    layout="vertical",
-    allow_flagging=False,
-)
-
-my_app.launch()
+        explanation_text = "O angiograma não tem doença."
+    return explanation_text
+
+def process_input_image(input_img, model_name, model):
+    start_time = time.time()
+    
+    preprocessed_img = preprocess_image(input_img, model_name)
+    processed_img = process_image(preprocessed_img, model)
+    disease_flag = has_disease(processed_img)
+    explanation_text = explanation(disease_flag)
+    
+    elapsed_time = f"{time.time() - start_time:.3f} segundos"
+    return elapsed_time, processed_img, disease_flag, explanation_text
+
+def main():
+    # Load models
+    model_paths = {
+        'SE-RegUNet 4GF': './model/SERegUNet4GF.pt',
+        # ... (other model paths here)
+    }
+    models = {model_name: load_model(model_path) for model_name, model_path in model_paths.items()}
+    
+    # Create Gradio interface
+    gr_interface = gr.Interface(
+        fn=process_input_image,
+        inputs=[
+            gr.inputs.Image(label="Angiograma:", shape=(512, 512)),
+            gr.inputs.Dropdown(list(model_paths.keys()), label='Modelo', default='SE-RegUNet 4GF'),
+        ],
+        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. 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.",
+        theme="default",
+        layout="vertical",
+        allow_flagging=False,
+    )
+    
+    # Launch Gradio interface
+    gr_interface.launch()
+
+if __name__ == "__main__":
+    main()