""" 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'{h}' 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'{v}') trs.append("" + "".join(tds) + "") return ('
' '' f'{th}{"".join(trs)}
') 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()