trufor_runner.py

from __future__ import annotations

import logging
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional

import numpy as np
from PIL import Image

LOGGER = logging.getLogger(__name__)


class TruForUnavailableError(RuntimeError):
    """Raised when the TruFor assets are missing or inference fails."""


@dataclass
class TruForResult:
    score: Optional[float]
    map_overlay: Optional[Image.Image]
    confidence_overlay: Optional[Image.Image]
    raw_scores: Dict[str, float]


class TruForEngine:
    """Wrapper that executes TruFor inference through docker or python backends."""

    def __init__(
        self,
        repo_root: Optional[Path] = None,
        weights_path: Optional[Path] = None,
        device: str = "cpu",
    ) -> None:
        self.base_dir = Path(__file__).resolve().parent
        self.device = device
        self.backend: Optional[str] = None
        self.status_message = "TruFor backend not initialized."

        backend_pref = os.environ.get("TRUFOR_BACKEND", "auto").lower()
        if backend_pref not in {"auto", "native", "docker"}:
            backend_pref = "auto"

        errors: list[str] = []

        if backend_pref in {"auto", "native"}:
            try:
                self._configure_native_backend(repo_root, weights_path)
                self.backend = "native"
                self.status_message = "TruFor ready (bundled python backend)."
            except TruForUnavailableError as exc:
                errors.append(f"Native backend unavailable: {exc}")
                if backend_pref == "native":
                    raise

        if self.backend is None and backend_pref in {"auto", "docker"}:
            try:
                self._configure_docker_backend()
                self.backend = "docker"
                self.status_message = f'TruFor ready (docker image "{self.docker_image}").'
            except TruForUnavailableError as exc:
                errors.append(f"Docker backend unavailable: {exc}")
                if backend_pref == "docker":
                    raise

        if self.backend is None:
            raise TruForUnavailableError(" | ".join(errors) if errors else "TruFor backend unavailable.")

    # ------------------------------------------------------------------
    # Backend configuration helpers
    # ------------------------------------------------------------------
    def _configure_docker_backend(self) -> None:
        if shutil.which("docker") is None:
            raise TruForUnavailableError("docker CLI not found on PATH.")

        test_docker_dir = self.base_dir / "test_docker"
        if not test_docker_dir.exists():
            raise TruForUnavailableError("test_docker directory not found in workspace.")

        image_name = os.environ.get("TRUFOR_DOCKER_IMAGE", "trufor")
        inspect = subprocess.run(
            ["docker", "image", "inspect", image_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
        if inspect.returncode != 0:
            raise TruForUnavailableError(
                f'Docker image "{image_name}" not found. Build it with "bash test_docker/docker_build.sh".'
            )

        weights_candidate = Path(os.environ.get("TRUFOR_DOCKER_WEIGHTS", self.base_dir / "weights")).expanduser()
        weight_file = weights_candidate / "trufor.pth.tar"
        self.docker_weights_dir: Optional[Path]
        self.docker_weights_dir = weight_file.parent if weight_file.exists() else None

        self.docker_runtime = os.environ.get("TRUFOR_DOCKER_RUNTIME")
        gpu_pref = os.environ.get("TRUFOR_DOCKER_GPU")
        if gpu_pref is None:
            gpu_pref = "-1" if self.device == "cpu" else "0"
        self.docker_gpu = gpu_pref

        gpus_arg = os.environ.get("TRUFOR_DOCKER_GPUS_ARG")
        if not gpus_arg and gpu_pref not in {"-1", "cpu", "none"}:
            gpus_arg = "all"
        self.docker_gpus_arg = gpus_arg

        self.docker_image = image_name

    def _configure_native_backend(self, _repo_root: Optional[Path], weights_path: Optional[Path]) -> None:
        try:
            from trufor_native import TruForBundledModel
        except ImportError as exc:  # pragma: no cover - packaging guard
            raise TruForUnavailableError("Bundled TruFor modules are not available.") from exc

        default_weights = self.base_dir / "weights" / "trufor.pth.tar"
        weight_candidate = weights_path or os.environ.get("TRUFOR_WEIGHTS") or default_weights
        weight_path = Path(weight_candidate).expanduser()
        if not weight_path.exists():
            raise TruForUnavailableError(
                f"TruFor weights missing at {weight_path}. Place trufor.pth.tar under weights/ or set TRUFOR_WEIGHTS."
            )

        try:
            self.native_model = TruForBundledModel(weight_path, device=self.device)
        except Exception as exc:  # pragma: no cover - propagate detailed failure
            raise TruForUnavailableError(f"Failed to initialise bundled TruFor model: {exc}") from exc

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def infer(self, image: Image.Image) -> TruForResult:
        if image is None:
            raise TruForUnavailableError("No image supplied to TruFor inference.")

        if self.backend == "docker":
            return self._infer_docker(image)
        if self.backend == "native":
            return self._infer_native(image)

        raise TruForUnavailableError("TruFor backend not configured.")

    # ------------------------------------------------------------------
    # Inference helpers
    # ------------------------------------------------------------------
    def _infer_native(self, image: Image.Image) -> TruForResult:
        outputs = self.native_model.predict(image)

        overlays: Dict[str, Optional[Image.Image]] = {"map": None, "conf": None}
        try:
            overlays["map"] = self._apply_heatmap(image, outputs.tamper_map)
        except Exception as exc:  # pragma: no cover - visualisation fallback
            LOGGER.debug("Failed to build tamper heatmap: %s", exc)

        if outputs.confidence_map is not None:
            try:
                overlays["conf"] = self._apply_heatmap(image, outputs.confidence_map)
            except Exception as exc:  # pragma: no cover
                LOGGER.debug("Failed to build confidence heatmap: %s", exc)

        raw_scores: Dict[str, float] = {
            "tamper_mean": float(np.mean(outputs.tamper_map)),
            "tamper_max": float(np.max(outputs.tamper_map)),
        }

        if outputs.confidence_map is not None:
            raw_scores["confidence_mean"] = float(np.mean(outputs.confidence_map))
            raw_scores["confidence_max"] = float(np.max(outputs.confidence_map))

        if outputs.detection_score is not None:
            raw_scores["tamper_score"] = float(outputs.detection_score)

        return TruForResult(
            score=outputs.detection_score,
            map_overlay=overlays["map"],
            confidence_overlay=overlays["conf"],
            raw_scores=raw_scores,
        )

    def _infer_docker(self, image: Image.Image) -> TruForResult:
        with tempfile.TemporaryDirectory(prefix="trufor_docker_") as workdir:
            workdir_path = Path(workdir)
            input_dir = workdir_path / "data"
            output_dir = workdir_path / "data_out"
            input_dir.mkdir(parents=True, exist_ok=True)
            output_dir.mkdir(parents=True, exist_ok=True)
            input_path = input_dir / "input.png"
            image.convert("RGB").save(input_path)

            cmd = ["docker", "run", "--rm"]
            if self.docker_runtime:
                cmd.extend(["--runtime", self.docker_runtime])

            gpu_flag = str(self.docker_gpu)
            if gpu_flag.lower() in {"cpu", "none"}:
                gpu_flag = "-1"
            if gpu_flag != "-1" and self.docker_gpus_arg:
                cmd.extend(["--gpus", self.docker_gpus_arg])

            cmd.extend([
                "-v",
                f"{input_dir.resolve()}:/data:ro",
                "-v",
                f"{output_dir.resolve()}:/data_out:rw",
            ])

            if self.docker_weights_dir is not None:
                cmd.extend([
                    "-v",
                    f"{self.docker_weights_dir.resolve()}:/weights:ro",
                ])

            cmd.append(self.docker_image)
            cmd.extend(
                [
                    "-gpu",
                    gpu_flag,
                    "-in",
                    "data/input.png",
                    "-out",
                    "data_out",
                ]
            )

            LOGGER.debug("Running TruFor docker command: %s", " ".join(cmd))
            result = subprocess.run(
                cmd,
                text=True,
                capture_output=True,
                check=False,
            )

            return self._process_results(result, output_dir, image)

    # ------------------------------------------------------------------
    # Result parsing
    # ------------------------------------------------------------------
    def _process_results(self, run_result: subprocess.CompletedProcess[str], output_dir: Path, image: Image.Image) -> TruForResult:
        if run_result.returncode != 0:
            stderr_tail = "\n".join(run_result.stderr.strip().splitlines()[-8:]) if run_result.stderr else ""
            LOGGER.error("TruFor stderr: %s", stderr_tail)
            raise TruForUnavailableError(
                "TruFor inference failed. Inspect dependencies and stderr:\n" + stderr_tail
            )

        npz_files = list(output_dir.rglob("*.npz"))
        if not npz_files:
            stdout_tail = "\n".join(run_result.stdout.strip().splitlines()[-8:]) if run_result.stdout else ""
            raise TruForUnavailableError(
                "TruFor inference produced no output files. Stdout tail:\n" + stdout_tail
            )

        data = np.load(npz_files[0], allow_pickle=False)
        tamper_map = data.get("map")
        conf_map = data.get("conf")
        score = float(data["score"]) if "score" in data.files else None

        overlays: Dict[str, Optional[Image.Image]] = {"map": None, "conf": None}
        try:
            overlays["map"] = self._apply_heatmap(image, tamper_map) if tamper_map is not None else None
        except Exception as exc:  # pragma: no cover
            LOGGER.debug("Failed to build tamper heatmap: %s", exc)

        try:
            overlays["conf"] = self._apply_heatmap(image, conf_map) if conf_map is not None else None
        except Exception as exc:  # pragma: no cover
            LOGGER.debug("Failed to build confidence heatmap: %s", exc)

        raw_scores: Dict[str, float] = {}
        if score is not None:
            raw_scores["tamper_score"] = score
        if tamper_map is not None:
            raw_scores["tamper_mean"] = float(np.mean(tamper_map))
            raw_scores["tamper_max"] = float(np.max(tamper_map))
        if conf_map is not None:
            raw_scores["confidence_mean"] = float(np.mean(conf_map))
            raw_scores["confidence_max"] = float(np.max(conf_map))

        return TruForResult(
            score=score,
            map_overlay=overlays["map"],
            confidence_overlay=overlays["conf"],
            raw_scores=raw_scores,
        )

    @staticmethod
    def _apply_heatmap(base: Image.Image, data: np.ndarray, alpha: float = 0.55) -> Image.Image:
        base_rgb = base.convert("RGB")
        if data is None or data.ndim != 2:
            raise ValueError("Expected a 2D map from TruFor")

        data = np.asarray(data, dtype=np.float32)
        if np.allclose(data.max(), data.min()):
            norm = np.zeros_like(data, dtype=np.float32)
        else:
            norm = (data - data.min()) / (data.max() - data.min())

        heat = np.zeros((*norm.shape, 3), dtype=np.uint8)
        heat[..., 0] = np.clip(norm * 255, 0, 255).astype(np.uint8)
        heat[..., 1] = np.clip(np.sqrt(norm) * 255, 0, 255).astype(np.uint8)
        heat[..., 2] = np.clip((1.0 - norm) * 255, 0, 255).astype(np.uint8)

        heat_img = Image.fromarray(heat, mode="RGB").resize(base_rgb.size, Image.BILINEAR)
        return Image.blend(base_rgb, heat_img, alpha)