Upload src/egg_damage/gradio_app.py

src/egg_damage/gradio_app.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+
+from .compare_models import load_best_model_record
+from .config import load_config
+from .inference import EggDamagePredictor, list_available_model_records
+from .utils import get_logger
+
+
+LOGGER = get_logger(__name__)
+
+
+def probability_figure(probabilities: dict[str, float]):
+    fig, ax = plt.subplots(figsize=(5.2, 3.2))
+    labels = list(probabilities.keys())
+    values = [probabilities[label] for label in labels]
+    colors = ["#2a9d8f", "#e76f51"]
+    ax.bar(labels, values, color=colors[: len(labels)])
+    ax.set_ylim(0, 1)
+    ax.set_ylabel("Probability")
+    ax.set_title("Class Confidence")
+    for idx, value in enumerate(values):
+        ax.text(idx, min(value + 0.03, 0.98), f"{value:.2f}", ha="center", va="bottom")
+    fig.tight_layout()
+    return fig
+
+
+def build_app(config: dict[str, Any]):
+    import gradio as gr
+
+    records = list_available_model_records(config)
+    if not records:
+        raise FileNotFoundError(
+            f"No trained models found in {config['paths']['model_dir']}. "
+            "Run python scripts/train_all.py and python scripts/evaluate_all.py first."
+        )
+    try:
+        best = load_best_model_record(config)
+        records = sorted(records, key=lambda r: 0 if r["model_name"] == best["model_name"] else 1)
+    except Exception:
+        best = records[0]
+    choices = [record["model_name"] for record in records]
+    by_name = {record["model_name"]: record for record in records}
+    cache: dict[str, EggDamagePredictor] = {}
+
+    def get_predictor(model_name: str) -> EggDamagePredictor:
+        if model_name not in cache:
+            cache[model_name] = EggDamagePredictor(by_name[model_name], config)
+        return cache[model_name]
+
+    def predict(image, model_name: str):
+        if image is None:
+            return {}, "Upload an egg image to classify it.", None, None, ""
+        predictor = get_predictor(model_name)
+        result = predictor.predict(image)
+        summary = (
+            f"Prediction: {result['predicted_label']}\n"
+            f"Confidence: {result['confidence']:.3f}\n"
+            f"Model: {result['model_name']}"
+        )
+        warning = ""
+        if result["low_confidence"]:
+            warning = (
+                "Confidence is modest. Use this as a screening signal, and review the image manually "
+                "if the egg surface, lighting, or background looks unusual."
+            )
+        explanation = None
+        if (
+            config.get("explainability", {}).get("enabled", True)
+            and predictor.model_type == "deep_learning"
+            and predictor.metadata.get("family", "cnn") == "cnn"
+        ):
+            try:
+                from .explainability import gradcam_overlay
+
+                explanation = gradcam_overlay(
+                    predictor.model,
+                    image,
+                    predictor.metadata.get("config", config),
+                    target_class=result["predicted_index"],
+                    device=predictor.device,
+                )
+            except Exception as exc:
+                LOGGER.warning("Grad-CAM not available for this prediction: %s", exc)
+        return result["probabilities"], summary, probability_figure(result["probabilities"]), explanation, warning
+
+    with gr.Blocks(title="Egg Damage Classifier") as demo:
+        gr.Markdown(
+            "# Egg Damage Classifier\n"
+            "Upload a clear egg image and choose a trained model. The default is the best-ranked model from evaluation."
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                image = gr.Image(type="pil", label="Egg image")
+                model_choice = gr.Dropdown(
+                    choices=choices,
+                    value=best.get("model_name", choices[0]),
+                    label="Model",
+                )
+                button = gr.Button("Classify", variant="primary")
+            with gr.Column(scale=1):
+                label = gr.Label(label="Prediction", num_top_classes=2)
+                summary = gr.Textbox(label="Result", lines=4)
+                plot = gr.Plot(label="Probability chart")
+                explanation = gr.Image(type="pil", label="Explanation")
+                warning = gr.Markdown()
+        button.click(predict, inputs=[image, model_choice], outputs=[label, summary, plot, explanation, warning])
+    return demo
+
+
+def launch(config: dict[str, Any]) -> None:
+    demo = build_app(config)
+    demo.launch(
+        server_name=str(config["gradio"].get("host", "127.0.0.1")),
+        server_port=int(config["gradio"].get("port", 7860)),
+        share=bool(config["gradio"].get("share", False)),
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Launch the local Gradio app.")
+    parser.add_argument("--config", default="configs/default.yaml")
+    args = parser.parse_args()
+    launch(load_config(args.config))
+
+
+if __name__ == "__main__":
+    main()
+