import logging import os from typing import List import torch import torchvision.transforms as transforms import torchvision.transforms.functional as F logger = logging.getLogger(__name__) MEAN = { "imagenet": [0.485, 0.456, 0.406], "clip": [0.48145466, 0.4578275, 0.40821073], "dino": [0.485, 0.456, 0.406], } STD = { "imagenet": [0.229, 0.224, 0.225], "clip": [0.26862954, 0.26130258, 0.27577711], "dino": [0.229, 0.224, 0.225], } RESIDUAL_MEAN: List[float] = [0.0, 0.0, 0.0] RESIDUAL_STD: List[float] = [0.5, 0.5, 0.5] # --------------------------------------------------------------------------- # Filesystem helpers # --------------------------------------------------------------------------- def recursively_read( rootdir: str, must_contain: str, exts=("png", "jpg", "jpeg", "JPEG", "bmp", "webp"), ) -> List[str]: """Recursively scan ``rootdir`` for image files. Args: rootdir: Root of the search. must_contain: Optional substring filter applied to the full path. exts: Allowed file extensions (case-insensitive on the lowercase form). Returns: List of absolute / joined paths to the matching files. """ out: List[str] = [] exts_lc = {e.lower() for e in exts} for r, _, files in os.walk(rootdir, followlinks=True): for file in files: # Defensive split: files without an extension are skipped instead # of raising IndexError. parts = file.split(".") if len(parts) < 2: continue ext = parts[-1].lower() full = os.path.join(r, file) if ext in exts_lc and must_contain in full: out.append(full) return out def get_list(path: str, must_contain: str = "") -> List[str]: """Convenience wrapper around :func:`recursively_read` with no extras.""" return recursively_read(path, must_contain) # --------------------------------------------------------------------------- # Geometry transforms # --------------------------------------------------------------------------- class PadRandomCrop: """Pad the image to ``size`` if needed, then take a random crop. Used at training time to preserve original pixel patterns (no resize distortion that would erase the per-pixel quantization residual). """ def __init__(self, size: int): self.size = size def __call__(self, img): w, h = img.size pad_h = max(0, self.size - h) pad_w = max(0, self.size - w) if pad_h > 0 or pad_w > 0: padding = ( pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2, ) img = F.pad(img, padding, fill=0) return transforms.RandomCrop(self.size)(img) class PadCenterCrop: """Pad the image to ``size`` if needed, then take a deterministic centre crop. Mirrors :class:`PadRandomCrop` exactly (same padding strategy) so train and test pipelines differ only in the crop position. Critical for forensic/PRNU-style features that can be wiped by ``Resize``. """ def __init__(self, size: int): self.size = size def __call__(self, img): w, h = img.size pad_h = max(0, self.size - h) pad_w = max(0, self.size - w) if pad_h > 0 or pad_w > 0: padding = ( pad_w // 2, pad_h // 2, pad_w - pad_w // 2, pad_h - pad_h // 2, ) img = F.pad(img, padding, fill=0) return transforms.CenterCrop(self.size)(img) # --------------------------------------------------------------------------- # Residual transform # --------------------------------------------------------------------------- class AppendResidual: """Append the quantization residual as 3 extra channels. The residual is computed on the RGB tensor in ``[0, 1]`` (BEFORE normalization). Normalization is then applied to the resulting 6-channel tensor, with different mean/std for the RGB block and the residual block. Output: ``[6, H, W]`` where channels ``0..2`` are RGB and channels ``3..5`` are the residual. The residual extractor used here is exactly the one used inside the model (``models.encoder.residual_extractor``). We import it lazily in ``__init__`` to keep this module importable when the model package is not yet on the import path (e.g. during unit tests). """ def __init__(self): # Lazy import so this module can still be imported when the # ``models`` package is absent (used by light-weight smoke tests). from models.encoder.residual_extractor import get_residual_extractor self.residual_extractor = get_residual_extractor() def __call__(self, img_tensor: torch.Tensor) -> torch.Tensor: """Compute the residual and concatenate with the RGB tensor. Args: img_tensor: ``[3, H, W]`` RGB tensor in ``[0, 1]``. Returns: ``[6, H, W]`` tensor: ``[RGB ∥ residual]``. """ assert img_tensor.size(0) == 3, ( f"Expected 3 channels, got {img_tensor.size(0)}" ) assert img_tensor.min() >= 0.0 and img_tensor.max() <= 1.0, ( f"Input must be in [0, 1], got " f"range [{img_tensor.min():.3f}, {img_tensor.max():.3f}]" ) residual = self.residual_extractor(img_tensor.unsqueeze(0)).squeeze(0) residual = residual.to(dtype=img_tensor.dtype) return torch.cat([img_tensor, residual], dim=0) def __repr__(self) -> str: return self.__class__.__name__ + "(input=[0,1])" # --------------------------------------------------------------------------- # Compose helpers # --------------------------------------------------------------------------- def create_train_transforms( image_size: int = 224, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), is_crop: bool = True, ) -> transforms.Compose: if is_crop: resize_func = PadRandomCrop(image_size) else: logger.info("Using Resize to %d x %d", image_size, image_size) resize_func = transforms.Resize((image_size, image_size)) return transforms.Compose( [ transforms.RandomHorizontalFlip(), transforms.RandomApply( [ transforms.ColorJitter( brightness=0.15, contrast=0.15, saturation=0.15, hue=0.05, ) ], p=0.2, ), resize_func, transforms.ToTensor(), AppendResidual(), transforms.Normalize( mean=list(mean) + RESIDUAL_MEAN, std=list(std) + RESIDUAL_STD, ), ] ) def create_eval_transforms( image_size: int = 224, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), ) -> transforms.Compose: logger.info( "[Eval] Using deterministic PadCenterCrop to %d x %d", image_size, image_size, ) return transforms.Compose( [ PadCenterCrop(image_size), transforms.ToTensor(), AppendResidual(), transforms.Normalize( mean=list(mean) + RESIDUAL_MEAN, std=list(std) + RESIDUAL_STD, ), ] )