Spaces:

Unearthed
/

artifact.classifier

Sleeping

Update app.py
55d4bcb verified 6 months ago
31.1 kB
"""
╔══════════════════════════════════════════════════════════════╗
║  UNEARTHED — Archaeological Lithic Artifact Classifier       ║
║  HuggingFace Spaces · Gradio UI                              ║
║  app.py — upload this alongside your .h5 model file          ║
╚══════════════════════════════════════════════════════════════╝
Files needed in your HF Space repo:
    app.py                          ← this file
    lithic_classifier_v2_final.h5  ← your trained model
    requirements.txt                ← see bottom of this file
"""

import gradio as gr
import numpy as np
import cv2
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from PIL import Image
from skimage.filters import gabor
from skimage.color import rgb2gray as sk_rgb2gray
from scipy.spatial import cKDTree
import tensorflow as tf
import keras
from keras.applications.efficientnet import preprocess_input
from keras import layers
import io
import os

# ═══════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════

MODEL_PATH   = "lithic_classifier_v2_clean.h5"
IMG_SIZE     = (224, 224)
CLASS_NAMES  = ["arrowheads", "axehammerstones", "notartifact", "scrapers"]

# Prediction gate thresholds
CONF_THRESHOLD        = 0.65
ENTROPY_THRESHOLD     = 1.30
TEXTURE_THRESHOLD     = 15.0
REGULARITY_THRESHOLD  = 0.30
SYMMETRY_THRESHOLD    = 0.38
ARTIFACT_OVERRIDE_THR = 0.55

# Per-class display info shown to end user
CLASS_INFO = {
    "arrowheads": {
        "icon": "🏹",
        "label": "Projectile Point / Arrowhead",
        "color": "#c0392b",
        "description": (
            "A hafted projectile point used as a weapon tip for arrows, darts, or spears. "
            "Produced by skilled knapping (chipping stone with antler or bone tools)."
        ),
        "what_to_look_for": [
            "Bifacial flaking visible on BOTH flat faces",
            "Thinned cross-section — blade is thinner than the base",
            "Stem, notch, or ear at the base for tying to a shaft (hafting)",
            "Bilateral symmetry — left and right sides mirror each other",
            "Regular, intentional serrations along the cutting edge",
            "Conchoidal (shell-like) fracture surface — smooth, glassy ripples",
        ],
        "common_materials": ["Chert", "Flint", "Obsidian", "Quartzite", "Jasper"],
        "typical_size": "2 – 10 cm long",
        "period": "Paleo-Indian through Late Prehistoric (13,000 BP – 500 BP)",
    },
    "scrapers": {
        "icon": "🪨",
        "label": "Scraper",
        "color": "#2980b9",
        "description": (
            "A stone tool used for scraping hides, wood, or plant material. "
            "Distinguished from projectile points by its thick, steep working edge on ONE face only."
        ),
        "what_to_look_for": [
            "Unifacial working — only ONE face shows intentional flaking",
            "Steep retouch angle on working edge (~60–80°)",
            "Cortex (original rock surface) often present on the dorsal face",
            "Thick, robust form — not thinned for hafting",
            "No stem or notching — not designed to be mounted on a shaft",
            "Wear polish visible on working edge under magnification",
        ],
        "common_materials": ["Chert", "Flint", "Quartzite", "Basalt"],
        "typical_size": "3 – 12 cm",
        "period": "Found across all prehistoric periods",
    },
    "axehammerstones": {
        "icon": "🪓",
        "label": "Axe / Hammerstone",
        "color": "#8e44ad",
        "description": (
            "A ground stone tool used for chopping, grinding, or pounding. "
            "Unlike knapped tools, these are shaped by abrasion and pecking rather than flaking."
        ),
        "what_to_look_for": [
            "Pecking marks — small circular pits from repeated impact use",
            "Ground or polished facets on working surfaces",
            "Heavier and more rounded form than flaked tools",
            "Battered or crushed areas at impact points",
            "May show groove for hafting (grooved axes)",
            "Surface smoothed by use — not from knapping",
        ],
        "common_materials": ["Granite", "Basalt", "Sandstone", "Quartzite"],
        "typical_size": "5 – 25 cm",
        "period": "Archaic through Woodland periods (8,000 BP – 1,000 BP)",
    },
    "notartifact": {
        "icon": "🪨",
        "label": "Natural Rock / Not an Artifact",
        "color": "#7f8c8d",
        "description": (
            "This appears to be a naturally occurring rock, unmodified stone, or debitage "
            "(waste flake from tool production) rather than an intentionally shaped artifact."
        ),
        "what_to_look_for": [
            "No systematic flaking pattern — fractures appear random",
            "No bilateral symmetry or intentional shaping",
            "No stem, notch, or hafting preparation",
            "No organised working edge",
            "Surface texture consistent with natural weathering",
            "Consider: may be debitage (waste from knapping nearby tools)",
        ],
        "common_materials": ["Any"],
        "typical_size": "Any",
        "period": "N/A",
    },
}


# ═══════════════════════════════════════════════════════════════
# LOAD MODEL
# ═══════════════════════════════════════════════════════════════

print("Loading model …")
model = keras.models.load_model(MODEL_PATH, compile=False)
print("Model loaded ✅")

# Try to get class names from model if available
try:
    detected_classes = model.output_shape[-1]
    print(f"Model output classes: {detected_classes}")
except:
    pass


# ═══════════════════════════════════════════════════════════════
# FEATURE EXTRACTION
# ═══════════════════════════════════════════════════════════════

def compute_laplacian_depth_channel(img_np):
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).astype(np.float32)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    lap  = np.abs(cv2.Laplacian(blur, cv2.CV_32F))
    lap  = (lap - lap.min()) / (lap.max() - lap.min() + 1e-8)
    return lap[:, :, np.newaxis]


def compute_gabor_texture_channel(img_np):
    gray  = sk_rgb2gray(img_np).astype(np.float64)
    resps = [np.abs(gabor(gray, frequency=0.15, theta=t)[0]) for t in [0, np.pi/4, np.pi/2]]
    gab   = np.mean(resps, axis=0).astype(np.float32)
    gab   = (gab - gab.min()) / (gab.max() - gab.min() + 1e-8)
    return gab[:, :, np.newaxis]


def build_5channel_input(img_np):
    rgb   = img_np.astype(np.float32) / 255.0
    depth = compute_laplacian_depth_channel(img_np)
    gab   = compute_gabor_texture_channel(img_np)
    return np.concatenate([rgb, depth, gab], axis=-1)


def texture_complexity_score(img_np):
    return float(cv2.Laplacian(cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY), cv2.CV_64F).var())


def edge_regularity_score(img_np):
    edges = cv2.Canny(cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY), 50, 150)
    ys, xs = np.where(edges > 0)
    if len(xs) < 20:
        return 0.0
    pts  = np.stack([xs, ys], axis=1).astype(np.float32)
    tree = cKDTree(pts)
    d, _ = tree.query(pts, k=2)
    return float(np.clip(1.0 / (1.0 + np.std(d[:, 1])), 0, 1))


