Fix: lazy loading, disable TTA on CPU, error handling, remove double inference

app.py

@@ -1,28 +1,31 @@
-"""Gradio app for WoundNetB7 DFU Analysis — Hugging Face Spaces deployment.
-
-Launch locally: python app.py
-Deploy to HF: push this repo to a Hugging Face Space (GPU recommended).
-"""
+"""Gradio app for WoundNetB7 DFU Analysis — Hugging Face Spaces deployment."""
 import gradio as gr
 import numpy as np
 import cv2
 import json
-from pipeline import WoundNetB7Pipeline
+import traceback
+
+# Lazy loading — don't crash at import time
+pipe = None
+
 
-# Initialize pipeline (loads all models once)
-pipe = WoundNetB7Pipeline(models_dir="models", use_tta=True)
+def get_pipeline():
+    global pipe
+    if pipe is None:
+        from pipeline import WoundNetB7Pipeline
+        pipe = WoundNetB7Pipeline(models_dir="models", use_tta=False)
+    return pipe
 
-CLASS_COLORS_RGB = {
-    0: (0, 0, 0),
-    1: (0, 255, 0),      # Foot: green
-    2: (255, 165, 0),     # Perilesion: orange
-    3: (255, 0, 0),       # Ulcer: red
-}
 
-FITZ_COLORS = {
-    "I": "#fef3c7", "II": "#fde68a", "III": "#fbbf24",
-    "IV": "#b45309", "V": "#78350f", "VI": "#451a03",
-}
+def create_overlay(img_rgb, classmap):
+    """Create segmentation overlay on RGB image."""
+    colors = {1: (0, 255, 0), 2: (255, 165, 0), 3: (255, 0, 0)}
+    overlay = img_rgb.astype(np.float32).copy()
+    for cid, color in colors.items():
+        mask = classmap == cid
+        if np.any(mask):
+            overlay[mask] = overlay[mask] * 0.5 + np.array(color, dtype=np.float32) * 0.5
+    return overlay.astype(np.uint8)
 
 
 def analyze_image(image):
@@ -30,29 +33,31 @@ def analyze_image(image):
     if image is None:
         return None, "Please upload an image.", "{}"
 
-    # Gradio provides RGB numpy array
-    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    try:
+        pipeline = get_pipeline()
+        img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        result = pipeline.analyze(img_bgr, use_tta=False)
 
-    # Run pipeline
-    result = pipe.analyze(img_bgr, use_tta=True)
+        # Create overlay from the segmentation already done (no re-run)
+        from src.segmentation import segment, CLASS_NAMES
+        seg = segment(pipeline.seg_model, img_bgr, pipeline.device, use_tta=False)
+        classmap = seg["classmap"]
 
-    # Create overlay visualization
-    overlay = pipe.visualize(img_bgr, result)
+        if classmap.shape[:2] != image.shape[:2]:
+            classmap = cv2.resize(classmap, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
 
-    # Build text summary
-    summary = result.summary()
+        overlay = create_overlay(image, classmap)
+        summary = result.summary()
+        json_out = json.dumps(result.to_dict(), indent=2, ensure_ascii=False)
 
-    # Build JSON output
-    json_out = json.dumps(result.to_dict(), indent=2, ensure_ascii=False)
+        return overlay, summary, json_out
 
-    return overlay, summary, json_out
+    except Exception as e:
+        error_msg = f"Error: {str(e)}\n\n{traceback.format_exc()}"
+        return None, error_msg, "{}"
 
 
-# Build Gradio interface
-with gr.Blocks(
-    title="WoundNetB7 DFU Analysis",
-    theme=gr.themes.Soft(),
-) as demo:
+with gr.Blocks(title="WoundNetB7 DFU Analysis", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         # WoundNetB7 — Diabetic Foot Ulcer Analysis
@@ -70,7 +75,6 @@ with gr.Blocks(
         with gr.Column(scale=1):
             input_image = gr.Image(label="DFU Image", type="numpy")
             analyze_btn = gr.Button("Analyze", variant="primary", size="lg")
-
         with gr.Column(scale=1):
             output_overlay = gr.Image(label="Segmentation Overlay")
 
@@ -80,11 +84,7 @@ with gr.Blocks(
         with gr.Column(scale=1):
             output_json = gr.Code(label="JSON Output", language="json")
 
-    analyze_btn.click(
-        fn=analyze_image,
-        inputs=[input_image],
-        outputs=[output_overlay, output_text, output_json],
-    )
+    analyze_btn.click(fn=analyze_image, inputs=[input_image], outputs=[output_overlay, output_text, output_json])
 
     gr.Markdown(
         """
@@ -94,8 +94,7 @@ with gr.Blocks(
         **PWAT Items:** 3=Necrotic Type, 4=Necrotic Amount, 5=Granulation Type,
         6=Granulation Amount, 7=Edges, 8=Periulcer Skin (0=best, 4=worst)
 
-        **Debiasing:** Scores adjusted by Fitzpatrick type to reduce skin-tone bias
-        (18% max group gap reduction, p < 10^-27).
+        **Debiasing:** Scores adjusted by Fitzpatrick type to reduce skin-tone bias.
         """
     )