app.py

"""
Bee Colony Survival — multi-modal (audio + virus) batch predictor (Hugging Face Space).

For every uploaded hive recording:
  1. AST (Audio Spectrogram Transformer) encodes the AUDIO -> 768-d embedding,
  2. the file name is matched to a row of the bundled AI_Data_Training.xlsx to read
     the viral loads (CBPV/DWV/KBV) + metadata + ground-truth labels,
  3. a class-balanced gradient-boosting model on [audio embedding + viral/metadata]
     predicts Terminal (T) / Survivor (S) for 0-3 and 4-6 months at a STRICT
     high-precision threshold (constraint: 0-3 T => 4-6 T),
  4. prediction and ground truth are shown side by side in a results table.

Batch upload is supported. Files not in the workbook still get a prediction from
audio + parsed metadata (viral loads default to 0) and show "not available" truth.
"""
import os
import json
import numpy as np
import joblib
import gradio as gr

import bee_features as bf
import ast_encoder as ae
import xls_lookup as xl

HERE = os.path.dirname(os.path.abspath(__file__))
MODELDIR = os.path.join(HERE, "model")
XLSX = os.path.join(HERE, "AI_Data_Training.xlsx")

CFG = json.load(open(os.path.join(MODELDIR, "config.json")))
M03 = joblib.load(os.path.join(MODELDIR, "model_03.joblib"))
M46 = joblib.load(os.path.join(MODELDIR, "model_46.joblib"))
TAB_ORDER = CFG["tabular_order"]
THR03, THR46 = CFG["threshold_03"], CFG["threshold_46"]

LUT = xl.build_lookup(XLSX)                      # ground truth + viral loads from Excel
FE, AST, AST_DEV = ae.load_ast()                 # AST weights download on first run

LBL = {"S": "Survivor (S)", "T": "Terminal (T)"}
HEADERS = ["File", "Colony", "Country", "Month", "Viral CBPV / DWV / KBV",
           "Predicted 0–3", "Actual 0–3", "Predicted 4–6", "Actual 4–6", "Match"]


def _results_html(rows):
    """Render results as a self-styled HTML table (full control, no truncation)."""
    th = "".join(
        f'<th style="border:1px solid #e3e3e3;padding:9px 14px;background:#f4f4f4;'
        f'color:#000;font-weight:600;text-align:left;white-space:nowrap;">{h}</th>'
        for h in HEADERS)
    trs = []
    for r in rows:
        tds = []
        for i, v in enumerate(r):
            extra = "white-space:nowrap;"
            if HEADERS[i] == "Match":
                c = {"✓": "#157347", "✗": "#b42318"}.get(v, "#888")
                extra += f"text-align:center;font-weight:700;font-size:16px;color:{c};"
            tds.append(f'<td style="border:1px solid #e3e3e3;padding:9px 14px;color:#000;'
                       f'{extra}">{v}</td>')
        trs.append("<tr>" + "".join(tds) + "</tr>")
    return ('<div style="overflow-x:auto;width:100%;">'
            '<table style="border-collapse:collapse;background:#fff;color:#000;'
            'font-family:-apple-system,Segoe UI,Roboto,Helvetica,Arial,sans-serif;font-size:14px;">'
            f'<thead><tr>{th}</tr></thead><tbody>{"".join(trs)}</tbody></table></div>')


def _predict_one(path):
    fname = os.path.basename(path)
    parsed = bf.parse_filename(fname)
    entry = xl.lookup(LUT, fname)

    colony = (entry or {}).get("colony") or parsed["colony"] or "?"
    country = (entry or {}).get("country") or parsed["country"] or "?"
    month = (entry or {}).get("month") or parsed["month"] or "?"
    cbpv = (entry or {}).get("cbpv", 0.0)
    dwv = (entry or {}).get("dwv", 0.0)
    kbv = (entry or {}).get("kbv", 0.0)

    emb = ae.encode_file(path, FE, AST, AST_DEV)               # AUDIO
    tabd = bf.build_tabular_dict(country, month, cbpv, dwv, kbv)  # VIRUS + metadata
    x = np.hstack([emb, [tabd[k] for k in TAB_ORDER]]).reshape(1, -1)

    p03 = float(M03.predict_proba(x)[0, 1]); p46 = float(M46.predict_proba(x)[0, 1])
    pred03 = "T" if p03 >= THR03 else "S"
    pred46 = "T" if (p46 >= THR46 or pred03 == "T") else "S"

    if entry is not None:
        gt03, gt46 = entry["y03"], entry["y46"]
        correct = "✓" if (gt03 == pred03 and gt46 == pred46) else "✗"
        t03, t46 = LBL[gt03], LBL[gt46]
    else:
        t03 = t46 = "not available"; correct = "—"
    viral = f"{cbpv:.2f} / {dwv:.2f} / {kbv:.2f}"
    return [fname, colony, country, month, viral,
            LBL[pred03], t03, LBL[pred46], t46, correct]