def symmetry_score(img_np):
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY).astype(np.float32)
    h, w = gray.shape
    l = gray[:, :w//2]
    r = gray[:, w//2:][:, ::-1]
    m = min(l.shape[1], r.shape[1])
    corr = np.corrcoef(l[:, :m].flatten(), r[:, :m].flatten())[0, 1]
    return float(np.clip((corr + 1) / 2, 0, 1))


def stem_notch_score(img_np):
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    _, bw = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    h, w  = bw.shape
    def row_width(region):
        widths = [cols[-1]-cols[0] for row in region for cols in [np.where(row>0)[0]] if len(cols)>2]
        return np.mean(widths) if widths else 0
    fw = row_width(bw)
    bw_ = row_width(bw[int(h*0.75):, :])
    return float(np.clip(1.0 - (bw_ / fw), 0, 1)) if fw > 1 else 0.0


def is_broken_artifact(tex, reg, sym):
    return tex > 80 and reg > 0.55 and sym < 0.70


# ═══════════════════════════════════════════════════════════════
# GRAD-CAM
# ═══════════════════════════════════════════════════════════════

def make_gradcam(img_5ch):
    try:
        gap_idx = next(i for i, l in enumerate(model.layers)
                       if isinstance(l, keras.layers.GlobalAveragePooling2D))
        feat_layer = model.layers[gap_idx - 1]
        grad_model = keras.models.Model(
            inputs=model.inputs,
            outputs=[feat_layer.output, model.output]
        )
        img_t = tf.cast(np.expand_dims(img_5ch, 0), tf.float32)
        with tf.GradientTape() as tape:
            conv_out, preds = grad_model(img_t)
            idx   = tf.argmax(preds[0])
            score = preds[:, idx]
        grads   = tape.gradient(score, conv_out)[0]
        weights = tf.reduce_mean(grads, axis=(0, 1))
        cam     = tf.reduce_sum(tf.multiply(weights, conv_out[0]), axis=-1)
        cam     = np.maximum(cam.numpy(), 0)
        cam     = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        return cam
    except:
        return None


def overlay_gradcam(img_np, cam, alpha=0.45):
    hm  = cv2.resize(cam, (img_np.shape[1], img_np.shape[0]))
    hm  = cv2.applyColorMap(np.uint8(255 * hm), cv2.COLORMAP_JET)
    hm  = cv2.cvtColor(hm, cv2.COLOR_BGR2RGB)
    return cv2.addWeighted(img_np, 1 - alpha, hm, alpha, 0)


# ═══════════════════════════════════════════════════════════════
# VISUALISATION PANELS
# ═══════════════════════════════════════════════════════════════

def make_feature_panel(img_np, x5, tex, reg, sym, stem):
    """
    Returns a matplotlib figure showing:
      Original | Depth channel | Gabor texture | Edge map
    """
    fig, axes = plt.subplots(1, 4, figsize=(16, 4))
    fig.patch.set_facecolor("#1a1a2e")

    titles = ["Original Photo", "Depth / Relief\n(Laplacian)", "Texture / Knapping\n(Gabor)", "Edge Regularity\n(Canny)"]
    imgs   = [
        img_np,
        x5[:, :, 3],
        x5[:, :, 4],
        cv2.Canny(cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY), 50, 150),
    ]
    cmaps = [None, "hot", "viridis", "gray"]

    for ax, title, im, cmap in zip(axes, titles, imgs, cmaps):
        ax.imshow(im, cmap=cmap)
        ax.set_title(title, color="white", fontsize=10, pad=6)
        ax.axis("off")

    # Score bar at the bottom
    fig.text(0.5, 0.01,
             f"Texture: {tex:.0f}   Edge Regularity: {reg:.2f}   Symmetry: {sym:.2f}   Stem/Notch: {stem:.2f}",
             ha="center", color="#aaaaaa", fontsize=9)

    plt.tight_layout(rect=[0, 0.06, 1, 1])
    return fig


def make_confidence_chart(preds, class_names, top_cls, override):
    """
    Horizontal bar chart of per-class confidence scores.
    """
    info    = CLASS_INFO
    colors  = [info.get(c, {}).get("color", "#888888") for c in class_names]
    labels  = [info.get(c, {}).get("label", c) for c in class_names]
    icons   = [info.get(c, {}).get("icon", "") for c in class_names]

    fig, ax = plt.subplots(figsize=(8, 3))
    fig.patch.set_facecolor("#1a1a2e")
    ax.set_facecolor("#16213e")

    y_pos = range(len(class_names))
    bars  = ax.barh(list(y_pos), preds, color=colors, height=0.55, edgecolor="none")

    # Value labels on bars
    for bar, val in zip(bars, preds):
        ax.text(min(val + 0.02, 0.98), bar.get_y() + bar.get_height() / 2,
                f"{val:.1%}", va="center", color="white", fontsize=10, fontweight="bold")

    ax.set_yticks(list(y_pos))
    ax.set_yticklabels([f"{ic} {lb}" for ic, lb in zip(icons, labels)], color="white", fontsize=9)
    ax.set_xlim(0, 1.15)
    ax.set_xlabel("Confidence", color="#aaaaaa", fontsize=9)
    ax.tick_params(colors="white")
    for spine in ax.spines.values():
        spine.set_visible(False)

    title = f"Model Confidence Scores"
    if override:
        title += "  ⚡ Override applied"
    ax.set_title(title, color="white", fontsize=11, pad=8)

    plt.tight_layout()
    return fig


def fig_to_pil(fig):
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=120, bbox_inches="tight")
    buf.seek(0)
    img = Image.open(buf).copy()
    plt.close(fig)
    return img


