Update app.py

app.py

@@ -1,7 +1,3 @@
-# preprocess.py
-# (No changes in this file)
-
-# main.py
 import os
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 
@@ -9,31 +5,22 @@ import gradio as gr
 import torch
 import cv2
 import numpy as np
-import time
-from gradio.components import Label, Image, Dropdown
-
-# Add the missing 'unsharp_masking' function from preprocess.py
 from preprocess import unsharp_masking
-
-# 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'),
-}
+import glob
+import time
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def preprocess_image(img, modelo):
-    img = img.copy()
-    if len(img.shape) == 3 and img.shape[-1] == 3:
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+print(
+    "torch: ", torch.__version__,
+)
+
+def ordenar_arquivos(img, modelo):
+    ori = img.copy()
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     h, w = img.shape
     img_out = preprocessamento(img, modelo)
-    return img_out, h, w, img
+    return img_out, h, w, img, ori
 
 def preprocessamento(img, modelo='SE-RegUNet 4GF'):
     img = cv2.resize(img, (512, 512))
@@ -57,11 +44,11 @@ def preprocessamento(img, modelo='SE-RegUNet 4GF'):
         img_out = np.stack((image1,) * 3, axis=0)
     return img_out
 
-def process_image(img, modelo):
-    model = models[modelo]
-    pipe = model.to(device).eval()
+def processar_imagem_de_entrada(img, modelo, pipe):
+    img = img.copy()
+    pipe = pipe.to(device).eval()
     start = time.time()
-    img, h, w, ori_gray = preprocess_image(img, modelo)
+    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':
@@ -74,33 +61,42 @@ def process_image(img, modelo):
         logit = cv2.resize(logit, (h, w))
 
     logit = logit.astype(bool)
-    img_out = img.cpu().numpy().squeeze().copy()
-    img_out[logit] = 1  # Set the mask to 1 (white)
-    img_out[~logit] = 0  # Set other regions to 0 (black)
-    return spent, img_out, detect_disease(logit)
+    img_out = ori.copy()
+    img_out[logit, 0] = 255
+    return spent, img_out
 
-def detect_disease(logit):
-    # Implement your disease detection logic here based on the logit mask
-    # For example, you can count the number of white pixels or use any other criteria.
-    # For the sake of simplicity, let's assume that we don't detect any disease for now.
-    return "No Disease"
+# 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'),
+}
+
+for model_name, model in models.items():
+    models[model_name] = model.to(device).eval()
+
+def processar_imagem_de_entrada_wrapper(img, modelo):
+    model = models[modelo]
+    return processar_imagem_de_entrada(img, modelo, model)
 
 my_app = gr.Interface(
-    fn=process_image,
+    fn=processar_imagem_de_entrada_wrapper,
     inputs=[
-        Image(label="Angiograma:", shape=(512, 512)),
-        Dropdown(['SE-RegUNet 4GF', 'SE-RegUNet 16GF', 'AngioNet', 'EffUNet++ B5', 'Reg-SA-UNet++', 'UNet3+'], label='Modelo', default='SE-RegUNet 4GF'),
+        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=[
-        Label(label="Tempo decorrido"),
-        Image(label="Imagem de Saída"),
-        Label(label="Status da Doença"),
+        gr.outputs.Label(label="Tempo decorrido"),
+        gr.outputs.Image(type="numpy", label="Imagem de Saída"),
     ],
     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.\n\nStatus da Doença:\n- 'Disease Detected': Indica que a segmentação detectou uma área significativa de estenose.\n- 'No Disease': Indica que a segmentação não detectou estenose significativa.\nCom base no resultado, é recomendado consultar um profissional de saúde para avaliação e orientação adicional, se necessário.",
+    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",
-    allow_flagging="never",
-    live=True,  # Add this line to enable live updates
+    layout="vertical",
+    allow_flagging=False,
 )
 
-my_app.launch()  # Remove this line if using the 'share=True' option
+my_app.launch()