| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
| from typing import Iterable, Protocol, Sequence |
|
|
| from PIL import Image, ImageOps |
|
|
| from .config import ( |
| DEFAULT_BACKEND, |
| DEFAULT_HF_MODEL_ID, |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| DEFAULT_MODEL_NAME, |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| DEFAULT_PRETRAINED, |
| FC_WEIGHT_PATH_IN_REPO, |
| IMAGE_EXTENSIONS, |
| MODEL_REPO_ID, |
| ) |
| from .nonescape_adapter import NonescapeFullDetector, NonescapeMiniDetector |
| from .sentry_adapter import SentryConvNeXtDetector |
| from .types import DetectionResult |
|
|
|
|
| class Detector(Protocol): |
| threshold: float |
|
|
| def predict_image(self, image: Image.Image, path: Path | None = None) -> DetectionResult: |
| ... |
|
|
| def predict_images( |
| self, |
| images: Sequence[Image.Image], |
| paths: Sequence[Path | None] | None = None, |
| ) -> list[DetectionResult]: |
| ... |
|
|
| def predict_path(self, path: str | Path) -> DetectionResult: |
| ... |
|
|
| def model_info(self) -> dict: |
| ... |
|
|
|
|
| def _select_device(torch_module, requested: str | None) -> str: |
| if requested and requested != "auto": |
| return requested |
| if torch_module.cuda.is_available(): |
| return "cuda" |
| mps = getattr(torch_module.backends, "mps", None) |
| if mps is not None and mps.is_available(): |
| return "mps" |
| return "cpu" |
|
|
|
|
| def _normalise_image(image: Image.Image) -> Image.Image: |
| return ImageOps.exif_transpose(image).convert("RGB") |
|
|
|
|
| def _load_torch_checkpoint(torch_module, weight_path: str | Path): |
| try: |
| return torch_module.load(str(weight_path), map_location="cpu", weights_only=True) |
| except TypeError: |
| return torch_module.load(str(weight_path), map_location="cpu") |
|
|
|
|
| def _build_detection_result( |
| *, |
| path: Path | None, |
| probability_ai: float, |
| raw_score: float, |
| threshold: float, |
| backend: str, |
| ) -> DetectionResult: |
| probability_real = 1.0 - probability_ai |
| label = "ai" if probability_ai >= threshold else "real" |
| confidence = probability_ai if label == "ai" else probability_real |
| return DetectionResult( |
| path=path, |
| label=label, |
| probability_ai=probability_ai, |
| probability_real=probability_real, |
| confidence=confidence, |
| raw_score=raw_score, |
| backend=backend, |
| ) |
|
|
|
|
| class AIImageDetector: |
| """UnivFD-style detector: CLIP image encoder plus a linear fake/real head.""" |
|
|
| backend_name = "univfd" |
|
|
| def __init__( |
| self, |
| device: str | None = None, |
| threshold: float = 0.5, |
| model_name: str = DEFAULT_MODEL_NAME, |
| pretrained: str = DEFAULT_PRETRAINED, |
| weight_path: str | Path | None = None, |
| repo_id: str = MODEL_REPO_ID, |
| ) -> None: |
| try: |
| import torch |
| except ImportError as exc: |
| raise RuntimeError("Missing dependency: install torch to use the UnivFD backend.") from exc |
| try: |
| import open_clip |
| except ImportError as exc: |
| raise RuntimeError("Missing dependency: install open_clip_torch to use the UnivFD backend.") from exc |
|
|
| self._torch = torch |
| self.device = torch.device(_select_device(torch, device)) |
| self.threshold = threshold |
| self.model_name = model_name |
| self.pretrained = pretrained |
| self.repo_id = repo_id |
|
|
| self.clip_model, _, self.preprocess = open_clip.create_model_and_transforms( |
| model_name, |
| pretrained=pretrained, |
| device=self.device, |
| ) |
| self.clip_model.eval() |
|
|
| output_dim = getattr(getattr(self.clip_model, "visual", None), "output_dim", 768) |
| self.head = torch.nn.Linear(int(output_dim), 1).to(self.device) |
| self._load_head(weight_path) |
| self.head.eval() |
|
|
| def _load_head(self, weight_path: str | Path | None) -> None: |
| if weight_path is None: |
| try: |
| from huggingface_hub import hf_hub_download |
| except ImportError as exc: |
| raise RuntimeError( |
| "Missing dependency: install huggingface_hub or pass --weight-path." |
| ) from exc |
| weight_path = hf_hub_download( |
| repo_id=self.repo_id, |
| filename=FC_WEIGHT_PATH_IN_REPO, |
| ) |
|
|
| state = _load_torch_checkpoint(self._torch, weight_path) |
| if isinstance(state, dict) and "state_dict" in state: |
| state = state["state_dict"] |
|
|
| if isinstance(state, dict): |
| cleaned = {} |
| for key, value in state.items(): |
| key = key.replace("module.", "") |
| key = key.replace("model.", "") |
| key = key.replace("fc.", "") |
| key = key.replace("head.", "") |
| if key in {"weight", "bias"}: |
| cleaned[key] = value |
| if cleaned: |
| self.head.load_state_dict(cleaned, strict=True) |
| return |
| try: |
| self.head.load_state_dict(state, strict=False) |
| return |
| except Exception as exc: |
| raise RuntimeError(f"Could not load UnivFD head weights from {weight_path}: {exc}") from exc |
| raise RuntimeError(f"Unsupported checkpoint format: {type(state)!r}") |
|
|
| def model_info(self) -> dict: |
| return { |
| "backend": self.backend_name, |
| "model_name": self.model_name, |
| "pretrained": self.pretrained, |
| "head_repo": self.repo_id, |
| "head_file": FC_WEIGHT_PATH_IN_REPO, |
| "device": str(self.device), |
| "threshold": self.threshold, |
| } |
|
|
| @property |
| def _torch_inference_mode(self): |
| return self._torch.inference_mode |
|
|
| def predict_image(self, image: Image.Image, path: Path | None = None) -> DetectionResult: |
| return self.predict_images([image], paths=[path])[0] |
|
|
| def predict_images( |
| self, |
| images: Sequence[Image.Image], |
| paths: Sequence[Path | None] | None = None, |
| ) -> list[DetectionResult]: |
| if not images: |
| return [] |
| if paths is None: |
| paths = [None] * len(images) |
| if len(paths) != len(images): |
| raise ValueError("paths must have the same length as images") |
|
|
| torch = self._torch |
| tensors = [self.preprocess(_normalise_image(image)) for image in images] |
| batch = torch.stack(tensors, dim=0).to(self.device) |
| with torch.inference_mode(): |
| features = self.clip_model.encode_image(batch) |
| features = features / features.norm(dim=-1, keepdim=True) |
| logits = self.head(features).flatten().float().detach().cpu() |
| probabilities = torch.sigmoid(logits) |
|
|
| results: list[DetectionResult] = [] |
| for probability_ai_tensor, logit_tensor, path in zip(probabilities, logits, paths, strict=True): |
| results.append( |
| _build_detection_result( |
| path=path, |
| probability_ai=float(probability_ai_tensor.item()), |
| raw_score=float(logit_tensor.item()), |
| threshold=self.threshold, |
| backend=self.backend_name, |
| ) |
| ) |
| return results |
|
|
| def predict_path(self, path: str | Path) -> DetectionResult: |
| path = Path(path) |
| with Image.open(path) as image: |
| return self.predict_image(image, path=path) |
|
|
| def predict_many( |
| self, |
| paths: Iterable[str | Path], |
| batch_size: int = 16, |
| ) -> list[DetectionResult]: |
| results: list[DetectionResult] = [] |
| batch_images: list[Image.Image] = [] |
| batch_paths: list[Path] = [] |
| for path_like in paths: |
| path = Path(path_like) |
| with Image.open(path) as image: |
| batch_images.append(_normalise_image(image)) |
| batch_paths.append(path) |
| if len(batch_images) >= batch_size: |
| results.extend(self.predict_images(batch_images, paths=batch_paths)) |
| batch_images = [] |
| batch_paths = [] |
| if batch_images: |
| results.extend(self.predict_images(batch_images, paths=batch_paths)) |
| return results |
|
|
|
|
| class HuggingFaceImageDetector: |
| """Generic Hugging Face image-classification backend.""" |
|
|
| backend_name = "hf" |
|
|
| def __init__( |
| self, |
| model_id: str = DEFAULT_HF_MODEL_ID, |
| device: str | None = None, |
| threshold: float = 0.5, |
| ) -> None: |
| try: |
| import torch |
| except ImportError as exc: |
| raise RuntimeError("Missing dependency: install torch to use the Hugging Face backend.") from exc |
| try: |
| from transformers import AutoImageProcessor, AutoModelForImageClassification |
| except ImportError as exc: |
| raise RuntimeError("Missing dependency: install transformers to use the HF backend.") from exc |
|
|
| self._torch = torch |
| self.device = torch.device(_select_device(torch, device)) |
| self.threshold = threshold |
| self.model_id = model_id |
| model_source = _resolve_local_hf_snapshot(model_id) or model_id |
| local_only = model_source != model_id |
| try: |
| self.model = AutoModelForImageClassification.from_pretrained( |
| model_source, |
| local_files_only=local_only, |
| ).to(self.device) |
| except Exception: |
| try: |
| self.model = AutoModelForImageClassification.from_pretrained(model_id).to(self.device) |
| except Exception as second_exc: |
| raise RuntimeError( |
| "Could not load the requested Hugging Face model as a standard " |
| "Transformers image-classification checkpoint. Some open-source " |
| "detectors, such as custom Sentry or Nonescape releases, need a " |
| "dedicated adapter instead of the generic --backend hf path. " |
| f"Model: {model_id}. Last error: {second_exc}" |
| ) from second_exc |
| self.model.eval() |
| self.id2label = self.model.config.id2label |
| try: |
| self.processor = AutoImageProcessor.from_pretrained( |
| model_source, |
| local_files_only=local_only, |
| ) |
| except Exception: |
| self.processor = _build_fallback_image_processor(self.model.config) |
|
|
| def model_info(self) -> dict: |
| return { |
| "backend": self.backend_name, |
| "model_id": self.model_id, |
| "device": str(self.device), |
| "threshold": self.threshold, |
| "id2label": {str(k): v for k, v in self.id2label.items()}, |
| } |
|
|
| def predict_image(self, image: Image.Image, path: Path | None = None) -> DetectionResult: |
| return self.predict_images([image], paths=[path])[0] |
|
|
| def predict_images( |
| self, |
| images: Sequence[Image.Image], |
| paths: Sequence[Path | None] | None = None, |
| ) -> list[DetectionResult]: |
| if not images: |
| return [] |
| if paths is None: |
| paths = [None] * len(images) |
| if len(paths) != len(images): |
| raise ValueError("paths must have the same length as images") |
|
|
| torch = self._torch |
| normalised = [_normalise_image(image) for image in images] |
| inputs = self.processor(images=normalised, return_tensors="pt") |
| inputs = {key: value.to(self.device) for key, value in inputs.items()} |
| with torch.inference_mode(): |
| logits = self.model(**inputs).logits.detach().cpu() |
| probabilities = torch.softmax(logits, dim=-1) |
|
|
| fake_indices = self._fake_label_indices() |
| if not fake_indices: |
| raise RuntimeError(f"Could not infer fake/AI label from model labels: {self.id2label}") |
|
|
| results: list[DetectionResult] = [] |
| for row, logit_row, path in zip(probabilities, logits, paths, strict=True): |
| results.append( |
| _build_detection_result( |
| path=path, |
| probability_ai=float(row[fake_indices].sum().item()), |
| raw_score=float(logit_row[fake_indices].mean().item()), |
| threshold=self.threshold, |
| backend=self.backend_name, |
| ) |
| ) |
| return results |
|
|
| def predict_path(self, path: str | Path) -> DetectionResult: |
| path = Path(path) |
| with Image.open(path) as image: |
| return self.predict_image(image, path=path) |
|
|
| def predict_many( |
| self, |
| paths: Iterable[str | Path], |
| batch_size: int = 16, |
| ) -> list[DetectionResult]: |
| results: list[DetectionResult] = [] |
| batch_images: list[Image.Image] = [] |
| batch_paths: list[Path] = [] |
| for path_like in paths: |
| path = Path(path_like) |
| with Image.open(path) as image: |
| batch_images.append(_normalise_image(image)) |
| batch_paths.append(path) |
| if len(batch_images) >= batch_size: |
| results.extend(self.predict_images(batch_images, paths=batch_paths)) |
| batch_images = [] |
| batch_paths = [] |
| if batch_images: |
| results.extend(self.predict_images(batch_images, paths=batch_paths)) |
| return results |
|
|
| def _fake_label_indices(self) -> list[int]: |
| fake_tokens = ("fake", "ai", "aigc", "synthetic", "generated") |
| real_tokens = ("real", "human", "authentic", "natural") |
| fake_indices: list[int] = [] |
| for index, label in self.id2label.items(): |
| label_l = str(label).lower() |
| is_fake = any(token in label_l for token in fake_tokens) |
| is_real = any(token in label_l for token in real_tokens) |
| if is_fake and not is_real: |
| fake_indices.append(int(index)) |
| return fake_indices |
|
|
|
|
| class HybridImageDetector: |
| """Blend UnivFD and HF image classifier scores for a stronger practical baseline.""" |
|
|
| backend_name = "hybrid" |
|
|
| def __init__( |
| self, |
| *, |
| device: str | None = None, |
| threshold: float = 0.5, |
| univfd_weight: float = DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| weight_path: str | Path | None = None, |
| model_name: str = DEFAULT_MODEL_NAME, |
| pretrained: str = DEFAULT_PRETRAINED, |
| hf_model: str = DEFAULT_HF_MODEL_ID, |
| ) -> None: |
| if not 0.0 <= univfd_weight <= 1.0: |
| raise ValueError("univfd_weight must be between 0 and 1") |
| self.threshold = threshold |
| self.univfd_weight = univfd_weight |
| self.univfd = AIImageDetector( |
| device=device, |
| threshold=threshold, |
| model_name=model_name, |
| pretrained=pretrained, |
| weight_path=weight_path, |
| ) |
| self.hf = HuggingFaceImageDetector( |
| model_id=hf_model, |
| device=device, |
| threshold=threshold, |
| ) |
|
|
| def model_info(self) -> dict: |
| return { |
| "backend": self.backend_name, |
| "threshold": self.threshold, |
| "univfd_weight": self.univfd_weight, |
| "hf_weight": round(1.0 - self.univfd_weight, 6), |
| "components": { |
| "univfd": self.univfd.model_info(), |
| "hf": self.hf.model_info(), |
| }, |
| } |
|
|
| def predict_image(self, image: Image.Image, path: Path | None = None) -> DetectionResult: |
| return self.predict_images([image], paths=[path])[0] |
|
|
| def predict_images( |
| self, |
| images: Sequence[Image.Image], |
| paths: Sequence[Path | None] | None = None, |
| ) -> list[DetectionResult]: |
| if not images: |
| return [] |
| if paths is None: |
| paths = [None] * len(images) |
| if len(paths) != len(images): |
| raise ValueError("paths must have the same length as images") |
|
|
| univfd_results = self.univfd.predict_images(images, paths=paths) |
| hf_results = self.hf.predict_images(images, paths=paths) |
| results: list[DetectionResult] = [] |
| for univfd_result, hf_result, path in zip(univfd_results, hf_results, paths, strict=True): |
| probability_ai = ( |
| self.univfd_weight * univfd_result.probability_ai |
| + (1.0 - self.univfd_weight) * hf_result.probability_ai |
| ) |
| raw_score = ( |
| self.univfd_weight * univfd_result.raw_score |
| + (1.0 - self.univfd_weight) * hf_result.raw_score |
| ) |
| results.append( |
| _build_detection_result( |
| path=path, |
| probability_ai=probability_ai, |
| raw_score=raw_score, |
| threshold=self.threshold, |
| backend=self.backend_name, |
| ) |
| ) |
| return results |
|
|
|
|
| class HybridPlusDetector: |
| """Blend the existing hybrid detector with Nonescape Mini for a stronger practical ensemble.""" |
|
|
| backend_name = "hybrid-plus" |
|
|
| def __init__( |
| self, |
| *, |
| device: str | None = None, |
| threshold: float = 0.5, |
| primary_weight: float = DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| weight_path: str | Path | None = None, |
| model_name: str = DEFAULT_MODEL_NAME, |
| pretrained: str = DEFAULT_PRETRAINED, |
| hf_model: str = DEFAULT_HF_MODEL_ID, |
| hybrid_univfd_weight: float = DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| ) -> None: |
| if not 0.0 <= primary_weight <= 1.0: |
| raise ValueError("primary_weight must be between 0 and 1") |
| self.threshold = threshold |
| self.primary_weight = primary_weight |
| self.primary = HybridImageDetector( |
| device=device, |
| threshold=threshold, |
| univfd_weight=hybrid_univfd_weight, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| ) |
| self.secondary = NonescapeMiniDetector( |
| threshold=threshold, |
| device=device, |
| ) |
|
|
| def model_info(self) -> dict: |
| return { |
| "backend": self.backend_name, |
| "threshold": self.threshold, |
| "primary_weight": self.primary_weight, |
| "secondary_weight": round(1.0 - self.primary_weight, 6), |
| "components": { |
| "primary": self.primary.model_info(), |
| "secondary": self.secondary.model_info(), |
| }, |
| } |
|
|
| def predict_image(self, image: Image.Image, path: Path | None = None) -> DetectionResult: |
| return self.predict_images([image], paths=[path])[0] |
|
|
| def predict_images( |
| self, |
| images: Sequence[Image.Image], |
| paths: Sequence[Path | None] | None = None, |
| ) -> list[DetectionResult]: |
| if not images: |
| return [] |
| if paths is None: |
| paths = [None] * len(images) |
| if len(paths) != len(images): |
| raise ValueError("paths must have the same length as images") |
|
|
| primary_results = self.primary.predict_images(images, paths=paths) |
| secondary_results = self.secondary.predict_images(images, paths=paths) |
| results: list[DetectionResult] = [] |
| for primary_result, secondary_result, path in zip(primary_results, secondary_results, paths, strict=True): |
| probability_ai = ( |
| self.primary_weight * primary_result.probability_ai |
| + (1.0 - self.primary_weight) * secondary_result.probability_ai |
| ) |
| raw_score = ( |
| self.primary_weight * primary_result.raw_score |
| + (1.0 - self.primary_weight) * secondary_result.raw_score |
| ) |
| results.append( |
| _build_detection_result( |
| path=path, |
| probability_ai=probability_ai, |
| raw_score=raw_score, |
| threshold=self.threshold, |
| backend=self.backend_name, |
| ) |
| ) |
| return results |
|
|
| def predict_path(self, path: str | Path) -> DetectionResult: |
| path = Path(path) |
| with Image.open(path) as image: |
| return self.predict_image(image, path=path) |
|
|
|
|
| class UltraDetector: |
| """Blend HybridPlus and Sentry ConvNeXt into the strongest current practical ensemble.""" |
|
|
| backend_name = "ultra" |
|
|
| def __init__( |
| self, |
| *, |
| device: str | None = None, |
| threshold: float = 0.5, |
| primary_weight: float = DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| weight_path: str | Path | None = None, |
| model_name: str = DEFAULT_MODEL_NAME, |
| pretrained: str = DEFAULT_PRETRAINED, |
| hf_model: str = DEFAULT_HF_MODEL_ID, |
| hybrid_univfd_weight: float = DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| hybrid_plus_primary_weight: float = DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| ) -> None: |
| if not 0.0 <= primary_weight <= 1.0: |
| raise ValueError("primary_weight must be between 0 and 1") |
| self.threshold = threshold |
| self.primary_weight = primary_weight |
| self.primary = HybridPlusDetector( |
| device=device, |
| threshold=threshold, |
| primary_weight=hybrid_plus_primary_weight, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| ) |
| self.secondary = SentryConvNeXtDetector( |
| threshold=threshold, |
| device=device, |
| ) |
|
|
| def model_info(self) -> dict: |
| return { |
| "backend": self.backend_name, |
| "threshold": self.threshold, |
| "primary_weight": self.primary_weight, |
| "secondary_weight": round(1.0 - self.primary_weight, 6), |
| "components": { |
| "primary": self.primary.model_info(), |
| "secondary": self.secondary.model_info(), |
| }, |
| } |
|
|
| def predict_image(self, image: Image.Image, path: Path | None = None) -> DetectionResult: |
| return self.predict_images([image], paths=[path])[0] |
|
|
| def predict_images( |
| self, |
| images: Sequence[Image.Image], |
| paths: Sequence[Path | None] | None = None, |
| ) -> list[DetectionResult]: |
| if not images: |
| return [] |
| if paths is None: |
| paths = [None] * len(images) |
| if len(paths) != len(images): |
| raise ValueError("paths must have the same length as images") |
|
|
| primary_results = self.primary.predict_images(images, paths=paths) |
| secondary_results = self.secondary.predict_images(images, paths=paths) |
| results: list[DetectionResult] = [] |
| for primary_result, secondary_result, path in zip(primary_results, secondary_results, paths, strict=True): |
| probability_ai = ( |
| self.primary_weight * primary_result.probability_ai |
| + (1.0 - self.primary_weight) * secondary_result.probability_ai |
| ) |
| raw_score = ( |
| self.primary_weight * primary_result.raw_score |
| + (1.0 - self.primary_weight) * secondary_result.raw_score |
| ) |
| results.append( |
| _build_detection_result( |
| path=path, |
| probability_ai=probability_ai, |
| raw_score=raw_score, |
| threshold=self.threshold, |
| backend=self.backend_name, |
| ) |
| ) |
| return results |
|
|
| def predict_path(self, path: str | Path) -> DetectionResult: |
| path = Path(path) |
| with Image.open(path) as image: |
| return self.predict_image(image, path=path) |
|
|
| def predict_many( |
| self, |
| paths: Iterable[str | Path], |
| batch_size: int = 16, |
| ) -> list[DetectionResult]: |
| results: list[DetectionResult] = [] |
| batch_images: list[Image.Image] = [] |
| batch_paths: list[Path] = [] |
| for path_like in paths: |
| path = Path(path_like) |
| with Image.open(path) as image: |
| batch_images.append(_normalise_image(image)) |
| batch_paths.append(path) |
| if len(batch_images) >= batch_size: |
| results.extend(self.predict_images(batch_images, paths=batch_paths)) |
| batch_images = [] |
| batch_paths = [] |
| if batch_images: |
| results.extend(self.predict_images(batch_images, paths=batch_paths)) |
| return results |
|
|
|
|
| def create_detector( |
| backend: str = DEFAULT_BACKEND, |
| *, |
| device: str | None = None, |
| threshold: float = 0.5, |
| weight_path: str | Path | None = None, |
| model_name: str = DEFAULT_MODEL_NAME, |
| pretrained: str = DEFAULT_PRETRAINED, |
| hf_model: str = DEFAULT_HF_MODEL_ID, |
| hybrid_univfd_weight: float = DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| hybrid_plus_primary_weight: float = DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| ultra_primary_weight: float = DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| ) -> Detector: |
| backend = backend.lower() |
| if backend == "univfd": |
| return AIImageDetector( |
| device=device, |
| threshold=threshold, |
| model_name=model_name, |
| pretrained=pretrained, |
| weight_path=weight_path, |
| ) |
| if backend in {"hf", "huggingface"}: |
| return HuggingFaceImageDetector(model_id=hf_model, device=device, threshold=threshold) |
| if backend in {"nonescape", "nonescape-full"}: |
| return NonescapeFullDetector( |
| threshold=threshold, |
| device=device, |
| weight_path=weight_path, |
| ) |
| if backend in {"nonescape-mini", "nonescape_mini"}: |
| return NonescapeMiniDetector( |
| threshold=threshold, |
| device=device, |
| weight_path=weight_path, |
| ) |
| if backend in {"sentry", "sentry-convnext", "sentry-convnext-small"}: |
| return SentryConvNeXtDetector( |
| threshold=threshold, |
| device=device, |
| weight_path=weight_path, |
| ) |
| if backend in {"hybrid", "ensemble"}: |
| return HybridImageDetector( |
| device=device, |
| threshold=threshold, |
| univfd_weight=hybrid_univfd_weight, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| ) |
| if backend in {"hybrid-plus", "hybrid_plus", "stacked"}: |
| return HybridPlusDetector( |
| device=device, |
| threshold=threshold, |
| primary_weight=hybrid_plus_primary_weight, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| ) |
| if backend in {"ultra", "hybrid-ultra", "sentry-plus"}: |
| return UltraDetector( |
| device=device, |
| threshold=threshold, |
| primary_weight=ultra_primary_weight, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ) |
| raise ValueError( |
| "Unsupported backend: " |
| f"{backend!r}. Choose 'univfd', 'hf', 'nonescape-mini', 'nonescape-full', 'sentry-convnext-small', 'hybrid', 'hybrid-plus', or 'ultra'." |
| ) |
|
|
|
|
| def iter_images(path: str | Path, recursive: bool = True) -> list[Path]: |
| root = Path(path) |
| if root.is_file(): |
| return [root] if root.suffix.lower() in IMAGE_EXTENSIONS else [] |
| if not root.exists(): |
| return [] |
| globber = root.rglob if recursive else root.glob |
| return sorted(p for p in globber("*") if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS) |
|
|
|
|
| def _build_fallback_image_processor(config): |
| try: |
| from transformers import ViTImageProcessor |
| except ImportError as exc: |
| raise RuntimeError("Missing dependency: install transformers to use the HF backend.") from exc |
|
|
| image_size = getattr(config, "image_size", 224) |
| if isinstance(image_size, int): |
| size = {"height": image_size, "width": image_size} |
| else: |
| size = image_size |
| image_mean = getattr(config, "image_mean", [0.5, 0.5, 0.5]) |
| image_std = getattr(config, "image_std", [0.5, 0.5, 0.5]) |
| return ViTImageProcessor( |
| do_resize=True, |
| size=size, |
| resample=3, |
| do_rescale=True, |
| rescale_factor=1 / 255, |
| do_normalize=True, |
| image_mean=image_mean, |
| image_std=image_std, |
| ) |
|
|
|
|
| def _resolve_local_hf_snapshot(model_id: str) -> str | None: |
| candidates = [] |
| env_hf_home = os.environ.get("HF_HOME") |
| if env_hf_home: |
| candidates.append(Path(env_hf_home)) |
| candidates.append(Path.home() / ".cache" / "huggingface") |
|
|
| for hf_home in candidates: |
| hub_dir = hf_home / "hub" / f"models--{model_id.replace('/', '--')}" |
| if not hub_dir.exists(): |
| continue |
|
|
| ref_path = hub_dir / "refs" / "main" |
| if ref_path.exists(): |
| revision = ref_path.read_text(encoding="utf-8").strip() |
| snapshot_dir = hub_dir / "snapshots" / revision |
| if snapshot_dir.exists(): |
| return str(snapshot_dir) |
|
|
| snapshots_dir = hub_dir / "snapshots" |
| if snapshots_dir.exists(): |
| for snapshot_dir in sorted(snapshots_dir.iterdir()): |
| if snapshot_dir.is_dir(): |
| return str(snapshot_dir) |
| return None |
|
|