# ═══════════════════════════════════════════════════════════════
# MAIN PREDICTION FUNCTION
# ═══════════════════════════════════════════════════════════════

def predict_artifact(image):
    """
    Main Gradio prediction function.
    Returns:
      - result_html     : full formatted HTML result card
      - gradcam_img     : PIL image of Grad-CAM overlay
      - feature_img     : PIL image of 4-panel feature analysis
      - confidence_img  : PIL image of confidence bar chart
    """
    if image is None:
        return "<p style='color:red'>Please upload an image.</p>", None, None, None

    img_np = np.array(image.convert("RGB").resize(IMG_SIZE))

    # ── Feature scores ────────────────────────────────────────
    tex_score  = texture_complexity_score(img_np)
    reg_score  = edge_regularity_score(img_np)
    sym_score  = symmetry_score(img_np)
    stem_score = stem_notch_score(img_np)
    broken     = is_broken_artifact(tex_score, reg_score, sym_score)

    # ── Gate 1: Realism check ─────────────────────────────────
    if tex_score < TEXTURE_THRESHOLD:
        html = _format_rejection(
            "Not a Real Stone Photo",
            "The image surface appears too uniform. Please upload a real photograph of a stone object.",
            tex_score, reg_score, sym_score, stem_score
        )
        return html, None, None, None

    # ── Model prediction ──────────────────────────────────────
    x5    = build_5channel_input(img_np)
    preds = model.predict(np.expand_dims(x5, 0), verbose=0)[0]

    top_idx  = int(np.argmax(preds))
    top_cls  = CLASS_NAMES[top_idx]
    top_conf = float(preds[top_idx])
    ent      = float(-np.sum(preds * np.log(preds + 1e-8)))

    notartifact_idx    = next((i for i, c in enumerate(CLASS_NAMES)
                               if "notartifact" in c.lower()), None)
    artifact_score_sum = sum(float(preds[i]) for i in range(len(CLASS_NAMES))
                             if i != notartifact_idx)

    # ── Archaeological override ───────────────────────────────
    override = False
    predicted_notartifact = (notartifact_idx is not None and
                             preds[notartifact_idx] == max(preds))

    if predicted_notartifact:
        if tex_score > 80 and reg_score > 0.55 and artifact_score_sum > ARTIFACT_OVERRIDE_THR:
            override = True
            art_scores = {c: float(preds[i]) for i, c in enumerate(CLASS_NAMES)
                          if i != notartifact_idx}
            top_cls    = max(art_scores, key=art_scores.get)
            top_conf   = art_scores[top_cls]

    # ── Gate 4: Confidence check ──────────────────────────────
    uncertain = (not override) and (top_conf < CONF_THRESHOLD or ent > ENTROPY_THRESHOLD)

    # ── Grad-CAM ──────────────────────────────────────────────
    cam     = make_gradcam(x5)
    gradcam_pil = None
    if cam is not None:
        overlay    = overlay_gradcam(img_np, cam)
        fig_gc, ax = plt.subplots(figsize=(5, 5))
        fig_gc.patch.set_facecolor("#1a1a2e")
        ax.imshow(overlay)
        ax.set_title("Grad-CAM — What the model focused on", color="white", fontsize=10)
        ax.axis("off")
        plt.tight_layout()
        gradcam_pil = fig_to_pil(fig_gc)

    # ── Feature panel ─────────────────────────────────────────
    feat_pil = fig_to_pil(make_feature_panel(img_np, x5, tex_score, reg_score, sym_score, stem_score))

    # ── Confidence chart ──────────────────────────────────────
    conf_pil = fig_to_pil(make_confidence_chart(preds, CLASS_NAMES, top_cls, override))

    # ── Build result HTML ─────────────────────────────────────
    result_html = _format_result(
        top_cls, top_conf, ent, preds,
        tex_score, reg_score, sym_score, stem_score,
        broken, override, uncertain, artifact_score_sum
    )

    return result_html, gradcam_pil, feat_pil, conf_pil


