Update app.py

app.py

@@ -3,53 +3,26 @@ import torch
 import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
-from torchvision.transforms.functional import to_pil_image
 from model import load_model
 import matplotlib.pyplot as plt
 import numpy as np
 from thop import profile
 import io
 
-def calculate_performance_metrics(model, device, input_size=(1,3,224,224)):
-    model.to(device)
-    inputs = torch.randn(input_size).to(device)
-    flops, params = profile(model, inputs=(inputs,), verbose=False)
-    params_million = params / 1e6
-    flops_billion = flops / 1e9
-    # timing
-    with torch.no_grad():
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-        start.record()
-        _ = model(inputs)
-        end.record()
-        torch.cuda.synchronize()
-    speed_gpu_ms = start.elapsed_time(end)
-    # CPU timing
-    inputs_cpu = inputs.to('cpu')
-    start_c = torch.cuda.Event(enable_timing=True)
-    end_c = torch.cuda.Event(enable_timing=True)
-    # use time.time as fallback for CPU
-    import time
-    t0 = time.time()
-    _ = model(inputs_cpu)
-    t1 = time.time()
-    speed_cpu_ms = (t1 - t0) * 1000
-    return {
-        'params_million': round(params_million,2),
-        'flops_billion': round(flops_billion,2),
-        'speed_cpu_ms': round(speed_cpu_ms,2),
-        'speed_gpu_ms': round(speed_gpu_ms,2)
-    }
+# Device selection
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Cache models to avoid repeated downloads
+models_cache = {}
 
-# Preprocess transform
-def get_transform():
-    return transforms.Compose([
-        transforms.Resize((224,224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
-    ])
+# Preprocess transform for 224x224 input
+transform = transforms.Compose([
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
+])
 
+# Class names
 class_names = [
     'Alzheimer Disease',
     'Mild Alzheimer Risk',
@@ -59,42 +32,80 @@ class_names = [
     'Parkinson Disease'
 ]
 
-# Gradio predict function
+# Performance metrics calculation outside predict to not block UI
+def calculate_performance(model):
+    model.eval()
+    dummy = torch.randn(1,3,224,224).to(device)
+    flops, params = profile(model, inputs=(dummy,), verbose=False)
+    params_m = round(params/1e6,2)
+    flops_b = round(flops/1e9,2)
+    # inference timing on CPU
+    import time
+    start = time.time()
+    _ = model(dummy.cpu())
+    cpu_ms = round((time.time() - start)*1000,2)
+    # inference timing on GPU if available
+    if device.type == 'cuda':
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)
+        start_event.record()
+        _ = model(dummy)
+        end_event.record()
+        torch.cuda.synchronize()
+        gpu_ms = round(start_event.elapsed_time(end_event),2)
+    else:
+        gpu_ms = None
+    return {'params_million':params_m, 'flops_billion':flops_b, 'cpu_ms':cpu_ms, 'gpu_ms':gpu_ms}
+
+# Prediction function
 def predict_and_monitor(version, image):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = load_model(version, device)
-    # preprocess
-    img = image.convert("RGB")
-    tensor = get_transform()(img).unsqueeze(0).to(device)
-    # inference
-    with torch.no_grad():
-        outputs = model(tensor)
-        probs = F.softmax(outputs, dim=1)[0]
-    # prepare top3
-    topk = torch.topk(probs, k=3)
-    pred_items = {class_names[i]: round(float(probs[i]),4) for i in range(len(class_names))}
-    # metrics
-    metrics = calculate_performance_metrics(model, device)
-    # plot input image with prediction
-    buf = io.BytesIO()
-    plt.figure(figsize=(4,4))
-    plt.imshow(img)
-    plt.title(f"Top1: {topk.indices[0]} ({topk.values[0]:.4f})")
-    plt.axis('off')
-    plt.savefig(buf, format='png')
-    plt.close()
-    buf.seek(0)
-    return pred_items, metrics, buf
+    try:
+        # load or get cached model
+        if version not in models_cache:
+            models_cache[version] = load_model(version, device)
+        model = models_cache[version]
+
+        # preprocess
+        if image is None:
+            raise gr.Error("Görsel yüklenmedi.")
+        img = image.convert("RGB")
+        tensor = transform(img).unsqueeze(0).to(device)
+
+        # inference
+        with torch.no_grad():
+            logits = model(tensor)
+            probs = F.softmax(logits, dim=1)[0]
+
+        # prepare outputs
+        pred_dict = {class_names[i]: round(float(probs[i]),4) for i in range(len(class_names))}
+        metrics = calculate_performance(model)
+
+        # plot image with top1 label
+        top1 = max(pred_dict, key=pred_dict.get)
+        buf = io.BytesIO()
+        plt.figure(figsize=(3,3))
+        plt.imshow(img)
+        plt.title(f"{top1}: {pred_dict[top1]*100:.1f}%")
+        plt.axis('off')
+        plt.savefig(buf, format='png')
+        plt.close()
+        buf.seek(0)
+
+        return pred_dict, metrics, buf
+    except Exception as e:
+        # show exception message
+        raise gr.Error(f"Tahmin hatası: {e}")
 
+# Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Vbai-DPA 2.1 Risk Classification & Monitoring")
+    gr.Markdown("# Vbai-DPA Risk Classification & Monitoring")
     with gr.Row():
         version = gr.Radio(['f','c','q'], value='c', label="Model Version")
         image_in = gr.Image(type="pil", label="Brain Slice (224x224)")
     with gr.Row():
         preds = gr.JSON(label="Prediction Probabilities")
         stats = gr.JSON(label="Performance Metrics")
-        plot = gr.Image(label="Input & Top-1");
+        plot = gr.Image(label="Input & Top1")
     btn = gr.Button("Run")
     btn.click(fn=predict_and_monitor, inputs=[version, image_in], outputs=[preds, stats, plot])