Update app.py

app.py

@@ -1,138 +1,23 @@
-import os
 import gradio as gr
-import torch
-import cv2
-import numpy as np
-from preprocess import unsharp_masking
-import time
-from sklearn.cluster import KMeans
+from PIL import Image
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Import the ObstructionDetector class from your module
+from obstruction_detector import ObstructionDetector
 
-def preprocess_image(img, model='SE-RegUNet 4GF'):
-    img = cv2.resize(img, (512, 512))
-    img = unsharp_masking(img).astype(np.uint8)
+# Create an instance of ObstructionDetector
+detector = ObstructionDetector()
 
-    def normalize_image(img):
-        return np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
+# Define a Gradio function to process the image and return the report
+def process_image(image):
+    # Call the detect_obstruction method of the ObstructionDetector with the PIL image
+    report = detector.detect_obstruction(image)
+    
+    return report
 
-    if model in ('AngioNet', 'UNet3+'):
-        img = normalize_image(img)
-        img_out = np.expand_dims(img, axis=0)
-    elif model == '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 = normalize_image(img)
-        image1 = normalize_image(image1)
-        image2 = normalize_image(image2)
-        img_out = np.stack((img, image1, image2), axis=0)
-    else:
-        clahe1 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-        image1 = clahe1.apply(img)
-        image1 = normalize_image(image1)
-        img_out = np.stack((image1,) * 3, axis=0)
+# Define the Gradio interface
+iface = gr.Interface(fn=process_image, 
+                     inputs=gr.inputs.Image(shape=(224, 224)),  # Adjust shape as needed
+                     outputs="text")
 
-    return img_out
-
-def process_input_image(img, model, save_result=False):
-    try:
-        img = img.copy()
-        pipe = models[model].to(device).eval()
-        start = time.time()
-        img, h, w, ori_gray, ori = preprocess_image(img, model)
-        img = torch.FloatTensor(img).unsqueeze(0).to(device)
-
-        with torch.no_grad():
-            if model == '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 = f"{time.time() - start:.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
-
-        if save_result:
-            file_name = os.path.join(caminho_salvar_resultado, f'resultado_{int(time.time())}.png')
-            cv2.imwrite(file_name, img_out)
-
-        return spent, img_out
-
-    except Exception as e:
-        return str(e), None
-
-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 process_input_image_wrapper(img, model, save_result=False):
-    resultado, img_out = process_input_image(img, model, save_result)
-
-    kmeans = KMeans(n_clusters=2, random_state=0)
-    flattened_img = img_out[:, :, 0].reshape((-1, 1))
-    kmeans.fit(flattened_img)
-    labels = kmeans.labels_
-    cluster_centers = kmeans.cluster_centers_
-
-    num_clusters = len(cluster_centers)
-    cluster_features = []
-
-    for i in range(num_clusters):
-        cluster_mask = labels == i
-        area = np.sum(cluster_mask)
-
-        if area == 0:
-            continue
-
-        contours, _ = cv2.findContours(np.uint8(cluster_mask), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-        if len(contours) > 0:
-            perimeter = cv2.arcLength(contours[0], True)
-            compactness = 4 * np.pi * area / (perimeter ** 2)
-
-            cluster_features.append({'area': area, 'compactness': compactness})
-
-    has_disease_flag = any(feature['area'] >= 200 and feature['compactness'] < 0.3 for feature in cluster_features)
-
-    status_doenca = "Sim" if has_disease_flag else "Não"
-    explanation = "A máquina detectou uma possível doença nos vasos sanguíneos." if has_disease_flag else "A máquina não detectou nenhuma doença nos vasos sanguíneos."
-
-    return resultado, img_out, status_doenca, explanation, f"{num_analises} análises realizadas"
-
-caminho_salvar_resultado = "/Segmento_de_Angio_Coronariana_v5/Salvar Resultado"
-num_analises = 0
-
-my_app = gr.Interface(
-    fn=process_input_image_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'),
-        gr.inputs.Checkbox(label="Salvar Resultado"),
-    ],
-    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"),
-        gr.outputs.Label(label="Análises Realizadas"),
-    ],
-    title="Segmentação de Angiograma Coronariano",
-    description="Esta aplicação segmenta angiogramas coronarianos usando modelos de segmentação pré-treinados.",
-    theme="default",
-    layout="vertical",
-    allow_flagging=False,
-)
-
-my_app.launch()
+# Launch the Gradio interface
+iface.launch()