Upload 3 files

app.py

@@ -0,0 +1,1283 @@
+"""
+Screenshot Anonymizer
+=====================
+Backend : gr.Server (Gradio + FastAPI)
+Frontend: Canvas-based image editor
+Model   : charles-first-org/second-model  (OpenAI Privacy Filter)
+OCR     : pytesseract (Tesseract 5)
+
+Drag in a screenshot of a chat / email / document. OCR extracts words with
+pixel positions, the privacy filter finds PII character spans, we map the
+spans back to pixel rectangles and render black bars on top. The canvas
+editor lets the user toggle, move, add, or delete bars, then export the
+redacted screenshot as a PNG.
+"""
+
+# ── stdlib ───────────────────────────────────────────────────────
+import base64
+import dataclasses
+import functools
+import io
+import json
+import math
+import os
+from bisect import bisect_left, bisect_right
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+# ── third-party ──────────────────────────────────────────────────
+import gradio as gr
+import pytesseract
+import spaces
+import tiktoken
+import torch
+import torch.nn.functional as F
+from fastapi import File, UploadFile
+from fastapi.responses import HTMLResponse, JSONResponse
+from huggingface_hub import snapshot_download
+from PIL import Image
+from safetensors import safe_open
+
+# ── configuration ────────────────────────────────────────────────
+MODEL_REPO = os.getenv("MODEL_ID", "charles-first-org/second-model")
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+MODEL_DIR = Path(snapshot_download(MODEL_REPO, token=HF_TOKEN))
+
+CATEGORIES_META = {
+    "private_person":  {"color": "#ef4444", "label": "Person"},
+    "private_address": {"color": "#06b6d4", "label": "Address"},
+    "private_email":   {"color": "#3b82f6", "label": "Email"},
+    "private_phone":   {"color": "#22c55e", "label": "Phone"},
+    "private_url":     {"color": "#eab308", "label": "URL"},
+    "private_date":    {"color": "#a855f7", "label": "Date"},
+    "account_number":  {"color": "#f97316", "label": "Account"},
+    "secret":          {"color": "#dc2626", "label": "Secret"},
+}
+
+# =====================================================================
+# MODEL  ARCHITECTURE  +  INFERENCE
+# =====================================================================
+
+PRIVACY_FILTER_MODEL_TYPE: Final[str] = "privacy_filter"
+REQUIRED_MODEL_CONFIG_KEYS: Final[tuple[str, ...]] = (
+    "model_type", "encoding", "num_hidden_layers", "num_experts",
+    "experts_per_token", "vocab_size", "num_labels", "hidden_size",
+    "intermediate_size", "head_dim", "num_attention_heads",
+    "num_key_value_heads", "sliding_window", "bidirectional_context",
+    "bidirectional_left_context", "bidirectional_right_context",
+    "default_n_ctx", "initial_context_length", "rope_theta",
+    "rope_scaling_factor", "rope_ntk_alpha", "rope_ntk_beta", "param_dtype",
+)
+BACKGROUND_CLASS_LABEL: Final[str] = "O"
+BOUNDARY_PREFIXES: Final[tuple[str, ...]] = ("B", "I", "E", "S")
+SPAN_CLASS_NAMES: Final[tuple[str, ...]] = (
+    BACKGROUND_CLASS_LABEL,
+    "account_number", "private_address", "private_date", "private_email",
+    "private_person", "private_phone", "private_url", "secret",
+)
+NER_CLASS_NAMES: Final[tuple[str, ...]] = (BACKGROUND_CLASS_LABEL,) + tuple(
+    f"{prefix}-{base}"
+    for base in SPAN_CLASS_NAMES if base != BACKGROUND_CLASS_LABEL
+    for prefix in BOUNDARY_PREFIXES
+)
+VITERBI_TRANSITION_BIAS_KEYS: Final[tuple[str, ...]] = (
+    "transition_bias_background_stay", "transition_bias_background_to_start",
+    "transition_bias_inside_to_continue", "transition_bias_inside_to_end",
+    "transition_bias_end_to_background", "transition_bias_end_to_start",
+)
+DEFAULT_VITERBI_CALIBRATION_PRESET: Final[str] = "default"
+
+
+def validate_model_config_contract(cfg: dict, *, context: str) -> None:
+    missing = [k for k in REQUIRED_MODEL_CONFIG_KEYS if k not in cfg]
+    if missing:
+        raise ValueError(f"{context} missing keys: {', '.join(missing)}")
+    if cfg.get("model_type") != PRIVACY_FILTER_MODEL_TYPE:
+        raise ValueError(f"{context} model_type must be {PRIVACY_FILTER_MODEL_TYPE!r}")
+    if cfg.get("bidirectional_context") is not True:
+        raise ValueError(f"{context} must use bidirectional_context=true")
+    lc, rc = cfg.get("bidirectional_left_context"), cfg.get("bidirectional_right_context")
+    if not isinstance(lc, int) or not isinstance(rc, int) or lc != rc or lc < 0:
+        raise ValueError(f"{context} bidirectional context must be equal non-negative ints")
+    sw = cfg.get("sliding_window")
+    if sw != 2 * lc + 1:
+        raise ValueError(f"{context} sliding_window must equal 2*context+1")
+    if cfg["num_labels"] != 33:
+        raise ValueError(f"{context} num_labels must be 33")
+    if cfg["param_dtype"] != "bfloat16":
+        raise ValueError(f"{context} param_dtype must be bfloat16")
+
+
+def expert_linear(x, weight, bias):
+    n, e, k = x.shape
+    _, _, _, o = weight.shape
+    out = torch.bmm(x.reshape(n * e, 1, k), weight.reshape(n * e, k, o)).reshape(n, e, o)
+    return out + bias if bias is not None else out
+
+
+@dataclass
+class ModelConfig:
+    num_hidden_layers: int; num_experts: int; experts_per_token: int
+    vocab_size: int; num_labels: int; hidden_size: int; intermediate_size: int
+    head_dim: int; num_attention_heads: int; num_key_value_heads: int
+    bidirectional_context_size: int; initial_context_length: int
+    rope_theta: float; rope_scaling_factor: float; rope_ntk_alpha: float; rope_ntk_beta: float
+
+    @classmethod
+    def from_checkpoint_config(cls, cfg, *, context):
+        cfg = dict(cfg)
+        cfg["bidirectional_context_size"] = cfg["bidirectional_left_context"]
+        fields = {f.name for f in dataclasses.fields(cls)}
+        return cls(**{k: v for k, v in cfg.items() if k in fields})
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, n, eps=1e-5, device=None):
+        super().__init__()
+        self.eps = eps
+        self.scale = torch.nn.Parameter(torch.ones(n, device=device, dtype=torch.float32))
+
+    def forward(self, x):
+        t = x.float()
+        return (t * torch.rsqrt(t.pow(2).mean(-1, keepdim=True) + self.eps) * self.scale).to(x.dtype)
+
+
+def apply_rope(x, cos, sin):
+    cos = cos.unsqueeze(-2).to(x.dtype); sin = sin.unsqueeze(-2).to(x.dtype)
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    return torch.stack((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1).reshape(x.shape)
+
+
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self, head_dim, base, dtype, *, initial_context_length=4096,
+                 scaling_factor=1.0, ntk_alpha=1.0, ntk_beta=32.0, device=None):
+        super().__init__()
+        self.head_dim, self.base, self.dtype = head_dim, base, dtype
+        self.initial_context_length = initial_context_length
+        self.scaling_factor, self.ntk_alpha, self.ntk_beta = scaling_factor, ntk_alpha, ntk_beta
+        self.device = device
+        mp = max(int(initial_context_length * scaling_factor), initial_context_length)
+        self.max_position_embeddings = mp
+        cos, sin = self._compute(mp, device=torch.device("cpu"))
+        target = device or torch.device("cpu")
+        self.register_buffer("cos_cache", cos.to(target), persistent=False)
+        self.register_buffer("sin_cache", sin.to(target), persistent=False)
+
+    def _inv_freq(self, device=None):
+        device = device or self.device
+        freq = self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float, device=device) / self.head_dim)
+        if self.scaling_factor > 1.0:
+            d_half = self.head_dim / 2
+            low = d_half * math.log(self.initial_context_length / (self.ntk_beta * 2 * math.pi)) / math.log(self.base)
+            high = d_half * math.log(self.initial_context_length / (self.ntk_alpha * 2 * math.pi)) / math.log(self.base)
+            interp = 1.0 / (self.scaling_factor * freq)
+            extrap = 1.0 / freq
+            ramp = (torch.arange(d_half, dtype=torch.float32, device=device) - low) / (high - low)
+            mask = 1 - ramp.clamp(0, 1)
+            return interp * (1 - mask) + extrap * mask
+        return 1.0 / freq
+
+    def _compute(self, n, device=None):
+        inv_freq = self._inv_freq(device)
+        t = torch.arange(n, dtype=torch.float32, device=device or self.device)
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        c = 0.1 * math.log(self.scaling_factor) + 1.0 if self.scaling_factor > 1.0 else 1.0
+        return (freqs.cos() * c).to(self.dtype), (freqs.sin() * c).to(self.dtype)
+
+    def forward(self, q, k):
+        n = q.shape[0]
+        if n > self.cos_cache.shape[0]:
+            cos, sin = self._compute(n, torch.device("cpu"))
+            self.cos_cache, self.sin_cache = cos.to(q.device), sin.to(q.device)
+        cc = self.cos_cache.to(q.device) if self.cos_cache.device != q.device else self.cos_cache
+        sc = self.sin_cache.to(q.device) if self.sin_cache.device != q.device else self.sin_cache
+        cos, sin = cc[:n], sc[:n]
+        q = apply_rope(q.view(n, -1, self.head_dim), cos, sin).reshape(q.shape)
+        k = apply_rope(k.view(n, -1, self.head_dim), cos, sin).reshape(k.shape)
+        return q, k
+
+
+def sdpa(Q, K, V, S, sm_scale, ctx):
+    n, nh, qm, hd = Q.shape
+    w = 2 * ctx + 1
+    Kp = F.pad(K, (0, 0, 0, 0, ctx, ctx)); Vp = F.pad(V, (0, 0, 0, 0, ctx, ctx))
+    Kw = Kp.unfold(0, w, 1).permute(0, 3, 1, 2); Vw = Vp.unfold(0, w, 1).permute(0, 3, 1, 2)
+    idx = torch.arange(w, device=Q.device) - ctx
+    pos = torch.arange(n, device=Q.device)[:, None] + idx[None, :]
+    valid = (pos >= 0) & (pos < n)
+    scores = torch.einsum("nhqd,nwhd->nhqw", Q, Kw).float() * sm_scale
+    scores = scores.masked_fill(~valid[:, None, None, :], -float("inf"))
+    sink = (S * math.log(2.0)).reshape(nh, qm)[None, :, :, None].expand(n, -1, -1, 1)
+    scores = torch.cat([scores, sink], dim=-1)
+    wt = torch.softmax(scores, dim=-1)[..., :-1].to(V.dtype)
+    return torch.einsum("nhqw,nwhd->nhqd", wt, Vw).reshape(n, -1)
+
+
+class AttentionBlock(torch.nn.Module):
+    def __init__(self, cfg, device=None):
+        super().__init__()
+        dt = torch.bfloat16
+        self.head_dim, self.nah, self.nkv = cfg.head_dim, cfg.num_attention_heads, cfg.num_key_value_heads
+        self.ctx = int(cfg.bidirectional_context_size)
+        self.sinks = torch.nn.Parameter(torch.empty(cfg.num_attention_heads, device=device, dtype=torch.float32))
+        self.norm = RMSNorm(cfg.hidden_size, device=device)
+        qkv_d = cfg.head_dim * (cfg.num_attention_heads + 2 * cfg.num_key_value_heads)
+        self.qkv = torch.nn.Linear(cfg.hidden_size, qkv_d, device=device, dtype=dt)
+        self.out = torch.nn.Linear(cfg.head_dim * cfg.num_attention_heads, cfg.hidden_size, device=device, dtype=dt)
+        self.qk_scale = 1 / math.sqrt(math.sqrt(cfg.head_dim))
+        self.rope = RotaryEmbedding(cfg.head_dim, int(cfg.rope_theta), torch.float32,
+                                     initial_context_length=cfg.initial_context_length,
+                                     scaling_factor=cfg.rope_scaling_factor,
+                                     ntk_alpha=cfg.rope_ntk_alpha, ntk_beta=cfg.rope_ntk_beta, device=device)
+
+    def forward(self, x):
+        t = self.norm(x).to(self.qkv.weight.dtype)
+        qkv = F.linear(t, self.qkv.weight, self.qkv.bias)
+        hd, nah, nkv = self.head_dim, self.nah, self.nkv
+        q = qkv[:, :nah * hd].contiguous()
+        k = qkv[:, nah * hd:(nah + nkv) * hd].contiguous()
+        v = qkv[:, (nah + nkv) * hd:(nah + 2 * nkv) * hd].contiguous()
+        q, k = self.rope(q, k)
+        q, k = q * self.qk_scale, k * self.qk_scale
+        n = q.shape[0]
+        q = q.view(n, nkv, nah // nkv, hd); k = k.view(n, nkv, hd); v = v.view(n, nkv, hd)
+        ao = sdpa(q, k, v, self.sinks, 1.0, self.ctx).to(self.out.weight.dtype)
+        return x + F.linear(ao, self.out.weight, self.out.bias).to(x.dtype)
+
+
+def swiglu(x, alpha=1.702, limit=7.0):
+    g, l = x.chunk(2, dim=-1)
+    g, l = g.clamp(max=limit), l.clamp(-limit, limit)
+    return g * torch.sigmoid(alpha * g) * (l + 1)
+
+
+class MLPBlock(torch.nn.Module):
+    def __init__(self, cfg, device=None):
+        super().__init__()
+        dt = torch.bfloat16
+        self.ne, self.ept = cfg.num_experts, cfg.experts_per_token
+        self.norm = RMSNorm(cfg.hidden_size, device=device)
+        self.gate = torch.nn.Linear(cfg.hidden_size, cfg.num_experts, device=device, dtype=dt)
+        self.mlp1_weight = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.hidden_size, cfg.intermediate_size * 2, device=device, dtype=dt))
+        self.mlp1_bias = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.intermediate_size * 2, device=device, dtype=dt))
+        self.mlp2_weight = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.intermediate_size, cfg.hidden_size, device=device, dtype=dt))
+        self.mlp2_bias = torch.nn.Parameter(torch.empty(cfg.num_experts, cfg.hidden_size, device=device, dtype=dt))
+
+    def forward(self, x):
+        t = self.norm(x)
+        gs = F.linear(t.float(), self.gate.weight.float(), self.gate.bias.float())
+        top = torch.topk(gs, k=self.ept, dim=-1, sorted=True)
+        ew = torch.softmax(top.values, dim=-1) / self.ept
+        ei = top.indices
+        ept = self.ept
+
+        def _chunk(tc, eic, ewc):
+            o = expert_linear(tc.float().unsqueeze(1).expand(-1, eic.shape[1], -1),
+                              self.mlp1_weight[eic].float(), self.mlp1_bias[eic].float())
+            o = swiglu(o)
+            o = expert_linear(o.float(), self.mlp2_weight[eic].float(), self.mlp2_bias[eic].float())
+            return (torch.einsum("bec,be->bc", o.to(ewc.dtype), ewc) * ept).to(x.dtype)
+
+        cs = 32
+        if t.shape[0] > cs:
+            parts = [_chunk(t[s:s+cs], ei[s:s+cs], ew[s:s+cs]) for s in range(0, t.shape[0], cs)]
+            return x + torch.cat(parts, 0)
+        return x + _chunk(t, ei, ew)
+
+
+class TransformerBlock(torch.nn.Module):
+    def __init__(self, cfg, device=None):
+        super().__init__()
+        self.attn = AttentionBlock(cfg, device=device)
+        self.mlp = MLPBlock(cfg, device=device)
+    def forward(self, x):
+        return self.mlp(self.attn(x))
+
+
+class Checkpoint:
+    @staticmethod
+    def build_param_name_map(n):
+        return ({f"block.{i}.mlp.mlp1_bias": f"block.{i}.mlp.swiglu.bias" for i in range(n)}
+              | {f"block.{i}.mlp.mlp1_weight": f"block.{i}.mlp.swiglu.weight" for i in range(n)}
+              | {f"block.{i}.mlp.mlp2_bias": f"block.{i}.mlp.out.bias" for i in range(n)}
+              | {f"block.{i}.mlp.mlp2_weight": f"block.{i}.mlp.out.weight" for i in range(n)})
+
+    def __init__(self, path, device, num_hidden_layers):
+        self.pnm = self.build_param_name_map(num_hidden_layers)
+        self.ds = device.type if device.index is None else f"{device.type}:{device.index}"
+        files = [os.path.join(path, f) for f in os.listdir(path) if f.endswith(".safetensors")]
+        self.map = {}
+        for sf in files:
+            with safe_open(sf, framework="pt", device=self.ds) as h:
+                for k in h.keys():
+                    self.map[k] = sf
+
+    def get(self, name):
+        mapped = self.pnm.get(name, name)
+        with safe_open(self.map[mapped], framework="pt", device=self.ds) as h:
+            return h.get_tensor(mapped)
+
+
+class Transformer(torch.nn.Module):
+    def __init__(self, cfg, device):
+        super().__init__()
+        dt = torch.bfloat16
+        self.embedding = torch.nn.Embedding(cfg.vocab_size, cfg.hidden_size, device=device, dtype=dt)
+        self.block = torch.nn.ModuleList([TransformerBlock(cfg, device=device) for _ in range(cfg.num_hidden_layers)])
+        self.norm = RMSNorm(cfg.hidden_size, device=device)
+        self.unembedding = torch.nn.Linear(cfg.hidden_size, cfg.num_labels, bias=False, device=device, dtype=dt)
+
+    def forward(self, token_ids):
+        x = self.embedding(token_ids)
+        for blk in self.block:
+            x = blk(x)
+        return F.linear(self.norm(x), self.unembedding.weight, None)
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint_dir, *, device):
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        torch.set_float32_matmul_precision("highest")
+        cp = json.loads((Path(checkpoint_dir) / "config.json").read_text())
+        validate_model_config_contract(cp, context=str(checkpoint_dir))
+        cfg = ModelConfig.from_checkpoint_config(cp, context=str(checkpoint_dir))
+        ckpt = Checkpoint(checkpoint_dir, device, cfg.num_hidden_layers)
+        m = cls(cfg, device); m.eval()
+        for name, param in m.named_parameters():
+            loaded = ckpt.get(name)
+            if param.shape != loaded.shape:
+                raise ValueError(f"Shape mismatch {name}: {param.shape} vs {loaded.shape}")
+            param.data.copy_(loaded)
+        return m
+
+
+@dataclass(frozen=True)
+class LabelInfo:
+    boundary_label_lookup: dict
+    token_to_span_label: dict
+    token_boundary_tags: dict
+    span_class_names: tuple
+    span_label_lookup: dict
+    background_token_label: int
+    background_span_label: int
+
+
+def labels_to_spans(labels_by_index, label_info):
+    spans, cur_label, start_idx, prev_idx = [], None, None, None
+    bg = label_info.background_span_label
+    for ti in sorted(labels_by_index):
+        lid = labels_by_index[ti]
+        sl = label_info.token_to_span_label.get(lid)
+        bt = label_info.token_boundary_tags.get(lid)
+        if prev_idx is not None and ti != prev_idx + 1:
+            if cur_label is not None and start_idx is not None:
+                spans.append((cur_label, start_idx, prev_idx + 1))
+            cur_label = start_idx = None
+        if sl is None:
+            prev_idx = ti; continue
+        if sl == bg:
+            if cur_label is not None and start_idx is not None:
+                spans.append((cur_label, start_idx, ti))
+            cur_label = start_idx = None; prev_idx = ti; continue
+        if bt == "S":
+            if cur_label is not None and start_idx is not None and prev_idx is not None:
+                spans.append((cur_label, start_idx, prev_idx + 1))
+            spans.append((sl, ti, ti + 1)); cur_label = start_idx = None
+        elif bt == "B":
+            if cur_label is not None and start_idx is not None and prev_idx is not None:
+                spans.append((cur_label, start_idx, prev_idx + 1))
+            cur_label, start_idx = sl, ti
+        elif bt == "I":
+            if cur_label is None or cur_label != sl:
+                if cur_label is not None and start_idx is not None and prev_idx is not None:
+                    spans.append((cur_label, start_idx, prev_idx + 1))
+                cur_label, start_idx = sl, ti
+        elif bt == "E":
+            if cur_label is None or cur_label != sl or start_idx is None:
+                if cur_label is not None and start_idx is not None and prev_idx is not None:
+                    spans.append((cur_label, start_idx, prev_idx + 1))
+                spans.append((sl, ti, ti + 1)); cur_label = start_idx = None
+            else:
+                spans.append((cur_label, start_idx, ti + 1)); cur_label = start_idx = None
+        else:
+            if cur_label is not None and start_idx is not None and prev_idx is not None:
+                spans.append((cur_label, start_idx, prev_idx + 1))
+            cur_label = start_idx = None
+        prev_idx = ti
+    if cur_label is not None and start_idx is not None and prev_idx is not None:
+        spans.append((cur_label, start_idx, prev_idx + 1))
+    return spans
+
+
+def token_spans_to_char_spans(spans, cs, ce):
+    out = []
+    for li, ts, te in spans:
+        if not (0 <= ts < te <= len(cs)):
+            continue
+        s, e = cs[ts], ce[te - 1]
+        if e > s:
+            out.append((li, s, e))
+    return out
+
+
+def trim_char_spans_whitespace(spans, text):
+    out = []
+    for li, s, e in spans:
+        if not (0 <= s < e <= len(text)):
+            continue
+        while s < e and text[s].isspace(): s += 1
+        while e > s and text[e - 1].isspace(): e -= 1
+        if e > s:
+            out.append((li, s, e))
+    return out
+
+
+@functools.lru_cache(maxsize=1)
+def get_viterbi_transition_biases():
+    cp = MODEL_DIR / "viterbi_calibration.json"
+    default = {k: 0.0 for k in VITERBI_TRANSITION_BIAS_KEYS}
+    if not cp.is_file():
+        return default
+    payload = json.loads(cp.read_text())
+    raw = payload
+    ops = payload.get("operating_points")
+    if isinstance(ops, dict):
+        preset = ops.get(DEFAULT_VITERBI_CALIBRATION_PRESET)
+        if isinstance(preset, dict):
+            raw = preset.get("biases", raw)
+    if not isinstance(raw, dict):
+        return default
+    return {k: float(raw.get(k, 0.0)) for k in VITERBI_TRANSITION_BIAS_KEYS}
+
+
+class Decoder:
+    def __init__(self, label_info):
+        nc = len(label_info.token_to_span_label)
+        self._start = torch.full((nc,), -1e9, dtype=torch.float32)
+        self._end = torch.full((nc,), -1e9, dtype=torch.float32)
+        self._trans = torch.full((nc, nc), -1e9, dtype=torch.float32)
+        biases = get_viterbi_transition_biases()
+        bg_tok, bg_sp = label_info.background_token_label, label_info.background_span_label
+        ttsl, tbt = label_info.token_to_span_label, label_info.token_boundary_tags
+        for i in range(nc):
+            tag, sl = tbt.get(i), ttsl.get(i)
+            if tag in {"B", "S"} or i == bg_tok: self._start[i] = 0.0
+            if tag in {"E", "S"} or i == bg_tok: self._end[i] = 0.0
+            for j in range(nc):
+                nt, ns = tbt.get(j), ttsl.get(j)
+                if self._valid(tag, sl, nt, ns, bg_tok, bg_sp, j):
+                    self._trans[i, j] = self._bias(tag, sl, nt, ns, bg_sp, biases)
+
+    @staticmethod
+    def _valid(pt, ps, nt, ns, bti, bsi, ni):
+        nb = ns == bsi or ni == bti
+        if (ns is None or nt is None) and not nb: return False
+        if pt is None or ps is None: return nb or nt in {"B", "S"}
+        if ps == bsi or pt in {"E", "S"}: return nb or nt in {"B", "S"}
+        if pt in {"B", "I"}: return ps == ns and nt in {"I", "E"}
+        return False
+
+    @staticmethod
+    def _bias(pt, ps, nt, ns, bsi, b):
+        nb, pb = ns == bsi, ps == bsi
+        if pb: return b["transition_bias_background_stay"] if nb else b["transition_bias_background_to_start"]
+        if pt in {"B", "I"}: return b["transition_bias_inside_to_continue"] if nt == "I" else b["transition_bias_inside_to_end"]
+        return b["transition_bias_end_to_background"] if nb else b["transition_bias_end_to_start"]
+
+    def decode(self, lp):
+        sl, nc = lp.shape
+        if sl == 0: return []
+        st = self._start.to(lp.device, lp.dtype)
+        en = self._end.to(lp.device, lp.dtype)
+        tr = self._trans.to(lp.device, lp.dtype)
+        scores = lp[0] + st
+        bp = torch.empty((sl - 1, nc), device=lp.device, dtype=torch.int64)
+        for i in range(1, sl):
+            t = scores.unsqueeze(1) + tr
+            bs, bi = t.max(dim=0)
+            scores = bs + lp[i]; bp[i - 1] = bi
+        if not torch.isfinite(scores).any(): return lp.argmax(dim=1).tolist()
+        scores += en
+        path = torch.empty(sl, device=lp.device, dtype=torch.int64)
+        path[-1] = scores.argmax()
+        for i in range(sl - 2, -1, -1): path[i] = bp[i, path[i + 1]]
+        return path.tolist()
+
+
+@dataclass(frozen=True)
+class InferenceRuntime:
+    model: Transformer; encoding: tiktoken.Encoding; label_info: LabelInfo
+    device: torch.device; n_ctx: int
+
+
+@functools.lru_cache(maxsize=1)
+def get_runtime():
+    cp = MODEL_DIR
+    cfg = json.loads((cp / "config.json").read_text())
+    validate_model_config_contract(cfg, context=str(cp))
+    device = torch.device("cuda")
+    encoding = tiktoken.get_encoding(str(cfg["encoding"]).strip())
+    scn = [BACKGROUND_CLASS_LABEL]; sll = {BACKGROUND_CLASS_LABEL: 0}
+    bll, ttsl, tbt = {}, {}, {}
+    bg_idx = None
+    for idx, name in enumerate(NER_CLASS_NAMES):
+        if name == BACKGROUND_CLASS_LABEL:
+            bg_idx = idx; ttsl[idx] = 0; tbt[idx] = None; continue
+        bnd, base = name.split("-", 1)
+        si = sll.get(base)
+        if si is None:
+            si = len(scn); scn.append(base); sll[base] = si
+        ttsl[idx] = si; tbt[idx] = bnd
+        bll.setdefault(base, {})[bnd] = idx
+    li = LabelInfo(bll, ttsl, tbt, tuple(scn), sll, bg_idx, 0)
+    m = Transformer.from_checkpoint(str(cp), device=device)
+    return InferenceRuntime(m, encoding, li, device, int(cfg["default_n_ctx"]))
+
+
+@torch.inference_mode()
+def predict_text(runtime, text, decoder):
+    tids = tuple(int(t) for t in runtime.encoding.encode(text, allowed_special="all"))
+    if not tids: return text, []
+    scores = []
+    for s in range(0, len(tids), runtime.n_ctx):
+        e = min(s + runtime.n_ctx, len(tids))
+        wt = torch.tensor(tids[s:e], device=runtime.device, dtype=torch.int32)
+        lp = F.log_softmax(runtime.model(wt).float(), dim=-1)
+        scores.extend(lp.unbind(0))
+    stacked = torch.stack(scores, 0)
+    dl = decoder.decode(stacked)
+    if len(dl) != len(tids): dl = stacked.argmax(dim=1).tolist()
+    pli = {i: int(l) for i, l in enumerate(dl)}
+    pts = labels_to_spans(pli, runtime.label_info)
+    tb = [runtime.encoding.decode_single_token_bytes(t) for t in tids]
+    dt = b"".join(tb).decode("utf-8", errors="replace")
+    cbs, cbe = [], []
+    bc = 0
+    for ch in dt: cbs.append(bc); bc += len(ch.encode("utf-8")); cbe.append(bc)
+    cs, ce = [], []
+    tbc = 0
+    for rb in tb:
+        tbs = tbc; tbe = tbs + len(rb); tbc = tbe
+        cs.append(bisect_right(cbe, tbs)); ce.append(bisect_left(cbs, tbe))
+    pcs = token_spans_to_char_spans(pts, cs, ce)
+    pcs = trim_char_spans_whitespace(pcs, dt if dt != text else text)
+    src = dt if dt != text else text
+    detected = []
+    for li, s, e in pcs:
+        lbl = (runtime.label_info.span_class_names[li]
+               if 0 <= li < len(runtime.label_info.span_class_names) else f"label_{li}")
+        detected.append({"label": lbl, "start": s, "end": e, "text": src[s:e]})
+    return src, detected
+
+
+@spaces.GPU
+def run_pii_analysis(text):
+    runtime = get_runtime()
+    decoder = Decoder(label_info=runtime.label_info)
+    return predict_text(runtime, text, decoder)
+
+
+# =====================================================================
+# OCR  +  SPAN → BOX MAPPING
+# =====================================================================
+
+def ocr_image(img: Image.Image) -> dict:
+    """Run Tesseract and return the concatenated text plus per-word boxes.
+
+    The text is reconstructed with a single space between words on the same
+    line and a newline between lines, matching the character offsets we emit
+    in the `words` list — so later char-span → box mapping is a pure lookup.
+    """
+    data = pytesseract.image_to_data(img, output_type=pytesseract.Output.DICT)
+    words, parts = [], []
+    pos = 0
+    last_line_key = None
+    for i in range(len(data["text"])):
+        text = data["text"][i]
+        if text is None:
+            continue
+        text = text.strip()
+        if not text:
+            continue
+        try:
+            conf = float(data["conf"][i])
+        except (TypeError, ValueError):
+            conf = -1.0
+        if conf < 0:
+            continue
+        line_key = (data["block_num"][i], data["par_num"][i], data["line_num"][i])
+        if last_line_key is None:
+            pass
+        elif line_key != last_line_key:
+            parts.append("\n"); pos += 1
+        else:
+            parts.append(" "); pos += 1
+        last_line_key = line_key
+        start = pos
+        parts.append(text); pos += len(text)
+        words.append({
+            "text": text, "start": start, "end": pos,
+            "x": int(data["left"][i]), "y": int(data["top"][i]),
+            "w": int(data["width"][i]), "h": int(data["height"][i]),
+        })
+    return {"text": "".join(parts), "words": words}
+
+
+def map_spans_to_boxes(words, spans, pad=3):
+    """Map each char span to one or more pixel boxes, splitting across lines."""
+    boxes = []
+    for span in spans:
+        ss, se, lbl = span["start"], span["end"], span["label"]
+        hits = [w for w in words if w["start"] < se and w["end"] > ss]
+        if not hits:
+            continue
+        by_line = {}
+        for w in hits:
+            yc = w["y"] + w["h"] // 2
+            matched = None
+            for key in by_line:
+                if abs(key - yc) < max(w["h"] * 0.6, 10):
+                    matched = key; break
+            key = matched if matched is not None else yc
+            by_line.setdefault(key, []).append(w)
+        for line_words in by_line.values():
+            x1 = min(w["x"] for w in line_words) - pad
+            y1 = min(w["y"] for w in line_words) - pad
+            x2 = max(w["x"] + w["w"] for w in line_words) + pad
+            y2 = max(w["y"] + w["h"] for w in line_words) + pad
+            boxes.append({
+                "x": max(0, x1), "y": max(0, y1),
+                "w": max(1, x2 - x1), "h": max(1, y2 - y1),
+                "label": lbl,
+                "text": " ".join(w["text"] for w in line_words),
+            })
+    return boxes
+
+
+# =====================================================================
+# SERVER
+# =====================================================================
+
+server = gr.Server()
+
+
+@server.get("/", response_class=HTMLResponse)
+async def homepage():
+    return FRONTEND_HTML
+
+
+@server.post("/api/detect")
+async def detect(file: UploadFile = File(...)):
+    suffix = Path(file.filename or "").suffix.lower()
+    if suffix not in (".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"):
+        return JSONResponse({"error": f"Unsupported image type: {suffix or '(none)'}"}, 400)
+    try:
+        img_bytes = await file.read()
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    except Exception as e:
+        return JSONResponse({"error": f"Could not read image: {e}"}, 400)
+
+    ocr = ocr_image(img)
+    if not ocr["text"].strip():
+        return JSONResponse({"error": "No text detected in the image."}, 400)
+
+    try:
+        source_text, spans = run_pii_analysis(ocr["text"])
+    except Exception as e:
+        return JSONResponse({"error": f"PII analysis failed: {e}"}, 500)
+
+    # If the model round-tripped the text differently, fall back to original
+    # so the word offsets still line up. In practice this only matters for
+    # exotic unicode that tesseract won't produce anyway.
+    if source_text != ocr["text"]:
+        spans = [s for s in spans if s["end"] <= len(ocr["text"])]
+
+    boxes = map_spans_to_boxes(ocr["words"], spans)
+
+    buf = io.BytesIO(); img.save(buf, format="PNG")
+    data_url = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+    return JSONResponse({
+        "filename": file.filename,
+        "image": data_url,
+        "width": img.width, "height": img.height,
+        "boxes": boxes,
+        "text": ocr["text"],
+        "spans": spans,
+        "categories_meta": {k: {"color": v["color"], "label": v["label"]}
+                            for k, v in CATEGORIES_META.items()},
+    })
+
+
+@server.api(name="anonymize_screenshot")
+def anonymize_screenshot_api(image_path: str) -> str:
+    """Gradio API: takes a path to an image, returns JSON with detected boxes."""
+    img = Image.open(image_path).convert("RGB")
+    ocr = ocr_image(img)
+    if not ocr["text"].strip():
+        return json.dumps({"boxes": [], "text": "", "spans": []})
+    _, spans = run_pii_analysis(ocr["text"])
+    boxes = map_spans_to_boxes(ocr["words"], spans)
+    return json.dumps({
+        "width": img.width, "height": img.height,
+        "boxes": boxes, "text": ocr["text"], "spans": spans,
+    }, ensure_ascii=False)
+
+
+# =====================================================================
+# FRONTEND
+# =====================================================================
+
+FRONTEND_HTML = r"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width,initial-scale=1">
+<title>Screenshot Anonymizer</title>
+<link rel="preconnect" href="https://fonts.googleapis.com">
+<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap" rel="stylesheet">
+<style>
+*,*::before,*::after{box-sizing:border-box;margin:0;padding:0}
+:root{
+  --bg:#0f172a;--surface:#1e293b;--surface2:#334155;--border:#334155;--border2:#475569;
+  --text:#f1f5f9;--text2:#cbd5e1;--text3:#94a3b8;
+  --primary:#6366f1;--primary-light:#818cf8;--danger:#ef4444;--success:#22c55e;
+  --radius:12px;--radius-sm:8px;--shadow:0 4px 12px rgba(0,0,0,.3);--shadow-lg:0 12px 40px rgba(0,0,0,.4);
+}
+html,body{height:100%}
+body{font-family:'Inter',system-ui,sans-serif;background:var(--bg);color:var(--text);min-height:100vh;line-height:1.5;overflow:hidden}
+
+/* ── upload view ── */
+#upload-view{display:flex;flex-direction:column;align-items:center;justify-content:center;min-height:100vh;padding:2rem;background:radial-gradient(circle at 30% 20%,#312e81 0%,var(--bg) 50%)}
+.upload-card{background:var(--surface);border:1px solid var(--border);border-radius:20px;padding:3rem;max-width:640px;width:100%;text-align:center;box-shadow:var(--shadow-lg)}
+.brand{display:flex;align-items:center;justify-content:center;gap:.75rem;margin-bottom:.5rem}
+.brand h1{font-size:1.75rem;font-weight:800;background:linear-gradient(135deg,#818cf8,#ec4899);-webkit-background-clip:text;-webkit-text-fill-color:transparent}
+.brand-icon{width:40px;height:40px;background:linear-gradient(135deg,var(--primary),#ec4899);border-radius:10px;display:flex;align-items:center;justify-content:center;color:#fff;font-size:1.3rem}
+.subtitle{color:var(--text2);margin-bottom:2rem;font-size:1rem}
+.dropzone{border:2px dashed var(--border2);border-radius:var(--radius);padding:3rem 2rem;cursor:pointer;transition:all .2s;position:relative;background:rgba(15,23,42,.4)}
+.dropzone:hover,.dropzone.dragover{border-color:var(--primary-light);background:rgba(99,102,241,.08)}
+.dropzone-icon{font-size:2.5rem;margin-bottom:.75rem}
+.dropzone-text{font-weight:600;font-size:1.05rem;margin-bottom:.25rem}
+.dropzone-hint{color:var(--text3);font-size:.85rem}
+.dropzone input{position:absolute;inset:0;opacity:0;cursor:pointer}
+.features{display:grid;grid-template-columns:repeat(3,1fr);gap:.75rem;margin-top:1.75rem;text-align:left}
+.feature{background:rgba(15,23,42,.5);border:1px solid var(--border);padding:.9rem;border-radius:var(--radius-sm)}
+.feature-title{font-weight:600;font-size:.8rem;margin-bottom:.2rem;color:var(--text)}
+.feature-desc{color:var(--text3);font-size:.72rem;line-height:1.45}
+.powered-by{margin-top:1.5rem;font-size:.78rem;color:var(--text3)}
+.powered-by strong{color:var(--text2)}
+
+/* ── editor view ── */
+#editor-view{display:none;flex-direction:column;height:100vh}
+.top-bar{background:var(--surface);border-bottom:1px solid var(--border);padding:.65rem 1.25rem;display:flex;align-items:center;gap:1rem;flex-shrink:0}
+.top-bar .brand{margin:0}
+.top-bar .brand h1{font-size:1.15rem}
+.top-bar .brand-icon{width:30px;height:30px;font-size:.95rem}
+.file-info{font-size:.85rem;color:var(--text3);margin-left:.25rem;flex:1;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}
+.btn{padding:.5rem 1rem;border-radius:var(--radius-sm);border:1px solid var(--border2);cursor:pointer;font-weight:600;font-size:.82rem;transition:all .15s;background:var(--surface2);color:var(--text);font-family:inherit}
+.btn:hover{background:var(--border2);border-color:var(--text3)}
+.btn-primary{background:var(--primary);border-color:var(--primary);color:#fff}
+.btn-primary:hover{background:var(--primary-light);border-color:var(--primary-light)}
+.btn-ghost{background:transparent}
+.btn-ghost:hover{background:var(--surface2)}
+
+.editor-layout{flex:1;display:flex;min-height:0}
+.canvas-area{flex:1;background:#020617;background-image:linear-gradient(45deg,#0f172a 25%,transparent 25%),linear-gradient(-45deg,#0f172a 25%,transparent 25%),linear-gradient(45deg,transparent 75%,#0f172a 75%),linear-gradient(-45deg,transparent 75%,#0f172a 75%);background-size:20px 20px;background-position:0 0,0 10px,10px -10px,10px 0px;overflow:auto;position:relative;display:flex;align-items:center;justify-content:center;padding:2rem}
+.canvas-wrap{position:relative;box-shadow:0 20px 60px rgba(0,0,0,.5);cursor:crosshair;flex-shrink:0}
+.canvas-wrap.mode-move{cursor:grab}
+.canvas-wrap.mode-move.dragging{cursor:grabbing}
+.canvas-wrap canvas{display:block;border-radius:4px}
+
+.sidebar{width:300px;background:var(--surface);border-left:1px solid var(--border);overflow-y:auto;flex-shrink:0;display:flex;flex-direction:column}
+.sidebar section{padding:1rem 1.25rem;border-bottom:1px solid var(--border)}
+.sidebar section:last-child{border-bottom:none}
+.sidebar h3{font-size:.68rem;text-transform:uppercase;letter-spacing:.8px;color:var(--text3);margin-bottom:.65rem;font-weight:700}
+
+.tool-row{display:grid;grid-template-columns:1fr 1fr;gap:.4rem;margin-bottom:.5rem}
+.tool-btn{padding:.55rem .3rem;border-radius:var(--radius-sm);border:1px solid var(--border);background:transparent;color:var(--text2);cursor:pointer;font-size:.78rem;font-weight:600;display:flex;flex-direction:column;align-items:center;gap:.25rem;transition:all .15s;font-family:inherit}
+.tool-btn:hover{background:var(--surface2);color:var(--text)}
+.tool-btn.active{background:rgba(99,102,241,.15);border-color:var(--primary);color:var(--primary-light)}
+.tool-btn .ico{font-size:1.1rem}
+.tool-hint{font-size:.72rem;color:var(--text3);line-height:1.4;margin-top:.35rem}
+
+.stat-grid{display:grid;grid-template-columns:repeat(2,1fr);gap:.5rem;margin-bottom:.65rem}
+.stat{background:rgba(15,23,42,.5);border:1px solid var(--border);padding:.6rem .75rem;border-radius:var(--radius-sm)}
+.stat .num{font-size:1.4rem;font-weight:800;color:var(--primary-light);line-height:1}
+.stat .lbl{font-size:.65rem;color:var(--text3);text-transform:uppercase;letter-spacing:.4px;margin-top:.2rem}
+
+.filter-item{display:flex;align-items:center;gap:.55rem;padding:.4rem .45rem;border-radius:var(--radius-sm);cursor:pointer;transition:background .15s;user-select:none}
+.filter-item:hover{background:rgba(15,23,42,.5)}
+.filter-item input{display:none}
+.filter-check{width:16px;height:16px;border-radius:4px;border:2px solid var(--border2);display:flex;align-items:center;justify-content:center;transition:all .15s;flex-shrink:0}
+.filter-item input:checked~.filter-check{border-color:currentColor;background:currentColor}
+.filter-item input:checked~.filter-check::after{content:'';display:block;width:4px;height:8px;border:solid #fff;border-width:0 2px 2px 0;transform:rotate(45deg) translateY(-1px)}
+.filter-dot{width:9px;height:9px;border-radius:50%;flex-shrink:0}
+.filter-label{flex:1;font-size:.82rem;font-weight:500;color:var(--text)}
+.filter-count{font-size:.72rem;color:var(--text3);font-weight:600;background:var(--surface2);padding:.1rem .4rem;border-radius:10px}
+
+.action-row{display:flex;flex-direction:column;gap:.45rem}
+.action-row .btn{width:100%;justify-content:center;text-align:center}
+
+.empty-state{color:var(--text3);font-size:.8rem;font-style:italic;padding:.25rem 0}
+
+/* ── loading ── */
+#loading{position:fixed;inset:0;background:rgba(15,23,42,.8);backdrop-filter:blur(8px);display:none;flex-direction:column;align-items:center;justify-content:center;z-index:9999}
+.spinner{width:44px;height:44px;border:3px solid var(--border2);border-top-color:var(--primary-light);border-radius:50%;animation:spin .8s linear infinite}
+@keyframes spin{to{transform:rotate(360deg)}}
+#loading p{margin-top:1rem;font-weight:600;color:var(--text)}
+.progress-text{font-size:.82rem;color:var(--text3);margin-top:.25rem}
+.error-banner{background:rgba(239,68,68,.1);border:1px solid rgba(239,68,68,.4);color:#fca5a5;padding:.75rem 1rem;border-radius:var(--radius-sm);margin:.75rem 1.25rem;font-size:.85rem;display:none}
+
+.toast{position:fixed;bottom:1.5rem;left:50%;transform:translateX(-50%) translateY(100px);background:var(--surface);border:1px solid var(--border2);color:var(--text);padding:.7rem 1.25rem;border-radius:var(--radius-sm);font-size:.85rem;font-weight:500;box-shadow:var(--shadow-lg);transition:transform .25s ease;z-index:9998}
+.toast.show{transform:translateX(-50%) translateY(0)}
+.toast.success{border-color:rgba(34,197,94,.5)}
+
+@media(max-width:900px){
+  .editor-layout{flex-direction:column}
+  .sidebar{width:100%;border-left:none;border-top:1px solid var(--border);max-height:40vh}
+  .features{grid-template-columns:1fr}
+}
+</style>
+</head>
+<body>
+
+<!-- Upload view -->
+<div id="upload-view">
+  <div class="upload-card">
+    <div class="brand">
+      <div class="brand-icon">&#x1f576;</div>
+      <h1>Screenshot Anonymizer</h1>
+    </div>
+    <p class="subtitle">Redact PII in images &middot; OCR + OpenAI Privacy Filter</p>
+    <div class="dropzone" id="dropzone">
+      <div class="dropzone-icon">&#x1f5bc;</div>
+      <div class="dropzone-text">Drop a screenshot here</div>
+      <div class="dropzone-hint">PNG, JPG, WebP &middot; chat, email, document</div>
+      <input type="file" id="file-input" accept="image/png,image/jpeg,image/webp,image/bmp,image/tiff">
+    </div>
+    <div class="features">
+      <div class="feature"><div class="feature-title">Auto-detect</div><div class="feature-desc">OCR reads the text, the model finds names, emails, phones, secrets.</div></div>
+      <div class="feature"><div class="feature-title">Edit freely</div><div class="feature-desc">Drag to add bars, click to toggle, Delete to remove. Nothing is sent back.</div></div>
+      <div class="feature"><div class="feature-title">Export</div><div class="feature-desc">Save as PNG or copy to clipboard &mdash; ready for Twitter.</div></div>
+    </div>
+    <div class="powered-by">Powered by <strong>OpenAI Privacy Filter</strong> + <strong>Tesseract OCR</strong></div>
+  </div>
+</div>
+
+<!-- Editor view -->
+<div id="editor-view">
+  <div class="top-bar">
+    <div class="brand"><div class="brand-icon">&#x1f576;</div><h1>Screenshot Anonymizer</h1></div>
+    <div class="file-info" id="file-info"></div>
+    <button class="btn btn-ghost" onclick="resetView()">New Screenshot</button>
+  </div>
+  <div class="error-banner" id="error-banner"></div>
+  <div class="editor-layout">
+    <div class="canvas-area" id="canvas-area">
+      <div class="canvas-wrap" id="canvas-wrap">
+        <canvas id="canvas"></canvas>
+      </div>
+    </div>
+    <div class="sidebar">
+
+      <section>
+        <h3>Tool</h3>
+        <div class="tool-row">
+          <button class="tool-btn active" data-mode="add" id="tool-add"><span class="ico">&#x270f;</span><span>Draw</span></button>
+          <button class="tool-btn" data-mode="move" id="tool-move"><span class="ico">&#x2725;</span><span>Select</span></button>
+        </div>
+        <div class="tool-hint" id="tool-hint">Drag anywhere to draw a black bar. Click a bar to select it.</div>
+      </section>
+
+      <section>
+        <h3>Detected</h3>
+        <div class="stat-grid">
+          <div class="stat"><div class="num" id="stat-boxes">0</div><div class="lbl">Total Bars</div></div>
+          <div class="stat"><div class="num" id="stat-cats">0</div><div class="lbl">Categories</div></div>
+        </div>
+      </section>
+
+      <section>
+        <h3>Categories</h3>
+        <div id="category-filters">
+          <div class="empty-state" id="cat-empty">No PII detected</div>
+        </div>
+      </section>
+
+      <section>
+        <h3>Export</h3>
+        <div class="action-row">
+          <button class="btn btn-primary" onclick="downloadImage()">&#x2b07; Download PNG</button>
+          <button class="btn" onclick="copyToClipboard()">&#x1f4cb; Copy to Clipboard</button>
+        </div>
+      </section>
+
+    </div>
+  </div>
+</div>
+
+<div id="loading"><div class="spinner"></div><p>Analyzing screenshot&hellip;</p><div class="progress-text">OCR &rarr; Privacy Filter &rarr; Map to pixels</div></div>
+<div class="toast" id="toast"></div>
+
+<script>
+const State = {
+  img: null,            // HTMLImageElement of the screenshot
+  width: 0, height: 0,
+  boxes: [],            // [{id, x, y, w, h, label, enabled, custom}]
+  nextId: 1,
+  activeCats: new Set(),
+  catMeta: {},
+  mode: 'add',          // 'add' | 'move'
+  selected: null,       // box id
+  drag: null,           // {type:'draw'|'move', startX, startY, origBox?, newBox?}
+  scale: 1,             // display scale (natural px -> screen px)
+  filename: '',
+  catCounts: {},
+};
+
+const LABEL_NAMES = {private_person:'Person',private_address:'Address',private_email:'Email',private_phone:'Phone',private_url:'URL',private_date:'Date',account_number:'Account',secret:'Secret'};
+
+// ── upload flow ───────────────────────────────────────────────────
+const dz = document.getElementById('dropzone');
+const fi = document.getElementById('file-input');
+['dragenter','dragover'].forEach(e=>dz.addEventListener(e,ev=>{ev.preventDefault();dz.classList.add('dragover')}));
+['dragleave','drop'].forEach(e=>dz.addEventListener(e,ev=>{ev.preventDefault();dz.classList.remove('dragover')}));
+dz.addEventListener('drop',ev=>{if(ev.dataTransfer.files[0]) uploadFile(ev.dataTransfer.files[0])});
+fi.addEventListener('change',ev=>{if(ev.target.files[0]) uploadFile(ev.target.files[0])});
+
+// Also accept a paste (ctrl+v) when on upload view
+document.addEventListener('paste',ev=>{
+  if(document.getElementById('upload-view').style.display==='none') return;
+  const items = ev.clipboardData && ev.clipboardData.items;
+  if(!items) return;
+  for(const it of items){
+    if(it.type && it.type.startsWith('image/')){
+      const f = it.getAsFile(); if(f) uploadFile(f); ev.preventDefault(); return;
+    }
+  }
+});
+
+async function uploadFile(file){
+  if(!file.type || !file.type.startsWith('image/')){showError('Please drop an image file.'); return;}
+  document.getElementById('loading').style.display='flex';
+  document.getElementById('upload-view').style.display='none';
+  const form = new FormData(); form.append('file', file);
+  try{
+    const r = await fetch('/api/detect', {method:'POST', body:form});
+    const d = await r.json();
+    if(d.error){showError(d.error); return;}
+    await initEditor(d);
+  }catch(e){showError('Analysis failed: '+e.message);}
+  finally{document.getElementById('loading').style.display='none';}
+}
+
+function showError(m){
+  document.getElementById('loading').style.display='none';
+  document.getElementById('editor-view').style.display='flex';
+  const b = document.getElementById('error-banner');
+  b.textContent = m; b.style.display = 'block';
+}
+
+function resetView(){
+  document.getElementById('editor-view').style.display='none';
+  document.getElementById('upload-view').style.display='flex';
+  document.getElementById('error-banner').style.display='none';
+  fi.value='';
+  State.boxes = []; State.selected = null; State.img = null;
+}
+
+// ── init editor ───────────────────────────────────────────────────
+async function initEditor(data){
+  document.getElementById('editor-view').style.display='flex';
+  document.getElementById('error-banner').style.display='none';
+  document.getElementById('file-info').textContent = data.filename || '';
+  State.filename = data.filename || 'redacted.png';
+  State.width = data.width; State.height = data.height;
+  State.catMeta = data.categories_meta || {};
+  State.catCounts = {};
+  State.boxes = (data.boxes || []).map(b => {
+    State.catCounts[b.label] = (State.catCounts[b.label]||0) + 1;
+    return {id: State.nextId++, x:b.x, y:b.y, w:b.w, h:b.h, label:b.label, text:b.text, enabled:true, custom:false};
+  });
+  State.activeCats = new Set(Object.keys(State.catCounts));
+
+  const img = new Image();
+  await new Promise((res, rej) => { img.onload = res; img.onerror = rej; img.src = data.image; });
+  State.img = img;
+
+  setupCanvas();
+  renderCategoryFilters();
+  updateStats();
+  draw();
+}
+
+function setupCanvas(){
+  const cv = document.getElementById('canvas');
+  cv.width = State.width;
+  cv.height = State.height;
+  // scale to fit view area
+  const area = document.getElementById('canvas-area');
+  const maxW = area.clientWidth - 64, maxH = area.clientHeight - 64;
+  const scale = Math.min(1, maxW/State.width, maxH/State.height);
+  State.scale = scale;
+  cv.style.width = (State.width * scale) + 'px';
+  cv.style.height = (State.height * scale) + 'px';
+}
+
+window.addEventListener('resize', ()=>{ if(State.img){ setupCanvas(); draw(); } });
+
+// ── drawing ───────────────────────────────────────────────────────
+function draw(){
+  const cv = document.getElementById('canvas');
+  const ctx = cv.getContext('2d');
+  ctx.clearRect(0,0,cv.width,cv.height);
+  ctx.drawImage(State.img, 0, 0);
+
+  // draw enabled bars
+  for(const b of State.boxes){
+    if(!isVisible(b)) continue;
+    ctx.fillStyle = '#000';
+    ctx.fillRect(b.x, b.y, b.w, b.h);
+  }
+  // draw selection outline
+  const sel = selectedBox();
+  if(sel){
+    ctx.save();
+    ctx.strokeStyle = '#818cf8';
+    ctx.lineWidth = Math.max(2, 2/State.scale);
+    ctx.setLineDash([6/State.scale, 4/State.scale]);
+    ctx.strokeRect(sel.x, sel.y, sel.w, sel.h);
+    ctx.restore();
+  }
+  // draw in-progress new rectangle
+  if(State.drag && State.drag.type==='draw' && State.drag.newBox){
+    const b = State.drag.newBox;
+    ctx.save();
+    ctx.fillStyle = 'rgba(0,0,0,.75)';
+    ctx.fillRect(b.x, b.y, b.w, b.h);
+    ctx.strokeStyle = '#818cf8';
+    ctx.lineWidth = Math.max(1.5, 1.5/State.scale);
+    ctx.strokeRect(b.x, b.y, b.w, b.h);
+    ctx.restore();
+  }
+}
+
+function isVisible(b){
+  if(!b.enabled) return false;
+  if(!b.custom && !State.activeCats.has(b.label)) return false;
+  return true;
+}
+
+function selectedBox(){
+  if(State.selected==null) return null;
+  return State.boxes.find(b => b.id === State.selected) || null;
+}
+
+// ── mouse interaction ─────────────────────────────────────────────
+const wrap = document.getElementById('canvas-wrap');
+
+function toCanvasXY(ev){
+  const rect = wrap.getBoundingClientRect();
+  const x = (ev.clientX - rect.left) / State.scale;
+  const y = (ev.clientY - rect.top) / State.scale;
+  return {x, y};
+}
+
+function hitTest(x, y){
+  // Topmost first (last drawn is on top)
+  for(let i = State.boxes.length - 1; i >= 0; i--){
+    const b = State.boxes[i];
+    if(!isVisible(b)) continue;
+    if(x >= b.x && x <= b.x + b.w && y >= b.y && y <= b.y + b.h) return b;
+  }
+  return null;
+}
+
+wrap.addEventListener('mousedown', ev => {
+  if(ev.button !== 0) return;
+  ev.preventDefault();
+  const {x, y} = toCanvasXY(ev);
+  const hit = hitTest(x, y);
+
+  if(State.mode === 'add' && !hit){
+    // start new rectangle
+    State.selected = null;
+    State.drag = {type:'draw', startX:x, startY:y, newBox:{x,y,w:0,h:0}};
+  } else if(hit){
+    State.selected = hit.id;
+    State.drag = {type:'move', startX:x, startY:y, origBox:{x:hit.x, y:hit.y, w:hit.w, h:hit.h}, boxId:hit.id};
+    wrap.classList.add('dragging');
+  } else {
+    State.selected = null;
+    State.drag = null;
+  }
+  draw();
+});
+
+window.addEventListener('mousemove', ev => {
+  if(!State.drag) return;
+  const {x, y} = toCanvasXY(ev);
+  if(State.drag.type === 'draw'){
+    const sx = State.drag.startX, sy = State.drag.startY;
+    State.drag.newBox = {
+      x: Math.min(sx, x), y: Math.min(sy, y),
+      w: Math.abs(x - sx), h: Math.abs(y - sy)
+    };
+  } else if(State.drag.type === 'move'){
+    const dx = x - State.drag.startX, dy = y - State.drag.startY;
+    const b = State.boxes.find(b => b.id === State.drag.boxId);
+    if(b){
+      const o = State.drag.origBox;
+      b.x = Math.max(0, Math.min(State.width - o.w, o.x + dx));
+      b.y = Math.max(0, Math.min(State.height - o.h, o.y + dy));
+    }
+  }
+  draw();
+});
+
+window.addEventListener('mouseup', ev => {
+  if(!State.drag) return;
+  wrap.classList.remove('dragging');
+  if(State.drag.type === 'draw'){
+    const b = State.drag.newBox;
+    if(b.w > 3 && b.h > 3){
+      const nb = {id: State.nextId++, x: Math.round(b.x), y: Math.round(b.y), w: Math.round(b.w), h: Math.round(b.h),
+                  label: 'custom', text: '', enabled: true, custom: true};
+      State.boxes.push(nb);
+      State.selected = nb.id;
+      updateStats();
+    }
+  }
+  State.drag = null;
+  draw();
+});
+
+// Keyboard: Delete to remove, Esc to deselect, D/S tool toggle
+window.addEventListener('keydown', ev => {
+  if(document.getElementById('editor-view').style.display === 'none') return;
+  if(ev.key === 'Delete' || ev.key === 'Backspace'){
+    if(State.selected != null){
+      State.boxes = State.boxes.filter(b => b.id !== State.selected);
+      State.selected = null;
+      updateStats();
+      draw();
+      ev.preventDefault();
+    }
+  } else if(ev.key === 'Escape'){
+    State.selected = null; State.drag = null; draw();
+  } else if(ev.key === 'd' || ev.key === 'D'){
+    setMode('add');
+  } else if(ev.key === 's' || ev.key === 'S' || ev.key === 'v' || ev.key === 'V'){
+    setMode('move');
+  }
+});
+
+// ── tool mode ─────────────────────────────────────────────────────
+document.querySelectorAll('.tool-btn').forEach(btn => {
+  btn.addEventListener('click', () => setMode(btn.dataset.mode));
+});
+
+function setMode(m){
+  State.mode = m;
+  document.querySelectorAll('.tool-btn').forEach(b => b.classList.toggle('active', b.dataset.mode === m));
+  wrap.classList.toggle('mode-move', m === 'move');
+  document.getElementById('tool-hint').textContent = (m === 'add')
+    ? 'Drag anywhere to draw a black bar. Click a bar to select it.'
+    : 'Click a bar to select, drag to move, Delete to remove.';
+}
+
+// ── category filters ──────────────────────────────────────────────
+function renderCategoryFilters(){
+  const ct = document.getElementById('category-filters');
+  const empty = document.getElementById('cat-empty');
+  const cats = Object.keys(State.catCounts);
+  if(!cats.length){ ct.innerHTML = ''; ct.appendChild(empty); empty.style.display='block'; return; }
+  ct.innerHTML = '';
+  for(const cat of cats){
+    const meta = State.catMeta[cat] || {color:'#888', label:cat};
+    const count = State.catCounts[cat];
+    const el = document.createElement('label');
+    el.className = 'filter-item';
+    el.style.color = meta.color;
+    el.innerHTML = `<input type="checkbox" ${State.activeCats.has(cat)?'checked':''}>
+      <span class="filter-check"></span>
+      <span class="filter-dot" style="background:${meta.color}"></span>
+      <span class="filter-label">${meta.label}</span>
+      <span class="filter-count">${count}</span>`;
+    el.querySelector('input').addEventListener('change', ev => {
+      if(ev.target.checked) State.activeCats.add(cat);
+      else State.activeCats.delete(cat);
+      draw();
+      updateStats();
+    });
+    ct.appendChild(el);
+  }
+}
+
+function updateStats(){
+  const visible = State.boxes.filter(isVisible).length;
+  document.getElementById('stat-boxes').textContent = visible;
+  const cats = new Set();
+  for(const b of State.boxes) if(isVisible(b)) cats.add(b.custom ? 'custom' : b.label);
+  document.getElementById('stat-cats').textContent = cats.size;
+}
+
+// ── export ────────────────────────────────────────────────────────
+function renderExportCanvas(){
+  const ec = document.createElement('canvas');
+  ec.width = State.width; ec.height = State.height;
+  const ctx = ec.getContext('2d');
+  ctx.drawImage(State.img, 0, 0);
+  ctx.fillStyle = '#000';
+  for(const b of State.boxes) if(isVisible(b)) ctx.fillRect(b.x, b.y, b.w, b.h);
+  return ec;
+}
+
+function downloadImage(){
+  const ec = renderExportCanvas();
+  ec.toBlob(blob => {
+    const a = document.createElement('a');
+    const base = (State.filename || 'screenshot').replace(/\.[^/.]+$/, '');
+    a.download = base + '-redacted.png';
+    a.href = URL.createObjectURL(blob);
+    a.click();
+    setTimeout(() => URL.revokeObjectURL(a.href), 1000);
+    toast('Saved ' + a.download, true);
+  }, 'image/png');
+}
+
+async function copyToClipboard(){
+  const ec = renderExportCanvas();
+  try{
+    await new Promise((res, rej) => {
+      ec.toBlob(async blob => {
+        try{
+          if(!navigator.clipboard || !window.ClipboardItem){ rej(new Error('Clipboard not supported')); return; }
+          await navigator.clipboard.write([new ClipboardItem({'image/png': blob})]);
+          res();
+        }catch(e){ rej(e); }
+      }, 'image/png');
+    });
+    toast('Copied to clipboard', true);
+  }catch(e){
+    toast('Copy failed: ' + e.message);
+  }
+}
+
+let toastTimer = null;
+function toast(msg, ok){
+  const t = document.getElementById('toast');
+  t.textContent = msg;
+  t.classList.toggle('success', !!ok);
+  t.classList.add('show');
+  clearTimeout(toastTimer);
+  toastTimer = setTimeout(() => t.classList.remove('show'), 2200);
+}
+</script>
+</body>
+</html>"""
+
+
+if __name__ == "__main__":
+    server.launch(server_name="0.0.0.0", server_port=7860)

packages.txt

@@ -0,0 +1 @@
+tesseract-ocr

requirements.txt

@@ -0,0 +1,11 @@
+tiktoken
+sentencepiece
+torch
+safetensors
+huggingface_hub
+gradio[mcp]
+pytesseract
+Pillow
+python-multipart
+accelerate
+spaces