Upload app.py

app.py

@@ -0,0 +1,434 @@
+# -*- coding: utf-8 -*-
+"""
+Vertebrates Track Classifier (EfficientNet-B0)
+
+- Input: one or more photographs
+- Output: top-3 most probable classes + probabilities
+- Classes:
+    Bear, Coyote, Deer, Fox, Turkey, Otter,
+    Squirrel, Raccoon, Sauropod, Theropod
+"""
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"  # safe no-op in most environments
+
+import tempfile
+import html
+
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+
+from PIL import Image
+import numpy as np
+import pandas as pd
+import gradio as gr
+
+# =========================
+# Config
+# =========================
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# IMAGE_SIZE from training script (typically 224 for EfficientNet-B0)
+IMAGE_SIZE = 224
+criterion = nn.CrossEntropyLoss()
+# ----- Class names -----
+# IMPORTANT: the order MUST match the class order used during training.
+# If you used torchvision.datasets.ImageFolder, this is the alphabetical
+# order of your training subfolders.
+CLASS_NAMES = [
+    "Bear",
+    "Coyote",
+    "Deer",
+    "Fox",
+    "Otter",
+    "Raccoon",
+    "Sauropod",
+    "Squirrel",
+    "Theropod",
+    "Turkey",
+]
+
+NUM_CLASSES = len(CLASS_NAMES)
+
+# ---- Checkpoint path (Hugging Face: relative 'checkpoints' folder) ----
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+CHECKPOINT_PATH = os.path.join(THIS_DIR, "checkpoints", "model_checkpoint_8.pth")
+# Put your .pth file in: ./checkpoints/wild_dino_tracks_efficientnet_b0.pth
+# or change the filename above to match your checkpoint.
+
+# =========================
+# Preprocessing (matches training)
+# =========================
+INFER_TRANSFORM = transforms.Compose([
+    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],   # ImageNet mean
+        std=[0.229, 0.224, 0.225],    # ImageNet std
+    ),
+])
+
+
+# =========================
+# Model definitions
+# =========================
+def create_efficientnet_b0(num_classes: int) -> nn.Module:
+    """
+    EfficientNet-B0 head adapted for num_classes.
+    Matches typical transfer-learning setup:
+        model = models.efficientnet_b0(pretrained=True)
+        in_features = model.classifier[1].in_features
+        model.classifier = nn.Sequential(
+            nn.Dropout(p=0.2),
+            nn.Linear(in_features, num_classes)
+        )
+    """
+    model = models.efficientnet_b0(pretrained=True)
+    in_features = model.classifier[1].in_features
+    model.classifier = nn.Sequential(
+        nn.Dropout(p=0.2),
+        nn.Linear(in_features, num_classes),
+    )
+    return model
+
+
+def _safe_torch_load(path: str):
+    """
+    Helper to handle PyTorch 2.6+ (weights_only=True by default) and older versions.
+    """
+    try:
+        # Newer PyTorch versions
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        # Older PyTorch versions (no weights_only argument)
+        return torch.load(path, map_location="cpu")
+
+
+def load_model(checkpoint_path: str) -> nn.Module:
+    """
+    Load EfficientNet-B0 model and checkpoint.
+    """
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(
+            f"Checkpoint not found: {checkpoint_path}\n"
+            "Make sure the .pth is in the 'checkpoints' folder."
+        )
+
+    model = create_efficientnet_b0(NUM_CLASSES)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+    ckpt = _safe_torch_load(checkpoint_path)
+
+    if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
+        state_dict = ckpt["model_state_dict"]
+
+        optimizer.load_state_dict(ckpt['optimizer_state_dict'])
+    
+    
+    else:
+        state_dict = ckpt
+
+    model.load_state_dict(state_dict)
+    model.to(DEVICE)
+    model.eval()
+    return model
+
+
+# Lazy global model
+_MODEL = None
+
+
+def get_model() -> nn.Module:
+    global _MODEL
+    if _MODEL is None:
+        _MODEL = load_model(CHECKPOINT_PATH)
+    return _MODEL
+
+
+# =========================
+# Prediction helpers
+# =========================
+@torch.no_grad()
+def predict_top3_from_pil(pil_img: Image.Image):
+    """
+    Input: PIL image
+    Output: (top3_class_names, top3_probs) where probs are floats in [0,1]
+    """
+    model = get_model()
+
+    img = pil_img.convert("RGB")
+    x = INFER_TRANSFORM(img).unsqueeze(0).to(DEVICE)  # [1,3,H,W]
+
+    logits = model(x)  # [1, num_classes]
+    probs = torch.softmax(logits, dim=1)[0].cpu().numpy()  # (num_classes,)
+
+    top_idx = np.argsort(-probs)[:3]
+    top_classes = [CLASS_NAMES[i] for i in top_idx]
+    top_probs = [float(probs[i]) for i in top_idx]
+
+    return top_classes, top_probs
+
+
+def df_to_html(df: pd.DataFrame) -> str:
+    """
+    Convert the predictions DataFrame into a styled HTML table.
+    """
+    if df.empty:
+        return "<p>No predictions to display yet.</p>"
+
+    headers = df.columns.tolist()
+    header_cells = "".join(f"<th>{html.escape(str(h))}</th>" for h in headers)
+    rows_html = []
+
+    for _, row in df.iterrows():
+        cells = []
+        for col in headers:
+            val = row[col]
+
+            if val is None or (isinstance(val, float) and np.isnan(val)):
+                disp = ""
+            elif isinstance(val, float):
+                # Round decimals for readability
+                disp = f"{val:.3f}"
+            else:
+                disp = str(val)
+
+            cells.append(f"<td>{html.escape(disp)}</td>")
+
+        rows_html.append("<tr>" + "".join(cells) + "</tr>")
+
+    table_html = (
+        "<div class='pred-table'>"
+        "<table>"
+        "<thead><tr>"
+        f"{header_cells}"
+        "</tr></thead>"
+        "<tbody>"
+        f"{''.join(rows_html)}"
+        "</tbody></table>"
+        "</div>"
+    )
+    return table_html
+
+
+def classify_batch(filepaths):
+    """
+    Gradio callback.
+    """
+    cols = [
+        "image_name",
+        "top1_class", "top1_prob",
+        "top2_class", "top2_prob",
+        "top3_class", "top3_prob",
+    ]
+
+    if not filepaths:
+        empty_df = pd.DataFrame(columns=cols)
+        html_table = df_to_html(empty_df)
+        return html_table, "Please upload at least one image.", None
+
+    rows = []
+
+    for path in filepaths:
+        try:
+            pil = Image.open(path).convert("RGB")
+            top_classes, top_probs = predict_top3_from_pil(pil)
+
+            rows.append({
+                "image_name": os.path.basename(str(path)),
+                "top1_class": top_classes[0],
+                "top1_prob": top_probs[0],
+                "top2_class": top_classes[1],
+                "top2_prob": top_probs[1],
+                "top3_class": top_classes[2],
+                "top3_prob": top_probs[2],
+            })
+        except Exception as e:
+            rows.append({
+                "image_name": os.path.basename(str(path)),
+                "top1_class": f"Error: {e}",
+                "top1_prob": None,
+                "top2_class": None,
+                "top2_prob": None,
+                "top3_class": None,
+                "top3_prob": None,
+            })
+
+    df = pd.DataFrame(rows)
+    status = f"Processed {len(rows)} photograph(s)."
+
+    tmpdir = tempfile.mkdtemp()
+    csv_path = os.path.join(tmpdir, "predictions_vert_tracks.csv")
+    df.to_csv(csv_path, index=False)
+
+    html_table = df_to_html(df)
+    return html_table, status, csv_path
+
+
+# =========================
+# Gradio UI (paleo + wildlife aesthetics)
+# =========================
+theme = gr.themes.Soft(
+    primary_hue="orange",
+    secondary_hue="amber",
+    neutral_hue="gray",
+)
+
+with gr.Blocks(theme=theme, css="""
+.gradio-container {
+    font-family: 'Georgia', 'Times New Roman', serif;
+}
+
+.app-wrapper {
+    max-width: 1100px;
+    margin: 0 auto;
+    padding: 1.5rem 1rem 2rem 1rem;
+}
+
+.app-header {
+    text-align: center;
+    margin-bottom: 1.2rem;
+}
+
+.app-header h1 {
+    font-size: 2.1rem;
+    margin-bottom: 0.3rem;
+}
+
+.app-header h2 {
+    font-size: 1.1rem;
+    font-weight: normal;
+    opacity: 0.9;
+}
+
+.app-panel {
+    background: rgba(255, 255, 255, 0.85);
+    border-radius: 14px;
+    padding: 1.2rem 1.5rem;
+    margin-bottom: 1rem;
+    border: 1px solid rgba(120, 82, 45, 0.18);
+}
+
+/* === predictions table wrapper === */
+.pred-table {
+    width: 100%;
+    overflow-x: auto;     /* horizontal scrollbar if needed */
+}
+
+/* Styled table for predictions */
+.pred-table table {
+    width: 100%;
+    min-width: 650px;
+    border-collapse: collapse;
+    margin-top: 0.5rem;
+    font-size: 0.9rem;
+}
+
+.pred-table thead {
+    background: #e0cfb3;
+}
+
+.pred-table th, .pred-table td {
+    border: 1px solid #d0b897;
+    padding: 0.4rem 0.6rem;
+    text-align: center;
+    color: #000000;
+    white-space: nowrap;
+}
+
+.pred-table th {
+    font-weight: 600;
+}
+
+.pred-table tbody tr:nth-child(even) {
+    background: #f7eee2;
+}
+
+.pred-table tbody tr:nth-child(odd) {
+    background: #fbf4ea;
+}
+
+/* first column (image name) left-aligned */
+.pred-table td:first-child {
+    text-align: left;
+}
+""") as demo:
+
+    gr.HTML("<div class='app-wrapper'>")
+
+    # ----- Header -----
+    gr.HTML("""
+    <div class="app-header">
+        <h1>🐾 Vertebrate Tracks Classifier</h1>
+        <h2>Deep-learning assisted ichnological identifications with EfficientNet-B0</h2>
+           Model finetuned from a model trained on data obtained by the
+           <a href="https://zenodo.org/records/15092442" target="_blank">Deep Tracks</a>
+           App.<br>
+           Developed by <b>Carolina S. Marques</b>
+           (<a href="https://orcid.org/0000-0002-5936-9342" target="_blank">ORCID</a>)
+           as part of her PhD research, funded by CEAUL through FCT - Fundação para a Ciência e Tecnologia
+           (<a href="https://doi.org/10.54499/UI/BD/154258/2022" target="_blank">DOI</a>).
+   </div>
+
+    """)
+  
+
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.HTML("<div class='app-panel'>")
+            gr.Markdown(
+                " This model distinguishes between footprints of <b>Bear, Coyote, Deer, Fox, Turkey, Otter, Squirrel, Raccoon<b> as well as dinosaur tracks attributed to <b>Sauropod</b> and <b>Theropod</b> trackmakers.\n"
+                "#### 1. Upload track photographs\n"
+                "You can upload one or more photos of footprints from different vertebrates. "
+                "The network will estimate, for each image, the probability of belonging to each of the ten classes:\n\n"
+                "- Bear, Coyote, Deer, Fox, Turkey, Otter, Squirrel, Raccoon\n"
+                "- Sauropod, Theropod (dinosaur tracks)\n"
+            )
+
+            img_files = gr.Files(
+                label="Track photographs (you can select multiple files)",
+                file_types=["image"],
+                file_count="multiple",
+                type="filepath",
+            )
+            classify_btn = gr.Button("Run classification", variant="primary")
+            gr.HTML("</div>")
+
+        with gr.Column(scale=1.4):
+            gr.HTML("<div class='app-panel'>")
+            gr.Markdown("#### Predicted classes and probabilities")
+            results_html = gr.HTML(label="Top-3 predictions per image")
+            gr.Markdown(
+                "_How to read the table:_\n"
+                "- **top1_class** / **top1_prob**: class with the highest predicted probability for that image, and the corresponding probability.\n"
+                "- **top2_class** / **top2_prob**: second most probable class and the corresponding probability.\n"
+                "- **top3_class** / **top3_prob**: third most probable class and the corresponding probability.\n"
+                "- Probabilities are between 0 and 1 and, for each image, they sum to 1 across all ten classes."
+            )
+            gr.HTML("</div>")
+
+            gr.HTML("<div class='app-panel'>")
+            status_md = gr.Markdown()
+            df_file = gr.File(
+                label="Download full predictions as CSV",
+                file_types=[".csv"],
+            )
+            gr.Markdown(
+                "_Note_: The CSV export is plain text, ready to be used in R, Python, or Excel "
+                "for further analysis (e.g., confusion matrices, ROC curves, etc.)."
+            )
+            gr.HTML("</div>")
+
+    gr.HTML("</div>")  # close app-wrapper
+
+    classify_btn.click(
+        fn=classify_batch,
+        inputs=[img_files],
+        outputs=[results_html, status_md, df_file],
+    )
+
+# For local dev / Hugging Face Spaces:
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch()