import streamlit as st
import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import EfficientNet_B0_Weights
from PIL import Image
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
import json
from pathlib import Path
import os
# ── Page config ─────────────────────────────────────────────────────────────
st.set_page_config(
page_title="WildEye",
page_icon="🦌",
layout="wide",
initial_sidebar_state="expanded",
)
# ── Custom CSS ───────────────────────────────────────────────────────────────
st.markdown("""
""", unsafe_allow_html=True)
# ── Constants ─────────────────────────────────────────────────────────────────
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
IMAGE_SIZE = 224
CLASSES = ['butterfly', 'cat', 'chicken', 'cow', 'dog',
'elephant', 'horse', 'sheep', 'spider', 'squirrel']
CLASS_EMOJI = {
'butterfly': '🦋', 'cat': '🐱', 'chicken': '🐔', 'cow': '🐄',
'dog': '🐶', 'elephant': '🐘', 'horse': '🐎', 'sheep': '🐑',
'spider': '🕷️', 'squirrel': '🐿️',
}
# ── Model ─────────────────────────────────────────────────────────────────────
@st.cache_resource
def load_model():
"""Load the domain-augmented model. Tries local path first, then HF Hub."""
num_classes = len(CLASSES)
model = models.efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)
in_features = model.classifier[1].in_features
model.classifier = nn.Sequential(
nn.Dropout(0.3),
nn.Linear(in_features, num_classes),
)
# Try local weights first (for Colab / local dev)
local_paths = [
Path('models/domain_aug_best.pth'),
Path('src/models/domain_aug_best.pth'),
]
for p in local_paths:
if p.exists():
ckpt = torch.load(p, map_location='cpu')
model.load_state_dict(ckpt['model_state_dict'])
model.eval()
return model, str(p)
return model, 'pretrained_only'
# ── Augmentation builder ──────────────────────────────────────────────────────
def build_transform(night_ir, motion_blur, low_light, noise, occlusion):
"""Build a transform from slider values (all 0-1 floats)."""
ops = [A.Resize(IMAGE_SIZE, IMAGE_SIZE)]
if night_ir > 0:
ops.append(A.ToGray(p=1.0))
ops.append(A.RandomBrightnessContrast(
brightness_limit=(-night_ir * 0.5, -night_ir * 0.1),
contrast_limit=(-0.1, 0.1), p=1.0
))
if motion_blur > 0:
blur_limit = max(3, int(motion_blur * 20))
ops.append(A.MotionBlur(blur_limit=(blur_limit, blur_limit + 4), p=1.0))
if low_light > 0:
ops.append(A.RandomBrightnessContrast(
brightness_limit=(-low_light * 0.6, -low_light * 0.2),
contrast_limit=(-low_light * 0.3, 0.0), p=1.0
))
if noise > 0:
ops.append(A.GaussNoise(std_range=(noise * 0.05, noise * 0.15), p=1.0))
if occlusion > 0:
n_holes = max(1, int(occlusion * 10))
hole_size = max(10, int(occlusion * 50))
ops.append(A.CoarseDropout(
num_holes_range=(n_holes, n_holes + 2),
hole_height_range=(hole_size, hole_size + 10),
hole_width_range=(hole_size, hole_size + 10),
p=1.0
))
ops += [A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2()]
return A.Compose(ops)
def predict(model, img_np, transform):
tensor = transform(image=img_np)['image'].unsqueeze(0)
with torch.no_grad():
logits = model(tensor)[0]
probs = torch.softmax(logits, dim=0).numpy()
top_idx = int(probs.argmax())
return CLASSES[top_idx], probs
def denormalize(tensor):
img = tensor.squeeze(0).permute(1, 2, 0).numpy()
return np.clip(img * np.array(IMAGENET_STD) + np.array(IMAGENET_MEAN), 0, 1)
# ── UI ────────────────────────────────────────────────────────────────────────
# Hero
st.markdown('🦌 WildEye
', unsafe_allow_html=True)
st.markdown(
'Robust Wildlife Classifier · Transfer Learning + Domain Augmentation
',
unsafe_allow_html=True
)
# Sidebar
with st.sidebar:
st.markdown("### 🔬 Robustness Lab")
st.markdown(
"Simulate real-world camera-trap conditions. "
"Watch how each model handles them.
",
unsafe_allow_html=True
)
night_ir = st.slider("🌙 Night / IR", 0.0, 1.0, 0.0, 0.05)
motion_blur = st.slider("💨 Motion Blur", 0.0, 1.0, 0.0, 0.05)
low_light = st.slider("🔅 Low Light", 0.0, 1.0, 0.0, 0.05)
noise = st.slider("📡 Sensor Noise", 0.0, 1.0, 0.0, 0.05)
occlusion = st.slider("🌿 Occlusion", 0.0, 1.0, 0.0, 0.05)
if st.button("⚡ Worst-case scenario", use_container_width=True):
st.session_state['preset'] = 'worst'
st.rerun()
if st.button("✨ Reset to clean", use_container_width=True):
st.session_state['preset'] = 'clean'
st.rerun()
if st.session_state.get('preset') == 'worst':
night_ir = motion_blur = low_light = noise = occlusion = 0.8
elif st.session_state.get('preset') == 'clean':
night_ir = motion_blur = low_light = noise = occlusion = 0.0
st.divider()
st.markdown("### 📊 Project Results")
st.markdown("""
+19.3%
Robustness gain
worst-case conditions
""", unsafe_allow_html=True)
st.divider()
st.markdown(
""
"Model: EfficientNet-B0
"
"Training: Animals-10 (8,695 imgs)
"
"Augmentations: 6 domain-specific
"
"Mini Hackathon #1 — WildEye
",
unsafe_allow_html=True
)
# Main content
col_upload, col_results = st.columns([1, 1], gap="large")
with col_upload:
st.markdown("#### Upload an image")
uploaded = st.file_uploader(
"Camera trap image",
type=["jpg", "jpeg", "png", "webp"],
label_visibility="collapsed"
)
if not uploaded:
st.markdown(
"📷 Drop an animal photo here
"
"Supports: jpg · png · webp
",
unsafe_allow_html=True
)
else:
img_pil = Image.open(uploaded).convert("RGB")
img_np = np.array(img_pil)
tf = build_transform(night_ir, motion_blur, low_light, noise, occlusion)
import io
from PIL import Image as PILImage
aug_tensor = tf(image=img_np)['image']
aug_display = (denormalize(aug_tensor.unsqueeze(0)) * 255).astype(np.uint8)
any_aug = any([night_ir, motion_blur, low_light, noise, occlusion])
if any_aug:
c1, c2 = st.columns(2)
c1.image(img_pil, caption="Original", use_container_width=True)
c2.image(aug_display, caption="As model sees it", use_container_width=True)
else:
st.image(img_pil, caption="Original (no perturbations)", use_container_width=True)
# Active perturbation tags
active = []
if night_ir > 0: active.append(f"🌙 Night IR ({night_ir:.0%})")
if motion_blur > 0: active.append(f"💨 Blur ({motion_blur:.0%})")
if low_light > 0: active.append(f"🔅 Low Light ({low_light:.0%})")
if noise > 0: active.append(f"📡 Noise ({noise:.0%})")
if occlusion > 0: active.append(f"🌿 Occlusion ({occlusion:.0%})")
if active:
tags_html = " ".join(f"{t}" for t in active)
st.markdown(tags_html, unsafe_allow_html=True)
with col_results:
if uploaded:
st.markdown("#### Prediction")
model, source = load_model()
pred_class, probs = predict(model, img_np, tf)
conf = probs.max()
emoji = CLASS_EMOJI.get(pred_class, "🐾")
# Main prediction card
color = "#3fb950" if conf > 0.7 else "#d29922" if conf > 0.4 else "#f85149"
st.markdown(f"""
{emoji}
{pred_class.upper()}
Confidence: {conf:.1%}
""", unsafe_allow_html=True)
# Top-5 bar chart
st.markdown("
**Top predictions**", unsafe_allow_html=True)
top5_idx = probs.argsort()[::-1][:5]
top5_cls = [CLASSES[i] for i in top5_idx]
top5_prob = [probs[i] for i in top5_idx]
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(5, 2.8))
fig.patch.set_facecolor('#161b22')
ax.set_facecolor('#161b22')
colors_bar = ['#3fb950' if c == pred_class else '#30363d' for c in top5_cls]
bars = ax.barh(
[f"{CLASS_EMOJI.get(c,'🐾')} {c}" for c in top5_cls[::-1]],
top5_prob[::-1],
color=colors_bar[::-1], edgecolor='none', height=0.6
)
for bar, val in zip(bars, top5_prob[::-1]):
ax.text(min(val + 0.01, 0.95), bar.get_y() + bar.get_height()/2,
f'{val:.1%}', va='center', color='#e6edf3',
fontsize=8, fontfamily='monospace')
ax.set_xlim(0, 1.05)
ax.tick_params(colors='#8b949e', labelsize=8)
ax.spines[['top','right','bottom']].set_visible(False)
ax.spines['left'].set_color('#30363d')
ax.xaxis.set_visible(False)
plt.tight_layout(pad=0.5)
st.pyplot(fig, use_container_width=True)
plt.close()
# Robustness warning
if any_aug:
severity = np.mean([night_ir, motion_blur, low_light, noise, occlusion])
if conf < 0.5:
st.markdown(
"⚠️ Low confidence under these conditions — "
"this is exactly what domain augmentation trains against.
",
unsafe_allow_html=True
)
elif conf > 0.8:
st.markdown(
"✓ High confidence maintained despite perturbations — "
"domain augmentation working as intended.
",
unsafe_allow_html=True
)
else:
st.markdown(
""
"Upload an image to see predictions.
"
"Use the sidebar sliders to simulate
real camera-trap conditions.
",
unsafe_allow_html=True
)
# Footer
st.divider()
st.markdown(
"WildEye · Mini Hackathon #1 · "
"EfficientNet-B0 + Domain Augmentation · "
"Robustness gain: +19.3% under worst-case conditions
",
unsafe_allow_html=True
)