Update app.py

app.py

@@ -8,7 +8,6 @@ import time
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def preprocess_image(img, model_name):
-    # Preprocess the input image based on the selected model
     img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
     img = cv2.resize(img, (512, 512))
     img = cv2.GaussianBlur(img, (0, 0), 1.0)
@@ -23,7 +22,7 @@ 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]
+        logit = torch.softmax(model(img_tensor), dim=1).detach().cpu().numpy()[0, 0]
     return (logit > 0.5).astype(np.uint8)
 
 def has_disease(segmented_img):
@@ -39,10 +38,12 @@ def explanation(has_disease_flag):
         explanation_text = "O angiograma não tem doença."
     return explanation_text
 
-def process_input_image(input_img, model_name, model):
+def process_input_image(input_img, model_name):
     start_time = time.time()
     
     preprocessed_img = preprocess_image(input_img, model_name)
+    model_path = model_paths[model_name]  # Assume model_paths is defined
+    model = load_model(model_path)
     processed_img = process_image(preprocessed_img, model)
     disease_flag = has_disease(processed_img)
     explanation_text = explanation(disease_flag)
@@ -51,35 +52,37 @@ def process_input_image(input_img, model_name, model):
     return elapsed_time, processed_img, disease_flag, explanation_text
 
 def main():
-    # Load models
+    # Define model paths
     model_paths = {
         'SE-RegUNet 4GF': './model/SERegUNet4GF.pt',
-        # ... (other model paths here)
+        # ... other model paths
     }
-    models = {model_name: load_model(model_path) for model_name, model_path in model_paths.items()}
+    
+    # Define Gradio components
+    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"),
+    ]
     
     # 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"),
-        ],
+        inputs=inputs,
+        outputs=outputs,
         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...",
         theme="default",
-        layout="vertical",
-        allow_flagging=False,
     )
     
     # Launch Gradio interface
     gr_interface.launch()
 
 if __name__ == "__main__":
-    main()
+    main()