File size: 6,147 Bytes
a2a2e29 acb2ce3 5cd0eb1 acb2ce3 a2a2e29 acb2ce3 a2a2e29 acb2ce3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Offline model loading and non-overlapping tiling with boundary completion."""
from __future__ import annotations
import contextlib
import sys
from pathlib import Path
from PIL import Image
import torch
import yaml
from model.LTF_v3_3 import LTFNetV3_3
from utils.tile_dataset_srm import clip_transform, crop_image_to_tiles
ROOT = Path(__file__).absolute().parents[1]
BRANCHES = ('level1', 'level2', 'fusion')
def load_config(path=None):
path = Path(path) if path else ROOT / 'configs/eval.yaml'
with path.expanduser().resolve().open(encoding='utf-8') as stream:
config = yaml.safe_load(stream)
if not isinstance(config, dict):
raise ValueError('Configuration must be a YAML mapping')
return config
def package_path(path):
path = Path(path).expanduser()
return path.resolve() if path.is_absolute() else (ROOT / path).resolve()
def model_kwargs(config):
return dict(
num_heads=config.get('transformer_heads', 2), num_classes=2,
num_layers=config.get('transformer_layers', 0), dropout=config.get('dropout', .1),
tile_heads=config.get('tile_heads', 2), tile_layers=config.get('tile_layers', 1),
use_global_image=config.get('use_global_image', True),
use_patch_transformer=config.get('use_patch_transformer', False),
linear_classifier=config.get('linear_classifier', True),
mil_heads=config.get('mil_heads', 8), fusion_hidden=config.get('fusion_hidden', 512),
l1_cross_shuffle_p=float(config.get('v2_l1_cross_shuffle_p', .9)),
l1_replace_ratio_range=tuple(config.get('v2_l1_replace_ratio_range', [.5, .9])),
l1_with_replacement_p=float(config.get('v2_l1_with_replacement_p', .5)),
l2_cross_shuffle_p=float(config.get('v2_l2_cross_shuffle_p', .9)),
l2_replace_ratio_range=tuple(config.get('v2_l2_replace_ratio_range', [.1, .8])),
l2_with_replacement_p=float(config.get('v2_l2_with_replacement_p', .5)),
bag_token_shuffle_p=float(config.get('v2_bag_token_shuffle_p', .5)),
cross_shuffle_min_keep_original=int(config.get('v2_cross_shuffle_min_keep_original', 1)),
use_clip_intermediate_for_mil=config.get('v2_use_clip_intermediate_for_mil', False),
clip_mil_stop_after_block_1based=int(config.get('v2_clip_mil_stop_after_block_1based', 12)),
pool_gate_hidden=int(config.get('v3_3_pool_gate_hidden', 16)),
use_pool_gate_fusion=config.get('v3_3_use_pool_gate_fusion', True),
use_mil_cross_shuffle=config.get('mil_use_cross_shuffle', True),
)
def read_core_weights(path):
weights = torch.load(path, map_location='cpu', weights_only=True)
if not isinstance(weights, dict) or not weights or not all(torch.is_tensor(v) for v in weights.values()):
raise ValueError('Expected an LTF v3.3 core state_dict, such as epoch_20.pth')
return weights
def create_model(config, weights=None, device='cpu'):
import clip
backbone = package_path(config['backbone'])
if not backbone.is_file():
raise FileNotFoundError(f'Bundled CLIP backbone not found: {backbone}')
# Local checkpoint loading is explicit: this path never downloads a model.
with contextlib.redirect_stdout(sys.stderr):
clip_model, _ = clip.load(str(backbone), device='cpu', jit=False)
model = LTFNetV3_3(clip_model.visual, **model_kwargs(config))
del clip_model
if weights is not None:
model.core.load_state_dict(read_core_weights(package_path(weights)), strict=True)
return model.to(device)
def prepare_image(path, tile_size=224):
"""Original RGB and CLIP normalization; non-overlapping blocks with boundary completion."""
path = Path(path).expanduser().resolve()
with Image.open(path) as source:
image = source.convert('RGB')
transform = clip_transform(tile_size, blur_radius=None)
local_tiles = crop_image_to_tiles(image, tile_size)
return torch.stack([transform(image)] + [transform(tile) for tile in local_tiles])
class LTFDetector:
def __init__(self, config=None, weights=None, device='auto', chunk_size=256):
self.config = dict(config) if isinstance(config, dict) else load_config(config)
if device == 'auto':
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
self.device = torch.device(device)
if self.device.type == 'cuda' and not torch.cuda.is_available():
raise RuntimeError('CUDA is unavailable; use device="cpu" for CPU inference')
if int(chunk_size) < 1:
raise ValueError('chunk_size must be positive')
torch.manual_seed(int(self.config.get('seed', 42)))
self.weights = package_path(weights or self.config.get('weights', 'weights/epoch_20.pth'))
self.model = create_model(self.config, self.weights, self.device).eval()
self.model.eval_chunk_size = int(chunk_size)
@torch.inference_mode()
def predict(self, path, branch='fusion', threshold=.5):
if branch not in BRANCHES:
raise ValueError(f'branch must be one of {BRANCHES}')
if not 0 <= threshold <= 1:
raise ValueError('threshold must be between zero and one')
path = Path(path).expanduser().resolve()
images = prepare_image(path, int(self.model.input_resolution))
self.model.eval()
output = self.model(images.unsqueeze(0).to(self.device))
probabilities = {}
for name in BRANCHES:
probability = output[f'logits_{name}'].float().softmax(-1)[0].cpu()
if not torch.isfinite(probability).all():
raise ValueError(f'Nonfinite {name} prediction for {path}')
probabilities[name] = {'real': float(probability[0]), 'fake': float(probability[1])}
score = probabilities[branch]['fake']
label = int(score >= threshold)
return dict(path=str(path), prediction='fake' if label else 'real', label=label,
branch=branch, threshold=float(threshold), score_fake=score,
tiling_policy="nonoverlap_then_boundary",
views=int(images.shape[0]), probabilities=probabilities)
|