| """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}') |
| |
| 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) |
|
|