# ═══════════════════════════════════════════════════════════════
# HTML FORMATTERS
# ═══════════════════════════════════════════════════════════════

def _gate_badge(passed, label):
    color  = "#27ae60" if passed else "#e74c3c"
    symbol = "✅" if passed else "❌"
    return f'<span style="background:{color};color:white;padding:2px 8px;border-radius:10px;font-size:12px;margin:2px">{symbol} {label}</span>'


def _format_result(top_cls, top_conf, ent, preds,
                   tex, reg, sym, stem,
                   broken, override, uncertain, art_sum):

    info  = CLASS_INFO.get(top_cls, CLASS_INFO["notartifact"])
    color = info["color"]
    icon  = info["icon"]

    # Status banner
    if uncertain:
        status_bg    = "#7f8c8d"
        status_text  = "⚠️ UNCERTAIN — Manual review recommended"
        status_sub   = "The model is not confident. This may be a broken, stained, or unusual specimen."
    elif top_cls == "notartifact":
        status_bg    = "#7f8c8d"
        status_text  = "🪨 NOT AN ARTIFACT"
        status_sub   = "This appears to be a natural rock, debitage, or unmodified stone."
    else:
        status_bg    = color
        status_text  = f"{icon} {info['label'].upper()}"
        status_sub   = f"Confidence: {top_conf:.1%}"

    # Build checklist
    checklist_html = "".join(
        f'<li style="margin:4px 0;color:#ddd">✦ {item}</li>'
        for item in info["what_to_look_for"]
    )

    # Gate summary
    gate1 = tex > TEXTURE_THRESHOLD
    gate2 = reg > REGULARITY_THRESHOLD
    gate3 = sym > SYMMETRY_THRESHOLD
    gate4 = top_conf > CONF_THRESHOLD and ent < ENTROPY_THRESHOLD

    gates_html = (
        _gate_badge(gate1, f"Texture {tex:.0f}") +
        _gate_badge(gate2, f"Edge Regularity {reg:.2f}") +
        _gate_badge(gate3, f"Symmetry {sym:.2f}") +
        _gate_badge(gate4, f"Confidence {top_conf:.1%}") +
        (f'<span style="background:#e67e22;color:white;padding:2px 8px;border-radius:10px;font-size:12px;margin:2px">⚡ Override Applied</span>' if override else "") +
        (f'<span style="background:#e67e22;color:white;padding:2px 8px;border-radius:10px;font-size:12px;margin:2px">🔍 Fragment Detected</span>' if broken else "")
    )

    # Materials & metadata
    materials = ", ".join(info["common_materials"])

    html = f"""
<div style="font-family:'Segoe UI',sans-serif;background:#1a1a2e;border-radius:16px;overflow:hidden;max-width:820px">
  <!-- Header banner -->
  <div style="background:{status_bg};padding:20px 24px">
    <div style="font-size:26px;font-weight:700;color:white;letter-spacing:1px">{status_text}</div>
    <div style="font-size:14px;color:rgba(255,255,255,0.85);margin-top:4px">{status_sub}</div>
  </div>
  <!-- Body -->
  <div style="padding:20px 24px">
    <!-- Archaeological gates -->
    <div style="margin-bottom:18px">
      <div style="color:#aaa;font-size:11px;text-transform:uppercase;letter-spacing:1px;margin-bottom:8px">Archaeological Feature Gates</div>
      <div>{gates_html}</div>
    </div>
    <!-- Description -->
    <div style="background:#16213e;border-radius:10px;padding:14px 16px;margin-bottom:16px">
      <div style="color:{color};font-weight:600;font-size:14px;margin-bottom:6px">{icon} What is this?</div>
      <div style="color:#ccc;font-size:13px;line-height:1.6">{info['description']}</div>
    </div>
    <!-- Two columns: checklist + diagnostics -->
    <div style="display:grid;grid-template-columns:1fr 1fr;gap:16px;margin-bottom:16px">
      <!-- What to look for -->
      <div style="background:#16213e;border-radius:10px;padding:14px 16px">
        <div style="color:{color};font-weight:600;font-size:13px;margin-bottom:8px">🔍 What to look for</div>
        <ul style="margin:0;padding-left:4px;list-style:none;font-size:12px">
          {checklist_html}
        </ul>
      </div>
      <!-- Diagnostic scores -->
      <div style="background:#16213e;border-radius:10px;padding:14px 16px">
        <div style="color:{color};font-weight:600;font-size:13px;margin-bottom:10px">📊 Diagnostic Scores</div>
        <div style="margin-bottom:8px">
          <div style="display:flex;justify-content:space-between;color:#aaa;font-size:11px"><span>Texture Complexity</span><span style="color:white">{tex:.0f}</span></div>
          <div style="background:#0f3460;border-radius:4px;height:6px;margin-top:3px">
            <div style="background:{color};height:6px;border-radius:4px;width:{min(tex/500*100,100):.0f}%"></div>
          </div>
          <div style="color:#666;font-size:10px;margin-top:1px">&gt;80 = worked stone surface</div>
        </div>
        <div style="margin-bottom:8px">
          <div style="display:flex;justify-content:space-between;color:#aaa;font-size:11px"><span>Edge Regularity</span><span style="color:white">{reg:.2f}</span></div>
          <div style="background:#0f3460;border-radius:4px;height:6px;margin-top:3px">
            <div style="background:{color};height:6px;border-radius:4px;width:{reg*100:.0f}%"></div>
          </div>
          <div style="color:#666;font-size:10px;margin-top:1px">&gt;0.55 = systematic knapping edges</div>
        </div>
        <div style="margin-bottom:8px">
          <div style="display:flex;justify-content:space-between;color:#aaa;font-size:11px"><span>Bilateral Symmetry</span><span style="color:white">{sym:.2f}</span></div>
          <div style="background:#0f3460;border-radius:4px;height:6px;margin-top:3px">
            <div style="background:{color};height:6px;border-radius:4px;width:{sym*100:.0f}%"></div>
          </div>
          <div style="color:#666;font-size:10px;margin-top:1px">&gt;0.70 = intact hafted point</div>
        </div>
        <div style="margin-bottom:4px">
          <div style="display:flex;justify-content:space-between;color:#aaa;font-size:11px"><span>Stem / Notch Evidence</span><span style="color:white">{stem:.2f}</span></div>
          <div style="background:#0f3460;border-radius:4px;height:6px;margin-top:3px">
            <div style="background:{color};height:6px;border-radius:4px;width:{stem*100:.0f}%"></div>
          </div>
          <div style="color:#666;font-size:10px;margin-top:1px">&gt;0.25 = base constriction for hafting</div>
        </div>
      </div>
    </div>
    <!-- Material & period metadata -->
    <div style="background:#16213e;border-radius:10px;padding:12px 16px;display:grid;grid-template-columns:1fr 1fr 1fr;gap:10px">
      <div>
        <div style="color:#666;font-size:10px;text-transform:uppercase;letter-spacing:1px">Common Materials</div>
        <div style="color:#ddd;font-size:12px;margin-top:3px">{materials}</div>
      </div>
      <div>
        <div style="color:#666;font-size:10px;text-transform:uppercase;letter-spacing:1px">Typical Size</div>
        <div style="color:#ddd;font-size:12px;margin-top:3px">{info['typical_size']}</div>
      </div>
      <div>
        <div style="color:#666;font-size:10px;text-transform:uppercase;letter-spacing:1px">Time Period</div>
        <div style="color:#ddd;font-size:12px;margin-top:3px">{info['period']}</div>
      </div>
    </div>
    <!-- Disclaimer -->
    <div style="margin-top:14px;padding:10px 14px;background:#0f3460;border-radius:8px;border-left:3px solid {color}">
      <div style="color:#aaa;font-size:11px;line-height:1.5">
        ⚠️ <strong style="color:#ddd">Disclaimer:</strong> This tool provides an AI-assisted preliminary classification only.
        Always consult a certified archaeologist or lithic specialist for definitive identification.
        Do not disturb, collect, or move artifacts found in the field without proper permits.
      </div>
    </div>
  </div>
</div>
"""
    return html


