Add pipeline orchestrator

run_pipeline.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python3
+"""
+Red Team Pipeline Orchestrator - See README.md for full documentation.
+Workflow: Input -> Protection (Paper X) -> Deepfake Generation -> Evaluation -> Report
+"""
+import os, sys, json, argparse, time, logging, numpy as np
+from pathlib import Path
+from typing import List, Dict
+from datetime import datetime
+from dataclasses import dataclass, asdict
+import torch
+from PIL import Image
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
+logger = logging.getLogger("RedTeamPipeline")
+
+@dataclass
+class PipelineConfig:
+    input_dir: str = "./data/test_images"
+    output_dir: str = "./results"
+    methods: List[str] = None
+    attacks: List[str] = None
+    fs_total_iter: int = 30
+    fs_noise_clamp: int = 12
+    fs_step_size: float = 1.0
+    fs_resize: int = 512
+    dg_iters: int = 800
+    dg_eps: float = 16.0/255.0
+    dg_step_size: float = 1.0/255.0
+    dg_batch_size: int = 1
+    dg_inference_steps: int = 4
+    vgm_steps: int = 1000
+    vgm_eps: float = 4.0/255.0
+    vgm_directed: bool = False
+    preprocessing_tests: List[str] = None
+    device: str = "cuda"
+    def __post_init__(self):
+        if self.methods is None: self.methods = ["faceshield", "diffusionguard", "vgmshield"]
+        if self.attacks is None: self.attacks = ["inpainting", "ip_adapter"]
+        if self.preprocessing_tests is None:
+            self.preprocessing_tests = ["jpeg_q75","jpeg_q50","jpeg_q25","blur_s1","blur_s2","resize_50","resize_75","noise_001","noise_005"]
+
+def run_faceshield(image_path, output_path, config):
+    logger.info(f"  [FaceShield] Protecting {os.path.basename(image_path)}")
+    try:
+        sys.path.insert(0, "/app/faceshield_src")
+        from modules.faceshield.protect import protect_image
+        return protect_image(image_path=image_path, output_path=output_path, device=config.device,
+            total_iter=config.fs_total_iter, noise_clamp=config.fs_noise_clamp, step_size=config.fs_step_size, resize_shape=config.fs_resize)
+    except Exception as e:
+        logger.error(f"  [FaceShield] Failed: {e}")
+        return {"method": "faceshield", "status": "error", "error": str(e)}
+
+def run_diffusionguard(image_path, mask_path, output_path, config):
+    logger.info(f"  [DiffusionGuard] Protecting {os.path.basename(image_path)}")
+    try:
+        sys.path.insert(0, "/app/diffusionguard_src")
+        from modules.diffusionguard.protect import protect_image
+        return protect_image(image_path=image_path, mask_path=mask_path, output_path=output_path, device=config.device,
+            iters=config.dg_iters, eps=config.dg_eps, step_size=config.dg_step_size, batch_size=config.dg_batch_size, num_inference_steps=config.dg_inference_steps)
+    except Exception as e:
+        logger.error(f"  [DiffusionGuard] Failed: {e}")
+        return {"method": "diffusionguard", "status": "error", "error": str(e)}
+
+def run_vgmshield(image_path, output_path, config):
+    logger.info(f"  [VGMShield] Protecting {os.path.basename(image_path)}")
+    try:
+        sys.path.insert(0, "/app/vgmshield_src")
+        from modules.vgmshield.protect import protect_image
+        return protect_image(input_path=image_path, output_path=output_path, directed=config.vgm_directed, eps=config.vgm_eps, steps=config.vgm_steps, device=config.device)
+    except Exception as e:
+        logger.error(f"  [VGMShield] Failed: {e}")
+        return {"method": "vgmshield", "status": "error", "error": str(e)}
+
+def run_inpainting_attack(image_path, mask_path, output_path, prompt="a photograph of a person", device="cuda"):
+    logger.info(f"  [Attack:Inpainting] On {os.path.basename(image_path)}")
+    try:
+        from diffusers import StableDiffusionInpaintPipeline
+        pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16).to(device)
+        src = Image.open(image_path).convert("RGB").resize((512,512))
+        mask = Image.open(mask_path).convert("RGB").resize((512,512))
+        result = pipe(prompt=prompt, image=src, mask_image=mask, num_inference_steps=50).images[0]
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        result.save(output_path)
+        del pipe; torch.cuda.empty_cache()
+        return {"attack": "inpainting", "status": "success", "output": output_path}
+    except Exception as e:
+        logger.error(f"  [Attack:Inpainting] Failed: {e}")
+        return {"attack": "inpainting", "status": "error", "error": str(e)}
+
+def run_ip_adapter_attack(image_path, output_path, prompt="a photo of a person", device="cuda"):
+    logger.info(f"  [Attack:IP-Adapter] On {os.path.basename(image_path)}")
+    try:
+        from diffusers import StableDiffusionPipeline
+        pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device)
+        pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")
+        pipe.set_ip_adapter_scale(0.8)
+        src = Image.open(image_path).convert("RGB").resize((512,512))
+        result = pipe(prompt=prompt, ip_adapter_image=src, num_inference_steps=50).images[0]
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        result.save(output_path)
+        del pipe; torch.cuda.empty_cache()
+        return {"attack": "ip_adapter", "status": "success", "output": output_path}
+    except Exception as e:
+        logger.error(f"  [Attack:IP-Adapter] Failed: {e}")
+        return {"attack": "ip_adapter", "status": "error", "error": str(e)}
+
+def apply_preprocessing(image_path, attack_name, output_path):
+    from modules.evaluation.evaluate import PURIFICATION_ATTACKS
+    img = Image.open(image_path).convert("RGB")
+    processed = PURIFICATION_ATTACKS[attack_name](img)
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    processed.save(output_path)
+    return output_path
+
+def generate_face_mask(image_path, output_path):
+    try:
+        import cv2
+        img = cv2.imread(image_path)
+        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+        faces = face_cascade.detectMultiScale(gray, 1.3, 5)
+        mask = np.zeros_like(img)
+        if len(faces) > 0:
+            for (x, y, w, h) in faces:
+                p = int(max(w, h) * 0.3)
+                mask[max(0,y-p):min(img.shape[0],y+h+p), max(0,x-p):min(img.shape[1],x+w+p)] = 255
+        else:
+            h, w = img.shape[:2]; mask[h//4:3*h//4, w//4:3*w//4] = 255
+        cv2.imwrite(output_path, mask)
+    except Exception:
+        img = Image.open(image_path); w, h = img.size
+        mask = Image.new("RGB", (w, h), "black")
+        from PIL import ImageDraw; draw = ImageDraw.Draw(mask)
+        m = min(w, h) // 4; draw.rectangle([m, m, w-m, h-m], fill="white")
+        mask.save(output_path)
+    return output_path
+
+def run_pipeline(config):
+    start_time = time.time()
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    run_dir = os.path.join(config.output_dir, f"run_{timestamp}")
+    os.makedirs(run_dir, exist_ok=True)
+    with open(os.path.join(run_dir, "config.json"), 'w') as f: json.dump(asdict(config), f, indent=2, default=str)
+    input_dir = Path(config.input_dir)
+    test_images = sorted([str(p) for p in input_dir.iterdir() if p.suffix.lower() in {'.png','.jpg','.jpeg','.bmp','.webp'}])
+    if not test_images: logger.error(f"No images in {config.input_dir}"); return
+    logger.info(f"Found {len(test_images)} images | Methods: {config.methods} | Attacks: {config.attacks}")
+    all_results = []
+    for img_idx, image_path in enumerate(test_images):
+        img_name = Path(image_path).stem
+        logger.info(f"\n{'='*60}\nImage {img_idx+1}/{len(test_images)}: {img_name}\n{'='*60}")
+        mask_dir = os.path.join(run_dir, "masks"); os.makedirs(mask_dir, exist_ok=True)
+        mask_path = os.path.join(mask_dir, f"{img_name}_mask.png")
+        generate_face_mask(image_path, mask_path)
+        for method in config.methods:
+            method_dir = os.path.join(run_dir, method, img_name); os.makedirs(method_dir, exist_ok=True)
+            original_copy = os.path.join(method_dir, "original.png")
+            Image.open(image_path).convert("RGB").save(original_copy)
+            protected_path = os.path.join(method_dir, "protected.png")
+            logger.info(f"\n--- Method: {method} ---")
+            if method == "faceshield": protection_result = run_faceshield(image_path, protected_path, config)
+            elif method == "diffusionguard": protection_result = run_diffusionguard(image_path, mask_path, protected_path, config)
+            elif method == "vgmshield": protection_result = run_vgmshield(image_path, protected_path, config)
+            else: continue
+            with open(os.path.join(method_dir, "protection_result.json"), 'w') as f: json.dump(protection_result, f, indent=2, default=str)
+            if protection_result.get("status") != "success": logger.error(f"Protection failed for {method}"); continue
+            if config.preprocessing_tests:
+                pp_dir = os.path.join(method_dir, "preprocessed"); os.makedirs(pp_dir, exist_ok=True)
+                for pp in config.preprocessing_tests:
+                    try: apply_preprocessing(protected_path, pp, os.path.join(pp_dir, f"{pp}.png")); logger.info(f"  [Preprocess] {pp}")
+                    except Exception as e: logger.error(f"  [Preprocess] {pp} failed: {e}")
+            for attack in config.attacks:
+                attack_dir = os.path.join(method_dir, f"attack_{attack}"); os.makedirs(attack_dir, exist_ok=True)
+                df_clean = os.path.join(attack_dir, "deepfake_clean.png"); df_protected = os.path.join(attack_dir, "deepfake_protected.png")
+                if attack == "inpainting":
+                    run_inpainting_attack(original_copy, mask_path, df_clean, device=config.device)
+                    run_inpainting_attack(protected_path, mask_path, df_protected, device=config.device)
+                elif attack == "ip_adapter":
+                    run_ip_adapter_attack(original_copy, df_clean, device=config.device)
+                    run_ip_adapter_attack(protected_path, df_protected, device=config.device)
+            try:
+                from modules.evaluation.evaluate import evaluate_protection
+                for attack in config.attacks:
+                    attack_dir = os.path.join(method_dir, f"attack_{attack}")
+                    df_c = os.path.join(attack_dir, "deepfake_clean.png"); df_p = os.path.join(attack_dir, "deepfake_protected.png")
+                    eval_result = evaluate_protection(original_copy, protected_path,
+                        df_c if os.path.exists(df_c) else None, df_p if os.path.exists(df_p) else None, method_name=method, device=config.device)
+                    ed = asdict(eval_result); ed["attack"] = attack; ed["image"] = img_name; all_results.append(ed)
+                    with open(os.path.join(attack_dir, "evaluation.json"), 'w') as f: json.dump(ed, f, indent=2)
+            except Exception as e: logger.error(f"  [Evaluation] Failed: {e}")
+    with open(os.path.join(run_dir, "report.json"), 'w') as f:
+        json.dump({"timestamp": timestamp, "config": asdict(config), "total_images": len(test_images), "methods": config.methods, "attacks": config.attacks, "time_s": time.time()-start_time, "results": all_results}, f, indent=2, default=str)
+    logger.info(f"\nPipeline complete! Results: {run_dir} | Time: {time.time()-start_time:.1f}s")
+    return all_results
+
+def main():
+    p = argparse.ArgumentParser(description="Red Team Pipeline")
+    p.add_argument("--input_dir", default="./data/test_images"); p.add_argument("--output_dir", default="./results")
+    p.add_argument("--methods", nargs="+", default=["faceshield","diffusionguard","vgmshield"])
+    p.add_argument("--attacks", nargs="+", default=["inpainting","ip_adapter"]); p.add_argument("--device", default="cuda")
+    p.add_argument("--fs_iter", type=int, default=30); p.add_argument("--fs_noise", type=int, default=12)
+    p.add_argument("--dg_iter", type=int, default=800); p.add_argument("--dg_eps", type=float, default=16/255)
+    p.add_argument("--vgm_steps", type=int, default=1000); p.add_argument("--vgm_eps", type=float, default=4/255)
+    p.add_argument("--skip_robustness", action="store_true")
+    a = p.parse_args()
+    config = PipelineConfig(input_dir=a.input_dir, output_dir=a.output_dir, methods=a.methods, attacks=a.attacks, device=a.device,
+        fs_total_iter=a.fs_iter, fs_noise_clamp=a.fs_noise, dg_iters=a.dg_iter, dg_eps=a.dg_eps, vgm_steps=a.vgm_steps, vgm_eps=a.vgm_eps,
+        preprocessing_tests=[] if a.skip_robustness else None)
+    run_pipeline(config)
+
+if __name__ == "__main__": main()