FraudDetection/trufor-splicing-detecto

1import os2import uuid3import shutil4import numpy as np5import gradio as gr6from PIL import Image7 8from trufor_api import SplicingDetector9 10SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))11TRUFOR_FOLDER = os.path.join(SCRIPT_DIR, "TruFor_train_test")12 13detector = SplicingDetector(trufor_dir=TRUFOR_FOLDER, python_executable="python")14 15 16def predict(img: Image.Image):17    """Runs TruFor on an uploaded image and returns score + anomaly/confidence maps."""18    req_id = uuid.uuid4().hex  # unique per request19 20    uploads_dir = os.path.join(SCRIPT_DIR, "temp_uploads")21    out_root = os.path.join(SCRIPT_DIR, "temp_results")22    os.makedirs(uploads_dir, exist_ok=True)23    os.makedirs(out_root, exist_ok=True)24 25    image_path = os.path.join(uploads_dir, f"input_{req_id}.jpg")26    out_dir = os.path.join(out_root, f"run_{req_id}")27    os.makedirs(out_dir, exist_ok=True)28 29    img.save(image_path)30 31    try:32        results = detector.analyze_image(image_path, output_dir=out_dir)33        if results is None:34            raise gr.Error("TruFor failed. Check Space logs for the underlying error.")35 36        score = results["global_score"]37 38        anomaly = (np.clip(results["anomaly_map"], 0, 1) * 255).astype(np.uint8)39        conf = (np.clip(results["confidence_map"], 0, 1) * 255).astype(np.uint8)40 41        anomaly_img = Image.fromarray(anomaly)42        conf_img = Image.fromarray(conf)43 44        return score, anomaly_img, conf_img45 46    finally:47        # Cleanup (recommended for free Spaces)48        try:49            if os.path.exists(image_path):50                os.remove(image_path)51            shutil.rmtree(out_dir, ignore_errors=True)52        except Exception:53            pass54 55 56demo = gr.Interface(57    fn=predict,58    inputs=gr.Image(type="pil", label="Upload Image"),59    outputs=[60        gr.Number(label="Global Score"),61        gr.Image(type="pil", label="Anomaly Map"),62        gr.Image(type="pil", label="Confidence Map"),63    ],64    title="TruFor Splicing Detector",65    description="Upload an image to get a manipulation score + heatmaps from TruFor.",66)67 68# IMPORTANT: launch must be at top-level (not inside if __name__ == '__main__')69demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False)