File size: 5,005 Bytes
03b9592 | 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 | import cv2
import torch
import torch.nn.functional as F
import numpy as np
from transformers import PreTrainedModel, PretrainedConfig
# ==========================================
# 1. Hugging Face Configuration
# ==========================================
class PhorensicsConfig(PretrainedConfig):
model_type = "phorensics"
def __init__(self, max_dim=600, patch_size=128, stride=64, **kwargs):
self.max_dim = max_dim
self.patch_size = patch_size
self.stride = stride
super().__init__(**kwargs)
# ==========================================
# 2. Native PyTorch Physics Engine
# ==========================================
class PhorensicsModel(PreTrainedModel):
config_class = PhorensicsConfig
def __init__(self, config):
super().__init__(config)
self.patch_size = config.patch_size
self.stride = config.stride
# THE FIX: Hugging Face requires at least one parameter to determine model dtype.
# This dummy parameter satisfies the push_to_hub internal checks.
self.dummy_param = torch.nn.Parameter(torch.empty(0, dtype=torch.float32))
# STRUCTURAL TENSORS: Stored permanently in the .safetensors file.
laplacian_kernel = torch.tensor([[[[0., 1., 0.], [1., -4., 1.], [0., 1., 0.]]]], dtype=torch.float32)
self.register_buffer("noise_kernel", laplacian_kernel)
self.register_buffer("safe_z_threshold", torch.tensor(2.5, dtype=torch.float32))
def _extract_patches(self, tensor):
"""Hardware-accelerated tensor unfolding (PyTorch native sliding window)"""
c, h, w = tensor.shape
pad_h = max(0, self.patch_size - h) if h < self.patch_size else 0
pad_w = max(0, self.patch_size - w) if w < self.patch_size else 0
if pad_h > 0 or pad_w > 0:
tensor = F.pad(tensor, (0, pad_w, 0, pad_h), mode='reflect')
patches = tensor.unfold(1, self.patch_size, self.stride).unfold(2, self.patch_size, self.stride)
return patches.contiguous().view(c, -1, self.patch_size, self.patch_size).permute(1, 0, 2, 3)
@torch.no_grad()
def forward(self, image_path):
"""
The native inference pass. Automatically handles CPU/CUDA routing.
"""
# 1. Load and Standardize Image
img_bgr = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
if img_bgr is None:
raise ValueError(f"Failed to load image: {image_path}")
h, w = img_bgr.shape[:2]
if max(h, w) > self.config.max_dim:
scale = self.config.max_dim / max(h, w)
img_bgr = cv2.resize(img_bgr, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
# 2. TENSOR CONVERSION & GPU ACCELERATION
gray_np = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
# Move the image tensor to whatever device the model is currently on (CPU/CUDA)
t_gray = torch.from_numpy(gray_np).to(self.device).float().unsqueeze(0).unsqueeze(0)
# 3. SPATIAL TENSOR CONVOLUTION (PRNU Noise Validation)
# Apply the .safetensors kernel across the image
noise_map = F.conv2d(t_gray, self.noise_kernel, padding=1).squeeze(0)
p_noise = self._extract_patches(noise_map)
# Calculate robust Z-Score using pure vectorized PyTorch math
patch_variances = torch.var(p_noise.view(p_noise.shape[0], -1), dim=1)
median_var = torch.median(patch_variances)
mad = torch.median(torch.abs(patch_variances - median_var)) + 1e-3
z_scores = 0.6745 * torch.abs(patch_variances - median_var) / mad
max_noise_z = torch.quantile(z_scores, 0.98).item()
# 4. FAST FOURIER TRANSFORM (Global Spectral Decay)
f = np.fft.fft2(gray_np)
power = np.abs(np.fft.fftshift(f))**2
h_fft, w_fft = power.shape
cy, cx = h_fft // 2, w_fft // 2
y, x = np.indices((h_fft, w_fft))
r = np.sqrt((x - cx)**2 + (y - cy)**2).astype(int)
radial_profile = np.bincount(r.ravel(), power.ravel()) / np.maximum(np.bincount(r.ravel()), 1)
r_vals = np.arange(1, len(radial_profile))
p_vals = radial_profile[1:]
slope, _ = np.polyfit(np.log10(r_vals), np.log10(p_vals + 1e-10), 1)
alpha = (-slope) * 0.70
# 5. ENGINE FUSION
threshold = self.safe_z_threshold.item()
noise_threat = max(0.0, min(((max_noise_z - threshold) / 1.5) * 100.0, 100.0))
fft_threat = 99.9 if (2.0 <= alpha <= 3.5) else max(((alpha - 1.0) / 1.0) * 49.9, 0.0)
final_prob = (noise_threat * 0.40) + (fft_threat * 0.60)
return {
"is_fake": final_prob >= 50.0,
"threat_probability": round(min(final_prob, 99.9), 2),
"max_noise_z": round(max_noise_z, 2),
"fft_alpha_value": round(alpha, 3)
} |