""" Unified AI-Image & Deepfake Detector =================================== • Combines a generic AI-image detector (Swin-V2 + SuSy) *and* a deepfake-specialist face detector (Inception-ResNet V1). • Always runs both experts → fuses their calibrated scores. • Works on images **and** short videos (≤ 30 s). Add/keep in requirements.txt (versions pinned earlier): torch torchvision facenet-pytorch transformers torchcam captum timm mediapipe opencv-python-headless pillow scikit-image matplotlib gradio fpdf pandas numpy absl-py ttach """ # ───────────────────── bootstrap for extra wheels ────────────────────── import os, uuid, warnings, math, tempfile from pathlib import Path from typing import List, Tuple warnings.filterwarnings("ignore") def _ensure_deps(): try: import mediapipe, fpdf # noqa: F401 except ImportError: os.system("pip install --quiet --upgrade mediapipe fpdf") _ensure_deps() # ─────────────────────────────── imports ─────────────────────────────── import cv2 import gradio as gr import numpy as np import torch import torch.nn.functional as F from PIL import Image from fpdf import FPDF import mediapipe as mp from facenet_pytorch import InceptionResnetV1, MTCNN from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget from torchvision import transforms from transformers import AutoImageProcessor, AutoModelForImageClassification from torchcam.methods import GradCAM as TCGradCAM from captum.attr import Saliency from skimage.feature import graycomatrix, graycoprops import matplotlib.pyplot as plt import pandas as pd import spaces # ───────────────────────── runtime / models ──────────────────────────── plt.set_loglevel("ERROR") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Deep-fake specialist _face_det = MTCNN(select_largest=False, post_process=False, device=device).eval().to(device) _df_model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1, device=device) _df_model.load_state_dict(torch.load("resnet_inception.pth", map_location="cpu")["model_state_dict"]) _df_model.to(device).eval() _df_cam = GradCAM(_df_model, target_layers=[_df_model.block8.branch1[-1]], use_cuda=device.type == "cuda") # Helper: robust layer fetch def _get_layer(model, name: str): mods = dict(model.named_modules()) return mods.get(name) or next(m for n, m in mods.items() if n.endswith(name)) # Binary AI-image detector (Swin-V2) BIN_ID = "haywoodsloan/ai-image-detector-deploy" _bin_proc = AutoImageProcessor.from_pretrained(BIN_ID) _bin_mod = AutoModelForImageClassification.from_pretrained(BIN_ID).to(device).eval() _CAM_LAYER_BIN = "encoder.layers.3.blocks.1.layernorm_after" _bin_cam = TCGradCAM(_bin_mod, target_layer=_get_layer(_bin_mod, _CAM_LAYER_BIN)) # Generator classifier (SuSy — ScriptModule → Captum only) _susy_mod = torch.jit.load("SuSy.pt").to(device).eval() _GEN_CLASSES = ["Stable Diffusion 1.x", "DALL·E 3", "MJ V5/V6", "Stable Diffusion XL", "MJ V1/V2"] _PATCH, _TOP = 224, 5 _to_tensor = transforms.ToTensor() _to_gray = transforms.Compose([transforms.PILToTensor(), transforms.Grayscale()]) # ─────────────── calibration placeholders (optional tune) ────────────── _calib_df_slope, _calib_df_inter = 1.0, 0.0 _calib_ai_slope, _calib_ai_inter = 1.0, 0.0 # def _calibrate_df(p: float) -> float: # def _calibrate_ai(p: float) -> float: # return 1 / (1 + math.exp(-(_calib_ai_slope * (p + _calib_ai_inter)))) def _calibrate_df(p: float) -> float: # keep raw score for now return p def _calibrate_ai(p: float) -> float: return p # ───────────────────────────── misc helpers ──────────────────────────── UNCERTAIN_GAP = 0.10 MIN_FRAMES, MAX_SAMPLES = 4, 20 def _extract_landmarks(rgb: np.ndarray) -> Tuple[np.ndarray, np.ndarray | None]: mesh = mp.solutions.face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1) res = mesh.process(rgb); mesh.close() if not res.multi_face_landmarks: return rgb, None h, w, _ = rgb.shape out = rgb.copy() for lm in res.multi_face_landmarks[0].landmark: cx, cy = int(lm.x * w), int(lm.y * h) cv2.circle(out, (cx, cy), 1, (0, 255, 0), -1) return out, None def _overlay_cam(cam, base): # ---- NEW: make sure 'cam' is a NumPy array on CPU ---- if torch.is_tensor(cam): # covers torchcam output cam = cam.detach().cpu().numpy() # ------------------------------------------------------ cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-6) heat = Image.fromarray( (plt.cm.jet(cam)[:, :, :3] * 255).astype(np.uint8) ).resize((base.shape[1], base.shape[0]), Image.BICUBIC) return Image.blend( Image.fromarray(base).convert("RGBA"), heat.convert("RGBA"), alpha=0.45, ) def _render_pdf(title: str, verdict: str, conf: dict, pages: List[Image.Image]) -> str: out = Path(f"/tmp/report_{uuid.uuid4().hex}.pdf") pdf = FPDF(); pdf.set_auto_page_break(True, 15); pdf.add_page() pdf.set_font("Helvetica", size=14); pdf.cell(0, 10, title, ln=True, align="C") pdf.ln(4); pdf.set_font("Helvetica", size=12) pdf.multi_cell(0, 6, f"Verdict: {verdict}\n" f"Confidence -> Real {conf['real']:.3f} Fake {conf['fake']:.3f}") for idx, img in enumerate(pages): pdf.ln(4); pdf.set_font("Helvetica", size=11) pdf.cell(0, 6, f"Figure {idx+1}", ln=True) tmp = Path(tempfile.mktemp(suffix=".jpg")) img.convert("RGB").save(tmp, format="JPEG") # ← add .convert("RGB") pdf.image(str(tmp), x=10, w=90) tmp.unlink(missing_ok=True) pdf.output(out) return str(out) # ────────────────────────── SuSy helpers (saliency) ──────────────────── def _susy_cam(tensor: torch.Tensor, class_idx: int) -> np.ndarray: sal = Saliency(_susy_mod) grad = sal.attribute(tensor, target=class_idx).abs().mean(1, keepdim=True) return grad.squeeze().detach().cpu().numpy() @spaces.GPU def _susy_predict(img: Image.Image): w, h = img.size npx, npy = max(1, w // _PATCH), max(1, h // _PATCH) patches = np.zeros((npx * npy, _PATCH, _PATCH, 3), dtype=np.uint8) for i in range(npx): for j in range(npy): x, y = i * _PATCH, j * _PATCH patches[i*npy + j] = np.array(img.crop((x, y, x+_PATCH, y+_PATCH)) .resize((_PATCH, _PATCH))) contrasts = [] for p in patches: g = _to_gray(Image.fromarray(p)).squeeze(0).numpy() glcm = graycomatrix(g, [5], [0], 256, symmetric=True, normed=True) contrasts.append(graycoprops(glcm, "contrast")[0, 0]) idx = np.argsort(contrasts)[::-1][:_TOP] tens = torch.from_numpy(patches[idx].transpose(0, 3, 1, 2)).float() / 255.0 with torch.no_grad(): probs = _susy_mod(tens.to(device)).softmax(-1).mean(0).cpu().numpy()[1:] return dict(zip(_GEN_CLASSES, probs)) # ───────────────────────────── fusion math ───────────────────────────── def _fuse(p_ai: float, p_df: float) -> float: return 1 - (1 - p_ai) * (1 - p_df) def _verdict(p: float) -> str: return "uncertain" if abs(p - 0.5) <= UNCERTAIN_GAP else ("fake" if p > 0.5 else "real") # ─────────────────────────── IMAGE PIPELINE ──────────────────────────── @spaces.GPU def _predict_image(pil: Image.Image): gallery: List[Image.Image] = [] # Deep-fake path try: face = _face_det(pil) except Exception: face = None if face is not None: ft = F.interpolate(face.unsqueeze(0), (256, 256), mode="bilinear", align_corners=False).float() / 255.0 p_df_raw = torch.sigmoid(_df_model(ft.to(device))).item() p_df = _calibrate_df(p_df_raw) crop_np = (ft.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) cam_df = _df_cam(ft, [ClassifierOutputTarget(0)])[0] gallery.append(_overlay_cam(cam_df, crop_np)) gallery.append(Image.fromarray(_extract_landmarks( cv2.cvtColor(np.array(pil), cv2.COLOR_BGR2RGB))[0])) else: p_df = 0.5 # Binary AI model inp_bin = _bin_proc(images=pil, return_tensors="pt").to(device) logits = _bin_mod(**inp_bin).logits.softmax(-1)[0] p_ai_raw = logits[0].item() p_ai = _calibrate_ai(p_ai_raw) winner_idx = 0 if p_ai_raw >= logits[1].item() else 1 inp_bin_h = {k: v.clone().detach().requires_grad_(True) for k, v in inp_bin.items()} cam_bin = _bin_cam(winner_idx, scores=_bin_mod(**inp_bin_h).logits)[0] gallery.append(_overlay_cam(cam_bin, np.array(pil))) # Generator breakdown (SuSy) if AI bar_plot = gr.update(visible=False) if p_ai_raw > logits[1].item(): gen_probs = _susy_predict(pil) bar_plot = gr.update(value=pd.DataFrame(gen_probs.items(), columns=["class", "prob"]), visible=True) susy_in = _to_tensor(pil.resize((224, 224))).unsqueeze(0).to(device) g_idx = _susy_mod(susy_in)[0, 1:].argmax().item() + 1 cam_susy = _susy_cam(susy_in, g_idx) gallery.append(_overlay_cam(cam_susy, np.array(pil))) # Fusion p_final = _fuse(p_ai, p_df) verdict = _verdict(p_final) conf = {"real": round(1-p_final, 4), "fake": round(p_final, 4)} pdf = _render_pdf("Unified Detector", verdict, conf, gallery[:3]) return verdict, conf, gallery, bar_plot, pdf # ─────────────────────────── VIDEO PIPELINE ──────────────────────────── def _sample_idx(n): # max 20 evenly spaced return list(range(n)) if n <= MAX_SAMPLES else np.linspace(0, n-1, MAX_SAMPLES, dtype=int) @spaces.GPU def _predict_video(path: str): cap = cv2.VideoCapture(path); total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1 probs, frames = [], [] for i in _sample_idx(total): cap.set(cv2.CAP_PROP_POS_FRAMES, i) ok, frm = cap.read() if not ok: continue pil = Image.fromarray(cv2.cvtColor(frm, cv2.COLOR_BGR2RGB)) verdict, conf, _, _, _ = _predict_image(pil) probs.append(conf["fake"]) if len(frames) < MIN_FRAMES: frames.append(Image.fromarray(frm)) cap.release() if not probs: blank = Image.new("RGB", (256, 256)) return "No frames analysed", {"real": 0, "fake": 0}, [blank] p_final = float(np.mean(probs)) return _verdict(p_final), {"real": round(1-p_final, 4), "fake": round(p_final, 4)}, frames # ───────────────────────────────── UI ────────────────────────────────── _css = "footer{visibility:hidden!important}.logo,#logo{display:none!important}" with gr.Blocks(css=_css, title="Unified AI-Fake & Deepfake Detector") as demo: gr.Markdown(""" ## Unified AI-Fake & Deepfake Detector Upload an **image** or a short **video**. The app fuses two complementary models, then shows heat-maps & a PDF report. """) with gr.Tab("Image"): with gr.Row(): with gr.Column(scale=1): img_in = gr.Image(label="Upload image", type="pil") btn_i = gr.Button("Analyze") with gr.Column(scale=2): txt_v = gr.Textbox(label="Verdict", interactive=False) lbl_c = gr.Label(label="Confidence") gal = gr.Gallery(label="Explanations", columns=3, height=320) bar = gr.BarPlot(x="class", y="prob", title="Likely generator", y_label="probability", visible=False) pdf_f = gr.File(label="Download PDF report") btn_i.click(_predict_image, img_in, [txt_v, lbl_c, gal, bar, pdf_f]) with gr.Tab("Video"): with gr.Row(): with gr.Column(scale=1): vid_in = gr.Video(label="Upload MP4/AVI", format="mp4") btn_v = gr.Button("Analyze") with gr.Column(scale=2): txt_vv = gr.Textbox(label="Verdict", interactive=False) lbl_cv = gr.Label(label="Confidence") gal_v = gr.Gallery(label="Sample frames", columns=4, height=240) btn_v.click(_predict_video, vid_in, [txt_vv, lbl_cv, gal_v]) demo.launch(share=True, show_api=False)