def predict_batch(files, progress=gr.Progress()):
    if not files:
        return "", "Upload one or more `.wav` files (or click an example), then press **Predict**."
    paths = [f if isinstance(f, str) else f.name for f in files]
    rows = []
    for p in progress.tqdm(paths, desc="Scoring recordings"):
        try:
            rows.append(_predict_one(p))
        except Exception as e:  # noqa
            rows.append([os.path.basename(p), "?", "?", "?", "",
                         "error", str(e)[:30], "error", "", "✗"])
    n_known = sum(1 for r in rows if r[9] != "—")
    n_ok = sum(1 for r in rows if r[9] == "✓")
    summary = (f"Scored **{len(rows)}** file(s) · **{n_known}** found in the Excel workbook "
               f"· **{n_ok}/{n_known}** predicted correctly (both horizons).")
    return _results_html(rows), summary


CSS = """
.gradio-container {background:#ffffff !important; color:#000000 !important;
  font-family: -apple-system, Segoe UI, Roboto, Helvetica, Arial, sans-serif; max-width:1150px;}
.gradio-container * {color:#000000;}
h1,h2,h3,p,label {color:#000000 !important;}
button.primary, .primary {background:#000000 !important; color:#ffffff !important;
  border:1px solid #000000 !important; border-radius:6px !important;}
button.primary:hover, .primary:hover {background:#222222 !important;}
button.primary *, .primary * {color:#ffffff !important;}
.secondary {background:#ffffff !important; color:#000000 !important; border:1px solid #000000 !important;}
footer {display:none !important;}
"""
THEME = gr.themes.Base(primary_hue=gr.themes.colors.gray,
                       neutral_hue=gr.themes.colors.gray).set(
    body_background_fill="#ffffff", body_text_color="#000000",
    block_background_fill="#ffffff", block_border_color="#e5e5e5",
    button_primary_background_fill="#000000", button_primary_text_color="#ffffff",
    button_primary_background_fill_hover="#222222")

EXDIR = os.path.join(HERE, "examples")
example_files = (sorted(os.path.join(EXDIR, f) for f in os.listdir(EXDIR) if f.endswith(".wav"))
                 if os.path.isdir(EXDIR) else [])

with gr.Blocks(title="Bee Colony Survival — Multimodal AI Predictor", theme=THEME, css=CSS) as demo:
    gr.Markdown("# Bee Colony Survival — Multimodal AI Predictor")
    gr.Markdown("Upload one or more hive **audio recordings** (batch supported). The model "
                "combines the **audio** (AST embedding) with the recording's **viral loads and "
                "metadata** (read from the bundled Excel workbook) to predict **Survivor (S)** vs "
                "**Terminal (T)** for the next 0–3 and 4–6 months, shown against the ground truth.")
    with gr.Row():
        audio_in = gr.File(label="Hive recordings (.wav) — batch upload",
                           file_count="multiple", file_types=[".wav"], type="filepath")
    with gr.Row():
        btn = gr.Button("Predict", variant="primary")
        clear = gr.Button("Clear", variant="secondary")
    summary = gr.Markdown()
    gr.Markdown("### Results")
    table = gr.HTML()
    if example_files:
        gr.Markdown("**Example recordings** — click any one to load it and run the prediction "
                    "(all are Terminal colonies the model correctly detects). "
                    "You can also drag several into the upload box above for a batch.")
        gr.Examples(
            examples=[[[f]] for f in example_files],
            inputs=[audio_in],
            outputs=[table, summary],
            fn=predict_batch,
            run_on_click=True,
            cache_examples=False,
            examples_per_page=20,
            label="Bundled examples",
        )
    btn.click(predict_batch, inputs=audio_in, outputs=[table, summary])
    clear.click(lambda: (None, "", ""), outputs=[audio_in, table, summary])

if __name__ == "__main__":
    demo.launch()