""" shared.py — All shared logic for TrafficSense AI Models, preprocessing, prediction, helpers, CSS, constants """ import os, io, json, time, csv, datetime import numpy as np from PIL import Image, ImageFilter import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import torch import torch.nn as nn import torchvision.transforms as transforms from torchvision.models import efficientnet_b3 import streamlit as st try: import cv2 HAS_CV2 = True except ImportError: HAS_CV2 = False # ── Constants ────────────────────────────────────────────────────────────────── IMG_SIZE = 224 DEVICE = torch.device("cpu") MEAN = [0.485, 0.456, 0.406] STD = [0.229, 0.224, 0.225] NO_SIGN_THRESHOLD = 0.25 LOW_CONF = 0.60 HIGH_CONF = 0.95 MAX_HISTORY = 200 DANGER_ICONS = {"MANDATORY": "🔴", "WARNING": "🟡", "INFO": "🟢"} DANGER_COLORS = {"MANDATORY": "#e24b4a", "WARNING": "#f59e0b", "INFO": "#22c55e"} # ── Model ────────────────────────────────────────────────────────────────────── class TrafficSignNet(nn.Module): def __init__(self, num_classes): super().__init__() self.base = efficientnet_b3(weights=None) inf = self.base.classifier[1].in_features self.base.classifier = nn.Sequential( nn.Dropout(0.4), nn.Linear(inf, 512), nn.SiLU(), nn.BatchNorm1d(512), nn.Dropout(0.3), nn.Linear(512, 256), nn.SiLU(), nn.BatchNorm1d(256), nn.Dropout(0.2), nn.Linear(256, num_classes) ) def forward(self, x): return self.base(x) @st.cache_resource def load_models(): out = {} for key, fname, n in [("indian","indian_traffic_sign_model.pth",59),("gtsrb","gtsrb_traffic_sign_model.pth",43)]: m = TrafficSignNet(n); loaded = False for p in [fname, f"/app/{fname}", os.path.join(os.path.dirname(__file__), fname)]: if not os.path.exists(p): continue try: obj = torch.load(p, map_location=DEVICE, weights_only=False) state = obj if isinstance(obj, dict) else obj.state_dict() miss, _ = m.load_state_dict(state, strict=False) if 1 - len(miss) / len(m.state_dict()) >= 0.85: m.eval(); loaded = True; break except: pass out[key] = (m, loaded) return out @st.cache_data def load_all_labels(): out = {} for key, fname, n in [("indian","indian_labels.json",59),("gtsrb","gtsrb_labels.json",43)]: for p in [fname, f"/app/{fname}", os.path.join(os.path.dirname(__file__), fname)]: if os.path.exists(p): with open(p) as f: out[key] = json.load(f); break if key not in out: out[key] = {str(i): f"Class {i}" for i in range(n)} return out @st.cache_data def load_summary(): for f in ["training_summary.json", "/app/training_summary.json"]: if os.path.exists(f): with open(f) as fh: return json.load(fh) return {} # ── Sign metadata ────────────────────────────────────────────────────────────── GTSRB_INFO = { "0": {"danger":"MANDATORY","meaning":"Speed limit 20 km/h.","action":"Reduce to 20 km/h immediately."}, "1": {"danger":"MANDATORY","meaning":"Speed limit 30 km/h — residential zone.","action":"Reduce to 30 km/h. Watch for pedestrians."}, "2": {"danger":"MANDATORY","meaning":"Speed limit 50 km/h — standard urban limit.","action":"Maintain at or below 50 km/h."}, "3": {"danger":"MANDATORY","meaning":"Speed limit 60 km/h.","action":"Reduce to 60 km/h."}, "4": {"danger":"MANDATORY","meaning":"Speed limit 70 km/h.","action":"Reduce to 70 km/h."}, "5": {"danger":"MANDATORY","meaning":"Speed limit 80 km/h.","action":"Reduce to 80 km/h."}, "6": {"danger":"INFO","meaning":"End of 80 km/h restriction.","action":"Drive to conditions."}, "7": {"danger":"MANDATORY","meaning":"Speed limit 100 km/h.","action":"Do not exceed 100 km/h."}, "8": {"danger":"MANDATORY","meaning":"Speed limit 120 km/h.","action":"Do not exceed 120 km/h."}, "9": {"danger":"MANDATORY","meaning":"No overtaking for all vehicles.","action":"Stay in lane. Do not overtake."}, "10": {"danger":"MANDATORY","meaning":"No overtaking for vehicles over 3.5t.","action":"Heavy vehicles must not overtake."}, "11": {"danger":"INFO","meaning":"Right of way at next intersection.","action":"You have priority. Maintain speed."}, "12": {"danger":"INFO","meaning":"Priority road — continuous right of way.","action":"Maintain speed. You have priority."}, "13": {"danger":"WARNING","meaning":"Yield — give way to main road traffic.","action":"Slow down and yield to all traffic."}, "14": {"danger":"MANDATORY","meaning":"Stop sign — complete stop required.","action":"Stop completely before the line."}, "15": {"danger":"MANDATORY","meaning":"No vehicles permitted.","action":"Do not enter. Find alternative route."}, "16": {"danger":"MANDATORY","meaning":"Vehicles over 3.5t prohibited.","action":"Heavy vehicles must use alternative route."}, "17": {"danger":"MANDATORY","meaning":"No entry — one-way road.","action":"Do not enter. Turn around."}, "18": {"danger":"WARNING","meaning":"General hazard warning ahead.","action":"Reduce speed. Be alert."}, "19": {"danger":"WARNING","meaning":"Dangerous left curve ahead.","action":"Reduce speed before the curve."}, "20": {"danger":"WARNING","meaning":"Dangerous right curve ahead.","action":"Reduce speed before the curve."}, "21": {"danger":"WARNING","meaning":"Double curve ahead.","action":"Reduce speed. Keep to lane."}, "22": {"danger":"WARNING","meaning":"Bumpy road ahead.","action":"Reduce speed to protect suspension."}, "23": {"danger":"WARNING","meaning":"Slippery road ahead.","action":"Reduce speed. Avoid sudden inputs."}, "24": {"danger":"WARNING","meaning":"Road narrows on the right.","action":"Move centre-left. Watch oncoming traffic."}, "25": {"danger":"WARNING","meaning":"Roadworks ahead.","action":"Reduce speed. Watch for workers."}, "26": {"danger":"WARNING","meaning":"Traffic signals ahead.","action":"Prepare to stop at red light."}, "27": {"danger":"WARNING","meaning":"Pedestrian crossing ahead.","action":"Slow down. Yield to pedestrians."}, "28": {"danger":"WARNING","meaning":"Children crossing — school zone.","action":"Reduce speed sharply. Watch for children."}, "29": {"danger":"WARNING","meaning":"Bicycle crossing ahead.","action":"Check mirrors. Give cyclists space."}, "30": {"danger":"WARNING","meaning":"Ice or snow likely.","action":"Reduce speed. Increase following distance."}, "31": {"danger":"WARNING","meaning":"Wild animals crossing ahead.","action":"Reduce speed. Scan road edges."}, "32": {"danger":"INFO","meaning":"End of all restrictions.","action":"Normal rules apply. Drive to conditions."}, "33": {"danger":"MANDATORY","meaning":"Turn right — mandatory.","action":"Signal right. Prepare to turn right."}, "34": {"danger":"MANDATORY","meaning":"Turn left — mandatory.","action":"Signal left. Prepare to turn left."}, "35": {"danger":"MANDATORY","meaning":"Proceed straight only.","action":"Do not turn. Go straight only."}, "36": {"danger":"MANDATORY","meaning":"Go straight or turn right only.","action":"Left turns not permitted."}, "37": {"danger":"MANDATORY","meaning":"Go straight or turn left only.","action":"Right turns not permitted."}, "38": {"danger":"MANDATORY","meaning":"Keep right.","action":"Pass obstacle on the right side."}, "39": {"danger":"MANDATORY","meaning":"Keep left.","action":"Pass obstacle on the left side."}, "40": {"danger":"MANDATORY","meaning":"Roundabout — yield to circulating traffic.","action":"Yield before entering roundabout."}, "41": {"danger":"INFO","meaning":"End of no-overtaking zone.","action":"Overtaking now permitted where safe."}, "42": {"danger":"INFO","meaning":"End of no-overtaking zone for heavy vehicles.","action":"Heavy vehicles may overtake where safe."}, } INDIAN_DEFAULT = {"danger":"INFO","meaning":"Indian road sign detected.","action":"Follow standard road rules and local traffic laws."} def get_info(mk, cid): if mk == "gtsrb": return GTSRB_INFO.get(str(cid), {"danger":"INFO","meaning":"No description.","action":"Follow standard rules."}) return INDIAN_DEFAULT # ── Preprocessing ────────────────────────────────────────────────────────────── def preprocess_image(img_pil): warns = []; changed = False try: from PIL import ImageOps img_pil = ImageOps.exif_transpose(img_pil) except: pass if img_pil.mode == "P" and "transparency" in img_pil.info: img_pil = img_pil.convert("RGBA") if img_pil.mode in ("RGBA", "LA"): bg = Image.new("RGB", img_pil.size, (255,255,255)) mask = img_pil.split()[-1] if "A" in img_pil.mode else None bg.paste(img_pil.convert("RGB"), mask=mask); img_pil = bg elif img_pil.mode != "RGB": img_pil = img_pil.convert("RGB") arr = np.array(img_pil, dtype=np.float32) if arr.shape[2] == 3: ch_std = np.std([arr[:,:,0].mean(), arr[:,:,1].mean(), arr[:,:,2].mean()]) if ch_std < 5: warns.append(("⚠️","Grayscale image — colour images yield better accuracy.")) w, h = img_pil.size if w < 64 or h < 64: warns.append(("⚠️", f"Very small image ({w}×{h}px) — upscaling with Lanczos.")) scale = max(128/w, 128/h) img_pil = img_pil.resize((int(w*scale), int(h*scale)), Image.LANCZOS) arr = np.array(img_pil, dtype=np.float32); changed = True gray = np.mean(arr, axis=2).astype(np.float32) gy, gx = np.gradient(gray); lap_var = float(np.var(np.sqrt(gx**2 + gy**2))) if lap_var < 80: warns.append(("⚠️", f"Blurry image (sharpness:{lap_var:.0f}) — results may be less accurate.")) brightness = arr.mean() if brightness < 50: warns.append(("🌙", "Dark image — applying CLAHE contrast enhancement.")) if HAS_CV2: lab = cv2.cvtColor(arr.astype(np.uint8), cv2.COLOR_RGB2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8)) l = clahe.apply(l) arr = cv2.cvtColor(cv2.merge([l,a,b]), cv2.COLOR_LAB2RGB).astype(np.float32) else: arr = np.clip(arr * (120 / max(brightness, 1)), 0, 255) changed = True elif brightness > 210: warns.append(("☀️", "Overexposed image — applying gamma correction.")) arr = np.clip(np.power(arr/255.0, 1.4)*255.0, 0, 255); changed = True if HAS_CV2: arr_u8 = arr.astype(np.uint8) arr_u8 = cv2.bilateralFilter(arr_u8, d=9, sigmaColor=75, sigmaSpace=75) noise_est = float(np.std(arr - arr_u8.astype(np.float32))) if noise_est > 12: arr_u8 = cv2.fastNlMeansDenoisingColored(arr_u8, None, 10,10,7,21) arr = arr_u8.astype(np.float32); changed = True lab = cv2.cvtColor(arr.astype(np.uint8), cv2.COLOR_RGB2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8)) l = clahe.apply(l) arr = cv2.cvtColor(cv2.merge([l,a,b]), cv2.COLOR_LAB2RGB).astype(np.float32) else: img_tmp = Image.fromarray(arr.astype(np.uint8)).filter(ImageFilter.GaussianBlur(0.5)) arr = np.array(img_tmp, dtype=np.float32) img_s = Image.fromarray(arr.astype(np.uint8)) arr_s = np.array(img_s, dtype=np.float32) arr_b = np.array(img_s.filter(ImageFilter.GaussianBlur(1.0)), dtype=np.float32) arr = np.clip(arr_s + 0.4*(arr_s - arr_b), 0, 255) img_final = Image.fromarray(arr.astype(np.uint8)) W, H = img_final.size; side = min(W, H) img_final = img_final.crop(((W-side)//2, (H-side)//2, (W-side)//2+side, (H-side)//2+side)) img_final = img_final.resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS) # Always ensure solid RGB output (no transparency channel that causes checkerboard display) if img_final.mode != "RGB": bg = Image.new("RGB", img_final.size, (255, 255, 255)) if img_final.mode in ("RGBA", "LA"): bg.paste(img_final.convert("RGB"), mask=img_final.split()[-1]) else: bg.paste(img_final.convert("RGB")) img_final = bg return img_final, warns, changed val_transform = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize(MEAN, STD) ]) def predict(model, img_pil): t0 = time.time() tensor = val_transform(img_pil).unsqueeze(0).to(DEVICE) with torch.no_grad(): probs = torch.softmax(model(tensor), dim=1)[0] ms = (time.time() - t0) * 1000 topk_probs, topk_ids = probs.topk(5) return topk_ids.cpu().numpy(), topk_probs.cpu().numpy(), ms def make_csv(rows, cols): buf = io.StringIO(); w = csv.writer(buf); w.writerow(cols) for r in rows: w.writerow(r) return buf.getvalue().encode() # ── Render helpers ───────────────────────────────────────────────────────────── def badge(danger): icon = DANGER_ICONS.get(danger, "🟢") cls = {"MANDATORY":"bM","WARNING":"bW","INFO":"bI"}.get(danger, "bI") return f'{icon} {danger}' def conf_ind(c): if c >= HIGH_CONF: return '
✅ High Confidence — reliable prediction
' if NO_SIGN_THRESHOLD <= c < LOW_CONF: return '
⚠️ Low confidence — try better lighting or crop tighter
' return "" def render_warns(ws): for icon, msg in ws: st.markdown(f'
{icon}{msg}
', unsafe_allow_html=True) def render_pred(cid, conf, mk, mk_label, labels): name = labels.get(str(cid), f"Class {cid}"); info = get_info(mk, cid) color = DANGER_COLORS.get(info["danger"], "#22c55e"); pct = conf * 100 st.markdown(f"""
{name}
{mk_label} · Class #{cid}
{badge(info['danger'])}
{conf_ind(conf)}
Confidence {pct:.1f}%
{pct:.1f}%
What this sign means
{info['meaning']}
ADAS Action
{info['action']}
""", unsafe_allow_html=True) def render_no_sign(conf): st.markdown(f"""
🔍
No Traffic Sign Detected
Best confidence: {conf*100:.1f}% — below {NO_SIGN_THRESHOLD*100:.0f}% threshold
💡 TIPS TO IMPROVE DETECTION
📐 Crop tighter — sign should fill at least 60% of frame
💡 Ensure good lighting — avoid shadows or harsh backlight
📷 Hold steady — motion blur significantly reduces accuracy
🔄 Try the other model — Indian vs GTSRB sign sets differ
🖼 Use a clear front-facing angle — avoid extreme perspective
""", unsafe_allow_html=True) def render_top3(ids, probs, labels, mk): above20 = sum(1 for p in probs[:3] if p > 0.20) dangers = [get_info(mk, int(ids[i]))["danger"] for i in range(3) if probs[i] > 0.20] multi = above20 >= 3 and len(set(dangers)) >= 2 colors = ["#e24b4a", "#3b82f6", "#7c3aed"] st.markdown('
Top 3 Predictions
', unsafe_allow_html=True) for idx, prob, color in zip(ids[:3], probs[:3], colors): name = labels.get(str(idx), f"Class {idx}"); pct = prob * 100 icon = DANGER_ICONS.get(get_info(mk, int(idx))["danger"], "🟢") st.markdown(f'
{icon} {name}
{pct:.1f}%
', unsafe_allow_html=True) if multi: st.markdown('
⚠️ Multiple sign types detected — more than one sign may be in frame
', unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) def draw_sparkline(vals): if len(vals) < 2: return None fig, ax = plt.subplots(figsize=(4, 0.9)) fig.patch.set_facecolor("#080a10"); ax.set_facecolor("#080a10") x = list(range(len(vals))) ax.plot(x, vals, color="#e24b4a", lw=1.5, marker="o", markersize=3) ax.fill_between(x, vals, alpha=0.1, color="#e24b4a") ax.set_xlim(0, max(len(vals)-1, 1)); ax.set_ylim(0, 1); ax.axis("off") plt.tight_layout(pad=0) buf = io.BytesIO(); plt.savefig(buf, format="png", dpi=100, bbox_inches="tight", facecolor="#080a10") buf.seek(0); plt.close(); return buf def draw_gauge(pct, color, label): fig, ax = plt.subplots(figsize=(2.6, 1.5), subplot_kw=dict(projection="polar")) fig.patch.set_facecolor("#0d0f18") theta_val = np.pi - (np.pi * pct / 100) thetas = np.linspace(np.pi, 0, 100) ax.plot(thetas, [1]*100, color="#1a1d2e", lw=12, solid_capstyle="round") thetas2 = np.linspace(np.pi, theta_val, 100) ax.plot(thetas2, [1]*100, color=color, lw=12, solid_capstyle="round") ax.text(0, 0, f"{pct:.1f}%", ha="center", va="center", fontsize=13, fontweight="700", color=color, fontfamily="monospace") ax.text(0, -0.55, label, ha="center", va="center", fontsize=7, color="#4b5563") ax.set_ylim(0, 1.3); ax.axis("off") plt.tight_layout(pad=0) buf = io.BytesIO(); plt.savefig(buf, format="png", dpi=120, bbox_inches="tight", facecolor="#0d0f18") buf.seek(0); plt.close(); return buf # ── Shared CSS (base styles used by both desktop and mobile) ─────────────────── BASE_CSS = """ """