def _format_rejection(title, message, tex, reg, sym, stem):
    return f"""
<div style="font-family:'Segoe UI',sans-serif;background:#1a1a2e;border-radius:16px;overflow:hidden;max-width:820px">
  <div style="background:#c0392b;padding:20px 24px">
    <div style="font-size:22px;font-weight:700;color:white">❌ {title}</div>
  </div>
  <div style="padding:20px 24px;color:#ccc;font-size:14px">
    <p>{message}</p>
    <p style="margin-top:12px;color:#888;font-size:12px">
      Texture score: {tex:.1f} (minimum required: {TEXTURE_THRESHOLD}) |
      Edge regularity: {reg:.2f} | Symmetry: {sym:.2f} | Stem: {stem:.2f}
    </p>
  </div>
</div>
"""


# ═══════════════════════════════════════════════════════════════
# GRADIO UI
# ═══════════════════════════════════════════════════════════════

TITLE = "🏹 Unearthed — Archaeological Lithic Artifact Classifier"

DESCRIPTION = """
Upload a photograph of a stone object to identify whether it is a **projectile point (arrowhead)**,
**scraper**, **axe/hammerstone**, or a **natural rock / debitage**.
The classifier uses texture analysis, depth mapping, edge regularity, and bilateral symmetry —
the same visual criteria used by field archaeologists.
**Tips for best results:**
- Photograph the artifact against a plain, contrasting background (white or black)
- Include both faces if possible
- Use natural or diffuse lighting to show surface texture clearly
- Avoid harsh shadows across the face of the artifact
"""

