hiimwsta/AI-Source-Forensics
0
1import os2import json3import torch4import torch.nn as nn5from flask import Flask, request, jsonify, render_template, send_from_directory6from werkzeug.utils import secure_filename7from torchvision import transforms8from PIL import Image9import io10import numpy as np11import cv212import time13import mediapipe as mp 14from model_architecture import MultiBranchDetector15 16app = Flask(__name__)17project_root = os.path.dirname(os.path.abspath(__file__))18app.config['UPLOAD_FOLDER'] = os.path.join(project_root, 'static', 'uploads')19os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)20 21MODEL_PATH = "multi_branch_detector_genimage.pth"22CLASSES_JSON_PATH = "class_names.json"23IMAGE_SIZE = 22424device = torch.device("cuda" if torch.cuda.is_available() else "cpu")25 26mp_face_detection = mp.solutions.face_detection27face_detection = mp_face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5)28 29# --- Load Class Names ---30def load_classes():31 if os.path.exists(CLASSES_JSON_PATH):32 with open(CLASSES_JSON_PATH, 'r') as f:33 return json.load(f)34 return ['ai-gen', 'real'] 35 36CLASS_NAMES = load_classes()37NUM_CLASSES = len(CLASS_NAMES)38 39# --- Feature Extraction Tools ---40class FFTTransform:41 def __call__(self, img):42 img_np = np.array(img.convert('L'))43 f = np.fft.fft2(img_np)44 fshift = np.fft.fftshift(f)45 magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1e-9)46 magnitude_spectrum = np.clip(magnitude_spectrum, 0, 255).astype(np.uint8)47 return Image.fromarray(cv2.merge([magnitude_spectrum]*3))48 49class LBPTransform:50 def __call__(self, img):51 img_np = np.array(img.convert('L'), dtype=np.uint8)52 h, w = img_np.shape53 lbp = np.zeros((h, w), dtype=np.uint8)54 offsets = [(-1, -1), (-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1)]55 for i, (dy, dx) in enumerate(offsets):56 shifted = np.roll(img_np, shift=(-dy, -dx), axis=(0, 1))57 lbp += ((shifted >= img_np).astype(np.uint8) << (7 - i))58 return Image.fromarray(cv2.merge([lbp]*3))59 60# --- Load Model ---61def load_model():62 model = MultiBranchDetector(num_classes=NUM_CLASSES)63 if os.path.exists(MODEL_PATH):64 model.load_state_dict(torch.load(MODEL_PATH, map_location=device, weights_only=True))65 model = model.to(device)66 model.eval()67 return model68 69global_model = load_model()70 71# --- Dynamic Diagnostic Log Generator ---72def generate_diagnostic_log(class_name, highest_domain):73 target = class_name.lower().replace(' ', '_').replace('-', '')74 if 'real' in target:75 return "Success: Natural optical sensor noise profile verified. No synthetic upsampling artifacts detected."76 77 diagnostics = {78 'midjourney': {79 'Spatial': "High-contrast structural anomalies and unnatural edge gradients detected.",80 'Frequency': "High-frequency spectral artifacts consistent with diffusion denoising steps.",81 'Texture': "Unnatural micro-texture uniformity and hyper-smoothness detected in LBP analysis."82 },83 'stable_diffusion': {84 'Spatial': "Latent space decoding artifacts and pixel-level spatial inconsistencies identified.",85 'Frequency': "Periodic frequency anomalies indicative of VAE (Variational Autoencoder) upsampling.",86 'Texture': "Synthetic local binary patterns suggesting artificial texture mapping."87 },88 'stylegan': {89 'Spatial': "Asymmetric feature alignment and GAN-specific spatial droplet artifacts found.",90 'Frequency': "Pronounced checkerboard artifacts in the frequency spectrum characteristic of GAN upsampling.",91 'Texture': "Phase discrepancies and unnatural texture repetition detected."92 },93 'deepfacelab': {94 'Spatial': "Color blending inconsistencies and resolution mismatch along facial boundaries detected.",95 'Frequency': "Frequency spectrum mismatch between central facial region and background image.",96 'Texture': "Texture discontinuity along facial landmarks and blending mask edges."97 }98 }99 default_log = f"Synthetic generation patterns detected with primary anomalies in the {highest_domain} domain."100 specific_log = diagnostics.get(target, {}).get(highest_domain, default_log)101 return f"Warning: Source attributed to {class_name.upper()}. {specific_log}"102 103# ==========================================104# 💡 Smart Crop using Google MediaPipe105# ==========================================106def smart_crop(img_pil):107 img_np = np.array(img_pil)108 results = face_detection.process(img_np)109 110 if results.detections:111 largest_detection = max(results.detections, 112 key=lambda d: d.location_data.relative_bounding_box.width * d.location_data.relative_bounding_box.height)113 114 bbox = largest_detection.location_data.relative_bounding_box115 ih, iw, _ = img_np.shape116 117 x = int(bbox.xmin * iw)118 y = int(bbox.ymin * ih)119 w = int(bbox.width * iw)120 h = int(bbox.height * ih)121 122 pad_x, pad_y = int(w * 0.3), int(h * 0.3)123 x1, y1 = max(0, x - pad_x), max(0, y - pad_y)124 x2, y2 = min(iw, x + w + pad_x), min(ih, y + h + pad_y)125 126 face_crop = img_pil.crop((x1, y1, x2, y2))127 return face_crop.resize((IMAGE_SIZE, IMAGE_SIZE))128 else:129 fallback = transforms.Compose([130 transforms.Resize(IMAGE_SIZE),131 transforms.CenterCrop(IMAGE_SIZE)132 ])133 return fallback(img_pil)134 135# --- Save XAI Feature Maps ---136def save_feature_visualizations(img_cropped, base_filename):137 fft_img = FFTTransform()(img_cropped)138 lbp_img = LBPTransform()(img_cropped)139 140 timestamp = int(time.time())141 crop_name = f"crop_{timestamp}_{base_filename}.png"142 fft_name = f"fft_{timestamp}_{base_filename}.png"143 lbp_name = f"lbp_{timestamp}_{base_filename}.png"144 145 img_cropped.save(os.path.join(app.config['UPLOAD_FOLDER'], crop_name))146 fft_img.save(os.path.join(app.config['UPLOAD_FOLDER'], fft_name))147 lbp_img.save(os.path.join(app.config['UPLOAD_FOLDER'], lbp_name))148 149 return crop_name, fft_name, lbp_name150 151# --- Core Inference Logic ---152def predict_pil_image(img_cropped, model):153 to_tensor = transforms.ToTensor()154 norm = transforms.Normalize([0.5]*3, [0.5]*3)155 156 rgb_t = norm(to_tensor(img_cropped)).unsqueeze(0).to(device)157 fft_t = norm(to_tensor(FFTTransform()(img_cropped))).unsqueeze(0).to(device)158 lbp_t = norm(to_tensor(LBPTransform()(img_cropped))).unsqueeze(0).to(device)159 160 output, (s_rgb, s_fft, s_lbp) = model.get_domain_contribution(rgb_t, fft_t, lbp_t)161 probs = torch.nn.functional.softmax(output, dim=1).squeeze()162 163 top_prob, top_idx = torch.max(probs, 0)164 top_class_name = CLASS_NAMES[top_idx.item()]165 total_score = s_rgb + s_fft + s_lbp + 1e-9166 167 return {168 "source": top_class_name.upper(), 169 "confidence": top_prob.item() * 100,170 "p_rgb": (s_rgb / total_score) * 100,171 "p_fft": (s_fft / total_score) * 100,172 "p_lbp": (s_lbp / total_score) * 100173 }174 175def predict_image(image_bytes, model, filename):176 try:177 img = Image.open(io.BytesIO(image_bytes)).convert("RGB")178 img_cropped = smart_crop(img)179 180 crop_file, fft_file, lbp_file = save_feature_visualizations(img_cropped, secure_filename(filename))181 182 res = predict_pil_image(img_cropped, model)183 184 is_real = 'REAL' in res['source']185 prediction_status = 'real' if is_real else 'ai-gen'186 highest_domain = max([('Spatial', res['p_rgb']), ('Frequency', res['p_fft']), ('Texture', res['p_lbp'])], key=lambda x: x[1])[0]187 log_msg = generate_diagnostic_log(res['source'], highest_domain)188 189 return {190 "prediction": prediction_status, "source": res['source'], "confidence": res['confidence'],191 "rgb_score": res['confidence'] * (res['p_rgb']/33.3), "fft_score": res['confidence'] * (res['p_fft']/33.3), "lbp_score": res['confidence'] * (res['p_lbp']/33.3),192 "log": log_msg, "is_video": False, "frames": 1,193 "crop_file": crop_file, "fft_file": fft_file, "lbp_file": lbp_file194 }195 except Exception as e:196 print(f"Image Prediction Error: {e}")197 return None198 199def analyze_video(video_path, model, filename, num_frames=12):200 try:201 cap = cv2.VideoCapture(video_path)202 total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))203 if total_frames == 0: return None204 205 step = max(1, total_frames // num_frames)206 frame_results = []207 crop_file, fft_file, lbp_file = None, None, None208 209 for i in range(num_frames):210 frame_idx = i * step211 if frame_idx >= total_frames: break212 213 cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)214 ret, frame = cap.read()215 if not ret: break216 217 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)218 pil_img = Image.fromarray(frame_rgb)219 img_cropped = smart_crop(pil_img)220 221 if i == num_frames // 2:222 crop_file, fft_file, lbp_file = save_feature_visualizations(img_cropped, secure_filename(filename))223 224 res = predict_pil_image(img_cropped, model)225 frame_results.append(res)226 227 cap.release()228 if not frame_results: return None229 230 source_counts = {}231 total_conf, total_rgb, total_fft, total_lbp = 0, 0, 0, 0232 n = len(frame_results)233 calibrated_count = 0234 235 for res in frame_results:236 src = res['source']237 conf = res['confidence']238 239 if src in ['STABLE DIFFUSION', 'MIDJOURNEY', 'STYLEGAN'] and conf < 85.0:240 src = 'REAL'241 conf = 100.0 - conf242 calibrated_count += 1243 244 source_counts[src] = source_counts.get(src, 0) + 1245 total_conf += conf246 total_rgb += res['p_rgb']247 total_fft += res['p_fft']248 total_lbp += res['p_lbp']249 250 ai_threats = {k: v for k, v in source_counts.items() if 'REAL' not in k}251 if ai_threats:252 highest_ai_threat = max(ai_threats, key=ai_threats.get)253 if ai_threats[highest_ai_threat] >= (n // 4):254 top_source = highest_ai_threat255 else:256 top_source = max(source_counts, key=source_counts.get)257 else:258 top_source = 'REAL'259 260 avg_conf = total_conf / n261 avg_p_rgb = total_rgb / n262 avg_p_fft = total_fft / n263 avg_p_lbp = total_lbp / n264 265 is_real = 'REAL' in top_source266 prediction_status = 'real' if is_real else 'ai-gen'267 highest_domain = max([('Spatial', avg_p_rgb), ('Frequency', avg_p_fft), ('Texture', avg_p_lbp)], key=lambda x: x[1])[0]268 base_log = generate_diagnostic_log(top_source, highest_domain)269 270 if is_real and calibrated_count > 0:271 log_msg = f"[Video Analysis | Auto-Calibrated] MP4 compression filtered. " + base_log272 elif not is_real:273 log_msg = f"[Video Analysis | Threat Priority Triggered] Deepfake temporal anomalies detected. " + base_log274 else:275 log_msg = f"[Video Analysis | Examined {n} Keyframes] " + base_log276 277 return {278 "prediction": prediction_status, "source": top_source, "confidence": avg_conf,279 "rgb_score": avg_conf * (avg_p_rgb/33.3), "fft_score": avg_conf * (avg_p_fft/33.3), "lbp_score": avg_conf * (avg_p_lbp/33.3),280 "log": log_msg, "is_video": True, "frames": n,281 "crop_file": crop_file, "fft_file": fft_file, "lbp_file": lbp_file282 }283 except Exception as e:284 print(f"Video Processing Error: {e}")285 return None286 287# --- Routing ---288@app.route('/', methods=['GET', 'POST'])289def index():290 if request.method == 'POST':291 if 'media' not in request.files: return jsonify({"error": "No file part"}), 400292 file = request.files['media']293 if file.filename == '': return jsonify({"error": "No selected file"}), 400294 295 if file:296 filename = secure_filename(file.filename)297 path = os.path.join(app.config['UPLOAD_FOLDER'], filename)298 file.save(path)299 300 ext = filename.rsplit('.', 1)[1].lower() if '.' in filename else ''301 302 if ext in ['mp4', 'avi', 'mov', 'webm']:303 res = analyze_video(path, global_model, filename, num_frames=12)304 else:305 file.seek(0)306 res = predict_image(file.read(), global_model, filename)307 308 if res:309 res['media_file'] = filename310 return jsonify(res)311 else:312 return jsonify({"error": "Error processing media"}), 500313 314 return render_template('index.html', prediction=None)315 316@app.route('/static/uploads/<filename>')317def uploaded_file(filename):318 return send_from_directory(app.config['UPLOAD_FOLDER'], filename)319 320if __name__ == '__main__':321 app.run(host='0.0.0.0', port=7860)