FraudDetection/trufor-splicing-detecto

1import os2import sys3import glob4import subprocess5import numpy as np6from PIL import Image7 8 9class SplicingDetector:10    def __init__(self, trufor_dir, python_executable="python"):11        """12        Initialize the detector by pointing it to your TruFor installation.13 14        :param trufor_dir: The absolute or relative path to the 'TruFor_train_test' folder.15        :param python_executable: The python executable to run test.py with.16        """17        self.trufor_dir = os.path.abspath(trufor_dir)18        self.python_exec = python_executable19        self.test_script = os.path.join(self.trufor_dir, "test.py")20        self.weights_path = "pretrained_models/trufor.pth.tar"  # Relative to trufor_dir21 22    def analyze_image(self, image_path, output_dir="./temp_results"):23        """24        Runs the TruFor model on an image and returns the raw forensic data.25 26        Returns:27          dict with keys: anomaly_map, confidence_map, global_score28          or None if inference fails.29        """30        # Validate folder and test.py31        if not os.path.isdir(self.trufor_dir):32            raise NotADirectoryError(33                f"TruFor directory not found or invalid: {self.trufor_dir}\n"34                f"Tip: TRUFOR_FOLDER must point to the folder that contains test.py."35            )36 37        if not os.path.isfile(self.test_script):38            raise FileNotFoundError(39                f"Could not find test.py at: {self.test_script}\n"40                f"Tip: TRUFOR_FOLDER must be the 'TruFor_train_test' folder (the one that contains test.py)."41            )42 43        abs_image_path = os.path.abspath(image_path)44        abs_output_dir = os.path.abspath(output_dir)45 46        os.makedirs(abs_output_dir, exist_ok=True)47 48        print(f"🔍 Analyzing {os.path.basename(abs_image_path)} for splicing...")49 50        command = [51            self.python_exec,52            self.test_script,53            "-g",54            "-1",  # CPU55            "-in",56            abs_image_path,57            "-out",58            abs_output_dir,59            "-exp",60            "trufor_ph3",61            "TEST.MODEL_FILE",62            self.weights_path,63        ]64 65        # Ensure repo-local imports work (dataset/, lib/, etc.)66        env = os.environ.copy()67        env["PYTHONPATH"] = self.trufor_dir + os.pathsep + env.get("PYTHONPATH", "")68 69        try:70            subprocess.run(71                command,72                cwd=self.trufor_dir,73                env=env,74                check=True,75                capture_output=True,76                text=True,77            )78        except subprocess.CalledProcessError as e:79            print("🚨 TruFor Execution Failed!")80            # stderr is typically the most useful; stdout may contain progress logs.81            print("STDERR:\n", e.stderr)82            print("STDOUT:\n", e.stdout)83            return None84 85        # Robustly locate the output .npz (avoids stale file issues)86        npz_candidates = glob.glob(os.path.join(abs_output_dir, "*.npz"))87        if not npz_candidates:88            print(f"Error: No .npz output found in {abs_output_dir}")89            return None90 91        # pick the newest .npz in the output directory92        npz_path = max(npz_candidates, key=os.path.getmtime)93 94        print(f"✅ Analysis complete. Extracting forensic data from: {os.path.basename(npz_path)}")95        data = np.load(npz_path)96 97        results = {98            "anomaly_map": data["map"],99            "confidence_map": data["conf"],100            "global_score": float(data["score"]) if "score" in data else None,101        }102        return results103 104    def visualize_results(self, image_path, results):105        """106        Local-only visualization (not needed for Hugging Face Spaces).107        Kept for convenience when running on your machine.108        """109        import matplotlib.pyplot as plt  # lazy import so server does not require matplotlib110 111        original_img = Image.open(image_path).convert("RGB")112 113        plt.figure(figsize=(15, 5))114 115        plt.subplot(1, 3, 1)116        plt.imshow(original_img)117        plt.title("Original Image")118        plt.axis("off")119 120        plt.subplot(1, 3, 2)121        plt.imshow(results["anomaly_map"], cmap="jet", vmin=0, vmax=1)122        score = results.get("global_score")123        plt.title(f"Anomaly Map\nGlobal Score: {score:.4f}" if score is not None else "Anomaly Map")124        plt.colorbar(fraction=0.046, pad=0.04)125        plt.axis("off")126 127        plt.subplot(1, 3, 3)128        plt.imshow(results["confidence_map"], cmap="magma", vmin=0, vmax=1)129        plt.title("Confidence Map")130        plt.colorbar(fraction=0.046, pad=0.04)131        plt.axis("off")132 133        plt.tight_layout()134        plt.show()135 136 137if __name__ == "__main__":138    # Local test runner (optional)139    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))140    TRUFOR_FOLDER = os.path.join(SCRIPT_DIR, "TruFor_train_test")141 142    detector = SplicingDetector(trufor_dir=TRUFOR_FOLDER, python_executable=sys.executable)143 144    test_image = os.path.join(SCRIPT_DIR, "sample.jpg")145    forensic_data = detector.analyze_image(test_image, output_dir=os.path.join(SCRIPT_DIR, "temp_results_local"))146 147    if forensic_data:148        print(f"✅ Global AI Score: {forensic_data['global_score']:.4f}")149        detector.visualize_results(test_image, forensic_data)