src/app.py

import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
import tempfile, os, time
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

# ── Page config ───────────────────────────────────────────────
st.set_page_config(
    page_title="ShopGuard AI",
    page_icon="🛡️",
    layout="wide"
)

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Share+Tech+Mono&family=Rajdhani:wght@400;600;700&display=swap');

    .stApp {
        background-color: #080c10;
        color: #c9d1d9;
        font-family: 'Rajdhani', sans-serif;
    }
    .block-container { padding-top: 1.5rem; max-width: 1200px; }

    /* Header */
    .header-wrap {
        border-bottom: 1px solid #21262d;
        padding-bottom: 1rem;
        margin-bottom: 1.5rem;
    }
    .header-title {
        font-family: 'Share Tech Mono', monospace;
        font-size: 2rem;
        color: #58a6ff;
        letter-spacing: 0.05em;
        margin: 0;
    }
    .header-sub {
        color: #6e7681;
        font-size: 0.85rem;
        font-family: 'Share Tech Mono', monospace;
        margin-top: 0.2rem;
    }

    /* Panel cards */
    .panel {
        background: #0d1117;
        border: 1px solid #21262d;
        border-radius: 8px;
        padding: 1.2rem 1.4rem;
        margin-bottom: 1rem;
    }
    .panel-title {
        font-size: 0.7rem;
        letter-spacing: 0.15em;
        color: #6e7681;
        text-transform: uppercase;
        font-family: 'Share Tech Mono', monospace;
        margin-bottom: 0.8rem;
    }

    /* Result cards */
    .result-shoplifting {
        background: #1a0a0a;
        border: 1px solid #f85149;
        border-radius: 8px;
        padding: 1.5rem;
        text-align: center;
    }
    .result-normal {
        background: #0a1a0e;
        border: 1px solid #3fb950;
        border-radius: 8px;
        padding: 1.5rem;
        text-align: center;
    }
    .result-label-shop {
        font-family: 'Share Tech Mono', monospace;
        font-size: 2.2rem;
        color: #f85149;
        letter-spacing: 0.1em;
    }
    .result-label-norm {
        font-family: 'Share Tech Mono', monospace;
        font-size: 2.2rem;
        color: #3fb950;
        letter-spacing: 0.1em;
    }
    .result-conf {
        font-family: 'Share Tech Mono', monospace;
        font-size: 1.1rem;
        color: #e3b341;
        margin-top: 0.5rem;
    }
    .result-meta {
        color: #6e7681;
        font-size: 0.78rem;
        font-family: 'Share Tech Mono', monospace;
        margin-top: 0.4rem;
    }

    /* Prob bar */
    .prob-bar-bg {
        background: #161b22;
        border: 1px solid #21262d;
        border-radius: 6px;
        height: 24px;
        width: 100%;
        overflow: hidden;
        margin-top: 0.8rem;
    }

    /* Model badge */
    .model-badge {
        display: inline-block;
        background: #1f2937;
        border: 1px solid #374151;
        border-radius: 4px;
        padding: 2px 8px;
        font-family: 'Share Tech Mono', monospace;
        font-size: 0.75rem;
        color: #58a6ff;
        margin-bottom: 0.5rem;
    }

    /* Streamlit overrides */
    .stSelectbox label, .stSlider label, .stFileUploader label {
        color: #6e7681 !important;
        font-family: 'Share Tech Mono', monospace !important;
        font-size: 0.75rem !important;
        text-transform: uppercase;
        letter-spacing: 0.1em;
    }
    .stButton > button {
        background: #1f6feb;
        color: white;
        border: none;
        border-radius: 6px;
        font-family: 'Share Tech Mono', monospace;
        font-size: 0.9rem;
        letter-spacing: 0.05em;
        width: 100%;
        padding: 0.6rem;
        transition: background 0.2s;
    }
    .stButton > button:hover { background: #388bfd; }
    div[data-testid="stMetricValue"] {
        font-family: 'Share Tech Mono', monospace;
        color: #58a6ff;
    }
    .stSpinner > div { border-top-color: #58a6ff !important; }
    hr { border-color: #21262d; }
</style>
""", unsafe_allow_html=True)

# ── Config ────────────────────────────────────────────────────
FRAMES_PER_VIDEO = 16
IMG_SIZE         = 224
PERSON_CLASS     = 0
YOLO_CONF        = 0.3
PAD              = 0.10

MODEL_CONFIGS = {
    "Model A — General": {
        "repo_id":   "higsboson/shoplifting_exp_a",
        "filename":  "shoplifting_a.keras",
        "default_threshold": 0.50,
        "label": "A"
    },
    "Model B — Kitchen": {
        "repo_id":   "higsboson/shoplifting_exp_b",
        "filename":  "best_model.keras",
        "default_threshold": 0.50,
        "label": "B"
    },
    "Model C — Lab": {
        "repo_id":   "higsboson/shoplifting_exp_c",
        "filename":  "shoplifting_exp_c.keras",
        "default_threshold": 0.50,
        "label": "C"
    },
}

# ── Loaders ───────────────────────────────────────────────────
@st.cache_resource
def load_yolo():
    return YOLO("yolo11n.pt")

@st.cache_resource
def load_mobilenet():
    base = tf.keras.applications.MobileNetV2(
        input_shape=(IMG_SIZE, IMG_SIZE, 3),
        include_top=False, pooling="avg", weights="imagenet"
    )
    base.trainable = False
    return base

@st.cache_resource
def load_lstm(repo_id, filename):
    path = hf_hub_download(repo_id=repo_id, filename=filename)
    return tf.keras.models.load_model(path)

# ── Pipeline ──────────────────────────────────────────────────
def extract_frames(video_path, n=FRAMES_PER_VIDEO):
    cap   = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    idxs  = np.linspace(0, max(total - 1, 0), n, dtype=int)
    frames = {}
    for idx in idxs:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ret, frame = cap.read()
        if ret:
            frames[idx] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    cap.release()
    return frames, idxs

def crop_person(frame, yolo_model, last_box):
    h, w = frame.shape[:2]
    results = yolo_model(frame, conf=YOLO_CONF, classes=[PERSON_CLASS], verbose=False)
    boxes = results[0].boxes
    if boxes is not None and len(boxes):
        best = max(boxes, key=lambda b: b.conf.item())
        x1, y1, x2, y2 = map(int, best.xyxy[0].tolist())
        pw, ph = x2 - x1, y2 - y1
        x1 = max(0, x1 - int(PAD * pw))
        y1 = max(0, y1 - int(PAD * ph))
        x2 = min(w, x2 + int(PAD * pw))
        y2 = min(h, y2 + int(PAD * ph))
        last_box[0] = (x1, y1, x2, y2)
    if last_box[0]:
        x1, y1, x2, y2 = last_box[0]
        crop = frame[y1:y2, x1:x2]
    else:
        crop = frame
    return cv2.resize(crop, (IMG_SIZE, IMG_SIZE))

def run_inference(video_path, yolo_model, mobilenet, lstm_model, threshold):
    frames_dict, idxs = extract_frames(video_path)
    last_box = [None]
    crops    = []
    for idx in idxs:
        frame = frames_dict.get(idx, np.zeros((IMG_SIZE, IMG_SIZE, 3), np.uint8))
        crops.append(crop_person(frame, yolo_model, last_box))
    crops_arr = np.array(crops, dtype=np.float32)
    crops_pp  = tf.keras.applications.mobilenet_v2.preprocess_input(crops_arr)
    features  = mobilenet.predict(crops_pp, verbose=0)
    features  = features[np.newaxis, ...]
    prob      = lstm_model.predict(features, verbose=0)[0][0]
    label     = "SHOPLIFTING" if prob >= threshold else "NORMAL"
    return float(prob), label, crops

# ── Header ────────────────────────────────────────────────────
st.markdown("""
<div class="header-wrap">
    <p class="header-title">🛡️ SHOPGUARD AI</p>
    <p class="header-sub">YOLO11n → MobileNetV2 → Attention LSTM &nbsp;|&nbsp; FYP Demo System</p>
</div>
""", unsafe_allow_html=True)

# ── Layout: Left config | Right result ───────────────────────
col_left, col_right = st.columns([1, 1.6], gap="large")

with col_left:
    st.markdown('<div class="panel-title">⚙ Model Configuration</div>', unsafe_allow_html=True)

    model_choice = st.selectbox(
        "Select Model",
        list(MODEL_CONFIGS.keys()),
        help="Choose which trained model to run inference with"
    )
    cfg = MODEL_CONFIGS[model_choice]

    st.markdown(f'<div class="model-badge">HF: {cfg["repo_id"]}</div>', unsafe_allow_html=True)

    threshold = st.slider(
        "Decision Threshold",
        min_value=0.0,
        max_value=1.0,
        value=cfg["default_threshold"],
        step=0.01,
        help="Probability above this = Shoplifting. Adjust per your validation results."
    )

    st.caption(f"ℹ️ Prob ≥ {threshold:.2f} → 🚨 Shoplifting | Prob < {threshold:.2f} → ✅ Normal")

    st.divider()
    st.markdown('<div class="panel-title">📹 Video Input</div>', unsafe_allow_html=True)

    uploaded = st.file_uploader(
        "Upload Video",
        type=["mp4", "avi", "mov", "mkv"],
        help="Short clips (5–30s) work best"
    )

    run_btn = st.button("🔍 Run Inference", disabled=(uploaded is None))

with col_right:
    if uploaded is None:
        st.markdown("""
        <div style="border: 1px dashed #21262d; border-radius: 8px;
                    padding: 3rem; text-align: center; color: #6e7681;
                    font-family: 'Share Tech Mono', monospace; font-size: 0.85rem;">
            Upload a video on the left<br>and click Run Inference
        </div>
        """, unsafe_allow_html=True)
    else:
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
            tmp.write(uploaded.read())
            tmp_path = tmp.name

        st.video(tmp_path)

        if run_btn:
            with st.spinner("Loading YOLO + MobileNetV2..."):
                yolo      = load_yolo()
                mobilenet = load_mobilenet()

            with st.spinner(f"Downloading model from HuggingFace ({cfg['repo_id']})..."):
                lstm = load_lstm(cfg["repo_id"], cfg["filename"])

            with st.spinner("Running pipeline: frame extraction → YOLO crop → feature extraction → LSTM..."):
                t0 = time.time()
                prob, label, crops = run_inference(tmp_path, yolo, mobilenet, lstm, threshold)
                elapsed = time.time() - t0

            os.unlink(tmp_path)

            st.divider()

            # Result card
            is_shop = label == "SHOPLIFTING"
            card_cls = "result-shoplifting" if is_shop else "result-normal"
            lbl_cls  = "result-label-shop"  if is_shop else "result-label-norm"
            icon     = "🚨" if is_shop else "✅"
            bar_color = "#f85149" if is_shop else "#3fb950"

            st.markdown(f"""
            <div class="{card_cls}">
                <div class="{lbl_cls}">{icon} {label}</div>
                <div class="result-conf">Confidence: {prob:.4f}</div>
                <div class="result-meta">
                    Model {cfg['label']} &nbsp;|&nbsp; Threshold: {threshold:.2f}
                    &nbsp;|&nbsp; Inference: {elapsed:.2f}s
                </div>
                <div class="prob-bar-bg">
                    <div style="background:{bar_color}; width:{prob*100:.1f}%;
                                height:100%; display:flex; align-items:center;
                                padding-left:8px; color:#080c10;
                                font-family:'Share Tech Mono',monospace;
                                font-size:0.78rem; font-weight:bold;">
                        {prob*100:.1f}%
                    </div>
                </div>
            </div>
            """, unsafe_allow_html=True)

            # Metrics row
            st.divider()
            m1, m2, m3 = st.columns(3)
            m1.metric("Probability", f"{prob:.4f}")
            m2.metric("Threshold",   f"{threshold:.2f}")
            m3.metric("Inference",   f"{elapsed:.2f}s")

            # Sampled crops
            st.markdown('<div class="panel-title" style="margin-top:1rem;">🎞 YOLO-Cropped Frames</div>',
                        unsafe_allow_html=True)
            cols = st.columns(8)
            for i, crop in enumerate(crops[:8]):
                cols[i].image(crop, use_container_width=True, caption=f"f{i+1}")