app.py

import gradio as gr

# --- 1. Import Existing Baselines ---
# Wrapped in try-except so the app doesn't crash if files are temporarily missing
try:
    from baseline.baseline_convnext import predict_convnext
except ImportError:
    def predict_convnext(image): return {"Error": "ConvNeXt module missing"}

try:
    from baseline.baseline_infer import predict_baseline
except ImportError:
    def predict_baseline(image): return {"Error": "Baseline module missing"}

# --- 2. Import NEW SPA Approach ---
# This imports the function from: new_approach/spa_ensemble.py
try:
    from new_approach.spa_ensemble import predict_spa
except ImportError:
    def predict_spa(image): return {"Error": "SPA module missing. Check 'new_approach' folder."}


# --- Placeholder models (for future extensions) ---
def predict_placeholder_2(image):
    if image is None:
        return "Please upload an image."
    return "Model 4 is not available yet. Please check back later."

# --- Main Prediction Logic ---
def predict(model_choice, image):
    if image is None: return None

    if model_choice == "Herbarium Species Classifier":
        # Friend's ConvNeXt mix-stream CNN baseline
        return predict_convnext(image)
        
    elif model_choice == "Baseline (DINOv2 + LogReg)":
        # Original baseline
        return predict_baseline(image)
        
    elif model_choice == "SPA Ensemble (New Approach)":
        # YOUR NEW CODE: DINOv2 + BioCLIP + Handcrafted + SPA
        return predict_spa(image)
        
    elif model_choice == "Future Model 2 (Placeholder)":
        return predict_placeholder_2(image)
        
    else:
        return "Invalid model selected."

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
    with gr.Column(elem_id="app-wrapper"):
        # Header
        gr.Markdown(
            """
            <div id="app-header">
              <h1>🌿 Plant Species Classification</h1>
              <h3>AML Group Project – PsychicFireSong</h3>
            </div>
            """,
            elem_id="app-header",
        )
        
        # Badges row
        gr.Markdown(
            """
            <div id="badge-row">
              <span class="badge">Herbarium + Field images</span>
              <span class="badge">ConvNeXtV2</span>
              <span class="badge">SPA Ensemble</span>
            </div>
            """,
            elem_id="badge-row",
        )
        
        # Main card
        with gr.Row(elem_id="main-card"):
            # Left side: model + image
            with gr.Column(scale=1, elem_id="left-panel"):
                model_selector = gr.Dropdown(
                    label="Select model",
                    choices=[
                        "Herbarium Species Classifier",
                        "Baseline (DINOv2 + LogReg)",
                        "SPA Ensemble (New Approach)",
                        "Future Model 2 (Placeholder)",
                    ],
                    value="SPA Ensemble (New Approach)", # Default to your new model
                )
                
                gr.Markdown(
                    """
                    <div id="model-help">
                      <b>Herbarium Classifier</b> – ConvNeXtV2 CNN.<br>
                      <b>Baseline</b> – Simple DINOv2 + LogReg.<br>
                      <b>SPA Ensemble</b> – <i>(New)</i> DINOv2 + BioCLIP-2 + Handcrafted features.
                    </div>
                    """,
                    elem_id="model-help",
                )
                
                image_input = gr.Image(
                    type="pil",
                    label="Upload plant image",
                )
                
                submit_button = gr.Button("Classify 🌱", variant="primary")

            # Right side: predictions
            with gr.Column(scale=1, elem_id="right-panel"):
                output_label = gr.Label(
                    label="Top 5 predictions",
                    num_top_classes=5,
                )

        submit_button.click(
            fn=predict,
            inputs=[model_selector, image_input],
            outputs=output_label,
        )

        # Optional examples
        gr.Examples(
            examples=[],
            inputs=image_input,
            outputs=output_label,
            fn=lambda img: predict("SPA Ensemble (New Approach)", img),
            cache_examples=False,
        )

        gr.Markdown(
            "Built for the AML course – compare CNN vs. DINOv2 feature-extractor baselines.",
            elem_id="footer",
        )

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