with gr.Blocks(
    title=TITLE,
    theme=gr.themes.Base(
        primary_hue="orange",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
    ),
    css="""
        .gradio-container { max-width: 1100px !important; }
        #result-panel { min-height: 200px; }
        footer { display: none !important; }
    """
) as demo:

    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        # Left column: upload
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="Upload artifact photo",
                height=320,
                sources=["upload", "clipboard"],
            )
            submit_btn = gr.Button("🔍 Analyze Artifact", variant="primary", size="lg")

            gr.Markdown("""
---
**How it works**
| Feature | What it detects |
|---|---|
| 🌡️ Texture | Knapping scar density |
| 📐 Depth map | Biface ridge / relief |
| 〰️ Edge regularity | Systematic vs random edges |
| ↔️ Symmetry | Hafting geometry |
| 🪓 Stem/notch | Base constriction for shaft |
""")

        # Right column: results
        with gr.Column(scale=2):
            result_html = gr.HTML(
                label="Identification Result",
                elem_id="result-panel"
            )

    gr.Markdown("---")
    gr.Markdown("### 📊 Analysis Visualizations")

    with gr.Row():
        conf_chart  = gr.Image(label="Confidence Scores", height=240)
        gradcam_out = gr.Image(label="Grad-CAM — Model Focus Area", height=240)

    feature_out = gr.Image(label="Feature Analysis — Depth · Texture · Edges", height=260)

    # Examples
    gr.Markdown("---")
    gr.Markdown("### 📷 Try an example")
    gr.Examples(
        examples=[],   # add local image paths here once you have them e.g. ["examples/arrowhead1.jpg"]
        inputs=image_input,
        label="Example images (add your own to the /examples folder)"
    )

    # Footer
    gr.Markdown("""
---
<div style='text-align:center;color:#666;font-size:12px'>
Built with EfficientNetB3 · Gabor texture analysis · Laplacian depth mapping<br>
Grounded in criteria from <a href='https://projectilepoints.net' target='_blank'>projectilepoints.net</a>,
<a href='https://cartarchaeology.wordpress.com' target='_blank'>CART Archaeology</a>,
and <a href='https://rockhoundresource.com' target='_blank'>Rockhound Resource</a><br>
⚠️ For educational use only — consult a certified archaeologist for definitive identification
</div>
""")

    submit_btn.click(
        fn=predict_artifact,
        inputs=[image_input],
        outputs=[result_html, gradcam_out, feature_out, conf_chart],
    )

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