Update app.py

app.py

@@ -75,12 +75,25 @@ models = {
     'UNet3+': torch.jit.load('./model/UNet3plus.pt'),
 }
 
-for model_name, model in models.items():
-    models[model_name] = model.to(device).eval()
+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 processar_imagem_de_entrada_wrapper(img, modelo):
     model = models[modelo]
-    return processar_imagem_de_entrada(img, modelo, model)
+    spent, img_out = processar_imagem_de_entrada(img, modelo, model)
+    has_disease = has_disease(img_out)
+    return spent, img_out, has_disease
 
 my_app = gr.Interface(
     fn=processar_imagem_de_entrada_wrapper,
@@ -91,6 +104,7 @@ my_app = gr.Interface(
     outputs=[
         gr.outputs.Label(label="Tempo decorrido"),
         gr.outputs.Image(type="numpy", label="Imagem de Saída"),
+        gr.outputs.Label(label="Possui Doença?"),
     ],
     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.",