app.py

import io, numpy as np, tensorflow as tf
from PIL import Image, UnidentifiedImageError
import gradio as gr
from pathlib import Path

# ========= Config =========
BASE_DIR       = Path(__file__).parent
SAVEDMODEL_DIR = BASE_DIR / "models" / "efficientnetb0_best" / "savedmodel"
INPUT_SIZE     = (224, 224)
THRESHOLD      = 0.61

print(f"[startup] Looking for SavedModel at: {SAVEDMODEL_DIR}", flush=True)
assert (SAVEDMODEL_DIR / "saved_model.pb").exists(), (
    f"SavedModel not found at {SAVEDMODEL_DIR}. "
    "Expected saved_model.pb and variables/ inside that folder."
)

# Load SavedModel once (serving_default signature)
loaded = tf.saved_model.load(str(SAVEDMODEL_DIR))
infer  = loaded.signatures["serving_default"]
print("[startup] SavedModel loaded OK.", flush=True)

# ========= Readers & preprocessing =========
def _read_dicom_bytes(b: bytes):
    import pydicom
    dcm = pydicom.dcmread(io.BytesIO(b), force=True)
    arr = dcm.pixel_array.astype("float32")
    # Fix inverted grayscale if MONOCHROME1
    if "MONOCHROME1" in str(getattr(dcm, "PhotometricInterpretation", "")).upper():
        arr = arr.max() - arr
    # Normalize to [0,1] and expand to RGB
    if arr.max() > 0:
        arr = arr / arr.max()
    if arr.ndim == 2:
        arr = np.stack([arr]*3, axis=-1)
    elif arr.ndim == 3 and arr.shape[-1] == 1:
        arr = np.concatenate([arr]*3, axis=-1)
    return arr

def _read_image_bytes(b: bytes):
    img = Image.open(io.BytesIO(b)).convert("RGB")
    return (np.asarray(img).astype("float32") / 255.0)

def _prep_01_rgb(arr):
    x = tf.image.resize(arr, INPUT_SIZE).numpy().astype("float32")
    x = np.clip(x, 0.0, 1.0)
    xb = np.expand_dims(x, 0)  # [1,H,W,3]
    return tf.constant(xb), x  # batch tensor, preview image

def _get_prob(out_dict):
    t = next(iter(out_dict.values()))
    v = t.numpy()
    # Supports sigmoid single-output or softmax two-output heads
    return float(v.reshape(-1)[0]) if v.shape[-1] == 1 else float(v[0, 1])

# ========= Robust Gradio handler =========
def predict(file):
    if file is None:
        return None, {"NORMAL": 0.0, "PNEUMONIA": 0.0}, "No file uploaded."

    raw, name = None, "upload"

    # Newer Gradio often provides a dict with {"path": "...", "orig_name": "..."}
    try:
        if isinstance(file, dict) and "path" in file:
            p = Path(file["path"])
            name = file.get("orig_name", p.name)
            raw = p.read_bytes()
        elif isinstance(file, (str, Path)):
            p = Path(file)
            name = p.name
            raw = p.read_bytes()
        else:  # file-like object
            name = getattr(file, "name", "upload")
            raw = file.read()
    except Exception as e:
        return None, {"NORMAL": 0.0, "PNEUMONIA": 0.0}, f"Could not read upload: {e}"

    # Prefer reader by extension; fall back once to the other
    try:
        if str(name).lower().endswith((".dcm", ".dicom")):
            arr = _read_dicom_bytes(raw)
        else:
            arr = _read_image_bytes(raw)
    except (Exception, UnidentifiedImageError):
        try:
            arr = _read_dicom_bytes(raw)
        except Exception:
            try:
                arr = _read_image_bytes(raw)
            except Exception as e:
                return None, {"NORMAL": 0.0, "PNEUMONIA": 0.0}, f"Unreadable file: {e}"

    xb, x_disp = _prep_01_rgb(arr)
    out  = infer(xb)
    prob = _get_prob(out)
    decision = "PNEUMONIA" if prob >= THRESHOLD else "NORMAL"
    return x_disp, {"NORMAL": 1 - prob, "PNEUMONIA": prob}, \
           f"Decision @ {THRESHOLD:.2f}: {decision} (p={prob:.3f})"

# ========= UI =========
with gr.Blocks(title="Pneumonia Detector – EfficientNetB0") as demo:
    gr.Markdown("# Pneumonia Detector – EfficientNetB0")
    gr.Markdown(
        "Upload **PNG/JPG** or **DICOM**. Input is resized to **224×224 RGB [0,1]**; "
        "the SavedModel handles internal calibration. Operating threshold **0.61**."
    )
    inp = gr.File(label="Upload chest X-ray (.png, .jpg, .dcm)")
    with gr.Row():
        out_img = gr.Image(label="Model Input (224×224)")
        out_lbl = gr.Label(label="Probabilities")
    out_txt = gr.Textbox(label="Decision")
    gr.Button("Predict").click(predict, [inp], [out_img, out_lbl, out_txt])

if __name__ == "__main__":
    demo.launch()  # On Spaces, no need for share=True