| """Read actual LTF MIL attention and render its original-image coordinates.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image, ImageDraw, ImageFont |
|
|
|
|
| def capture_prediction(detector, path, branch, threshold): |
| """Observe one ordinary prediction; temporary hooks are always removed. |
| |
| Call under the app's model lock. return_viz only appends existing model |
| tensors to the output; it does not change logits or rerun the backbone. |
| """ |
| captured = {} |
|
|
| def enable_viz(module, args, kwargs): |
| return args, dict(kwargs, return_viz=True) |
|
|
| def collect(module, args, output): |
| viz = output["viz"] |
| for key in ("tile_attn", "level1_patch_attn", "level2_attn", "gate"): |
| captured[key] = viz[key].detach().float().cpu().numpy() |
|
|
| def collect_heads(module, args, output): |
| captured["head_attn"] = output[1][0, :, 0, :].detach().float().cpu().numpy() |
|
|
| handles = [] |
| try: |
| handles.append( |
| detector.model.register_forward_pre_hook(enable_viz, with_kwargs=True) |
| ) |
| handles.append(detector.model.register_forward_hook(collect)) |
| handles.append( |
| detector.model.core.level1_aggregate_attn.register_forward_hook( |
| collect_heads |
| ) |
| ) |
| result = detector.predict(path, branch, threshold) |
| finally: |
| for handle in handles: |
| handle.remove() |
| nt = result["views"] - 1 |
| patches = captured["level1_patch_attn"].reshape(nt, -1) |
| side = math.isqrt(patches.shape[1]) |
| if side * side != patches.shape[1]: |
| raise ValueError("Patch attention does not form a square grid.") |
| attention = dict( |
| l1=captured["tile_attn"].reshape(nt), |
| l1_heads=captured["head_attn"], |
| l2=captured["level2_attn"].reshape(nt + 1), |
| patch=patches.reshape(nt, side, side), |
| gate=float(captured["gate"].reshape(-1)[0]), |
| ) |
| for key in ("l1", "l2", "patch", "l1_heads"): |
| weights = attention[key] |
| if not np.isfinite(weights).all() or (weights < 0).any(): |
| raise ValueError("Attention contains invalid values.") |
| return result, attention |
|
|
|
|
| def tile_boxes(width, height, tile_size=224): |
| |
| from utils.tile_dataset_srm import get_start_positions |
|
|
| return [ |
| (x, y, x + tile_size, y + tile_size) |
| for y in get_start_positions(height, tile_size) |
| for x in get_start_positions(width, tile_size) |
| ] |
|
|
|
|
| def relative_scale(values): |
| values = np.asarray(values, dtype=np.float32) |
| low, high = float(values.min()), float(values.max()) |
| if high - low < 1e-12: |
| return np.full_like(values, 0.5) |
| return (values - low) / (high - low) |
|
|
|
|
| def colors(values): |
| """Fixed blue → neutral → orange-red scale, centered on uniform attention.""" |
| x = np.clip(np.asarray(values), 0, 1) |
| stops = np.array( |
| [[35, 57, 93], [104, 134, 171], [180, 191, 203], [247, 163, 66], [238, 80, 60]], |
| dtype=np.float32, |
| ) |
| first = np.minimum((x * 4).astype(int), 3) |
| fraction = (x * 4 - first)[..., None] |
| return (stops[first] * (1 - fraction) + stops[first + 1] * fraction).astype( |
| np.uint8 |
| ) |
|
|
|
|
| def attention_levels(values, reference_count): |
| """A common reference: 1/N is neutral; clip only outside [0.25×, 4×].""" |
| ratios = np.maximum(np.asarray(values, dtype=np.float32) * reference_count, 1e-12) |
| return np.clip((np.log2(ratios) + 2) / 4, 0, 1) |
|
|
|
|
| def preview_image(path): |
| with Image.open(path) as source: |
| image = source.convert("RGB") |
| image.thumbnail((1200, 1200)) |
| return image |
|
|
|
|
| def draw_grid(path, boxes): |
| image = preview_image(path) |
| with Image.open(path) as original: |
| sx, sy = image.width / original.width, image.height / original.height |
| draw = ImageDraw.Draw(image) |
| font = ImageFont.load_default(size=18) |
| labels = [] |
| for index, (x0, y0, x1, y1) in enumerate(boxes, 1): |
| rect = ( |
| int(x0 * sx), |
| int(y0 * sy), |
| min(image.width - 1, int(x1 * sx) - 1), |
| min(image.height - 1, int(y1 * sy) - 1), |
| ) |
| draw.rectangle(rect, outline=(11, 35, 59), width=4) |
| draw.rectangle(rect, outline=(234, 252, 255), width=2) |
| if len(boxes) > 80 and (index - 1) % math.ceil(len(boxes) / 80): |
| continue |
| label = f"T{index}" |
| left, top = rect[0] + 3, rect[1] + 3 |
| |
| |
| for candidate_top in range(top, max(top + 1, rect[3] - 21), 24): |
| extent = draw.textbbox((left, candidate_top), label, font=font) |
| if not any( |
| left < old[2] + 3 |
| and extent[2] + 3 > old[0] |
| and candidate_top < old[3] + 3 |
| and extent[3] + 3 > old[1] |
| for old in labels |
| ): |
| top = candidate_top |
| break |
| extent = draw.textbbox((left, top), label, font=font) |
| labels.append((left, top, extent[2], extent[3])) |
| draw.rectangle( |
| (left - 2, top - 2, extent[2] + 4, extent[3] + 4), fill=(11, 35, 59) |
| ) |
| draw.text((left, top), label, fill="white", font=font) |
| return image |
|
|
|
|
| def tile_overlay(path, boxes, weights, reference_count=None): |
| image = preview_image(path) |
| with Image.open(path) as original: |
| sx, sy = image.width / original.width, image.height / original.height |
| summed = np.zeros((image.height, image.width), dtype=np.float32) |
| coverage = np.zeros_like(summed) |
| for box, weight in zip(boxes, weights, strict=True): |
| x0, y0, x1, y1 = box |
| x0, y0 = int(x0 * sx), int(y0 * sy) |
| x1, y1 = ( |
| min(image.width, math.ceil(x1 * sx)), |
| min(image.height, math.ceil(y1 * sy)), |
| ) |
| summed[y0:y1, x0:x1] += float(weight) |
| coverage[y0:y1, x0:x1] += 1 |
| averaged = summed / np.maximum(coverage, 1) |
| count = reference_count or len(weights) |
| ratio = np.maximum(averaged * count, 1e-12) |
| heat = colors(attention_levels(averaged, count)).astype(np.float32) |
| alpha = np.clip(0.15 + np.abs(np.log2(ratio)) * 0.3, 0.15, 0.75)[..., None] |
| pixels = np.asarray(image, dtype=np.float32) * (1 - alpha) + heat * alpha |
| overlay = Image.fromarray(np.clip(pixels, 0, 255).astype(np.uint8)) |
| |
| if float(np.ptp(weights)) * count >= 0.10: |
| draw = ImageDraw.Draw(overlay) |
| for rank, index in enumerate( |
| np.argsort(-np.asarray(weights), kind="stable")[:3], 1 |
| ): |
| x0, y0, x1, y1 = boxes[index] |
| rect = ( |
| int(x0 * sx), |
| int(y0 * sy), |
| min(image.width - 1, int(x1 * sx) - 1), |
| min(image.height - 1, int(y1 * sy) - 1), |
| ) |
| draw.rectangle(rect, outline=(12, 26, 42), width=7) |
| draw.rectangle(rect, outline=(255, 198, 65), width=4) |
| return overlay |
|
|
|
|
| def patch_view(state, tile_number): |
| if not state: |
| return None, None, "Run analysis first." |
| index = int(tile_number) - 1 |
| if index < 0 or index >= len(state["boxes"]): |
| raise ValueError("Tile number is out of range.") |
| box = state["boxes"][index] |
| with Image.open(state["path"]) as source: |
| tile = source.convert("RGB").crop(box) |
| weights = state["attention"]["patch"][index] |
| heat = Image.fromarray(colors(attention_levels(weights, weights.size))).resize( |
| tile.size, Image.Resampling.NEAREST |
| ) |
| overlay = Image.blend(tile, heat, 0.65) |
| description = ( |
| f"T{index + 1} · original box {list(box)} · patch grid " |
| f"{weights.shape[0]} × {weights.shape[1]} · " |
| f"patch weights sum = {float(weights.sum()):.6f}. " |
| f" Uniform attention = {100/weights.size:.4f}%. " |
| "All tiles share a scale from 0.25× to 4× uniform attention. Small images retain the black padding used by the model." |
| ) |
| return tile, overlay, description |
|
|
|
|
| def top_tiles(state, k=4): |
| order = np.argsort(-state["attention"]["l1"], kind="stable")[:k] |
| gallery = [] |
| for index in order: |
| tile, _, _ = patch_view(state, int(index) + 1) |
| weight = float(state["attention"]["l1"][index]) |
| gallery.append( |
| (tile, f"#{len(gallery)+1} · T{index + 1} · weight {weight*100:.2f}%") |
| ) |
| return gallery |
|
|
|
|
| def head_view(state, selection): |
| if not state: |
| return None, "Run analysis first." |
| if selection == "mean": |
| weights = state["attention"]["l1"] |
| label = "Mean across L1 aggregation heads" |
| else: |
| index = int(selection.split("_")[-1]) - 1 |
| weights = state["attention"]["l1_heads"][index] |
| label = f"L1 aggregation head {index + 1}" |
| return tile_overlay(state["path"], state["boxes"], weights), ( |
| f"{label} · uniform reference = {100/len(weights):.2f}%. " |
| "Blue is below the reference, gray is near it, and orange-red is above it. Overlapping regions use mean tile attention." |
| "Tile attention is not a pixel-level forgery probability." |
| ) |
|
|
|
|
| def parameter_rows(detector, branch, threshold, views): |
| core, visual = detector.model.core, detector.model.visual |
| mil1, mil2 = core.mil_level1.head.head, core.mil_level2.head.head |
| rows = [ |
| ["Inference", "Branch / threshold", f"{branch} / {threshold:.2f}"], |
| [ |
| "Inference", |
| "Precision / device", |
| f"{visual.conv1.weight.dtype} / {detector.device}", |
| ], |
| ["Inference", "CLIP chunk size", str(detector.model.eval_chunk_size)], |
| ["Inference", "Views: global + local", f"1 + {views - 1} = {views}"], |
| ["Inference", "Tiling policy", "Non-overlapping full blocks; overlap only at right/bottom boundaries"], |
| ["Architecture", "Backbone", "Frozen CLIP ViT-L/14"], |
| [ |
| "Architecture", |
| "Local tile size", |
| f"{int(detector.model.input_resolution)} × {int(detector.model.input_resolution)}", |
| ], |
| ["Architecture", "CLIP patch size", str(visual.conv1.kernel_size[0])], |
| ["Architecture", "Feature dimension", str(detector.model.embed_dim)], |
| [ |
| "Architecture", |
| "Feature extraction layer", |
| str( |
| detector.model.clip_mil_stop_after_block_1based |
| if detector.model.use_clip_intermediate_for_mil |
| else len(visual.transformer.resblocks) |
| ), |
| ], |
| [ |
| "Architecture", |
| "L1 across-tile aggregation heads", |
| str(core.level1_aggregate_attn.num_heads), |
| ], |
| [ |
| "Architecture", |
| "L1 MIL pooling / hidden", |
| f"{mil1.pooling} / {mil1.attn_V.out_features} ({mil1.attn_bag_mode})", |
| ], |
| [ |
| "Architecture", |
| "L2 MIL pooling / hidden", |
| f"{mil2.pooling} / {mil2.attn_V.out_features} ({mil2.attn_bag_mode})", |
| ], |
| [ |
| "Architecture", |
| "Fusion / gate hidden", |
| f"pool-logit gated / {core.pool_gate[0].out_features}", |
| ], |
| ["Architecture", "Fusion MLP hidden", str(core.fusion_head[0].out_features)], |
| ["Inference", "Cross-shuffle / dropout", "Disabled in evaluation mode"], |
| ] |
| reference = Path(__file__).with_name("training_reference.json") |
| if reference.is_file(): |
| training = json.loads(reference.read_text()) |
| rows.extend( |
| ["Training only", key, json.dumps(value, ensure_ascii=False)] |
| for key, value in training.items() |
| ) |
| return rows |
|
|
|
|
| def public_attention(state): |
| weights = state["attention"] |
| return dict( |
| meaning="Relative model attention, not a forgery probability or causal explanation.", |
| display="Fixed log2 scale relative to uniform 1/N attention: 0.25x to 4x; overlap uses mean tile weight.", |
| local_tile_boxes=[list(box) for box in state["boxes"]], |
| l1_tile_weights=weights["l1"].tolist(), |
| l1_aggregation_head_weights=weights["l1_heads"].tolist(), |
| l2_view_weights=weights["l2"].tolist(), |
| l2_view_order=["global"] + [f"T{i + 1}" for i in range(len(state["boxes"]))], |
| patches_per_tile=list(weights["patch"].shape[1:]), |
| fusion_embedding_gate_l1=weights["gate"], |
| fusion_embedding_gate_l2=1 - weights["gate"], |
| ) |
|
|
