Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image | |
| import mediapipe as mp | |
| # Robust submodule imports to bypass __init__ issues on specific platforms | |
| try: | |
| import mediapipe.solutions.face_mesh as mp_face_mesh | |
| except ImportError: | |
| try: | |
| from mediapipe.python.solutions import face_mesh as mp_face_mesh | |
| except ImportError: | |
| import mediapipe.solutions.face_mesh as mp_face_mesh | |
| from collections import deque | |
| from huggingface_hub import hf_hub_download | |
| # --- ConvNeXt Model Architecture --- | |
| class LayerNorm(nn.Module): | |
| def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(normalized_shape)) | |
| self.bias = nn.Parameter(torch.zeros(normalized_shape)) | |
| self.eps = eps | |
| self.data_format = data_format | |
| if self.data_format not in ["channels_last", "channels_first"]: | |
| raise NotImplementedError | |
| self.normalized_shape = (normalized_shape, ) | |
| def forward(self, x): | |
| if self.data_format == "channels_last": | |
| return torch.nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) | |
| elif self.data_format == "channels_first": | |
| u = x.mean(1, keepdim=True) | |
| s = (x - u).pow(2).mean(1, keepdim=True) | |
| x = (x - u) / torch.sqrt(s + self.eps) | |
| x = self.weight[:, None, None] * x + self.bias[:, None, None] | |
| return x | |
| class Block(nn.Module): | |
| def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6): | |
| super().__init__() | |
| self.conv_dw = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) | |
| self.norm = LayerNorm(dim, eps=1e-6) | |
| self.mlp = nn.ModuleDict({ | |
| "fc1": nn.Linear(dim, 4 * dim), | |
| "act": nn.GELU(), | |
| "fc2": nn.Linear(4 * dim, dim) | |
| }) | |
| self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), | |
| requires_grad=True) if layer_scale_init_value > 0 else None | |
| self.drop_path = nn.Identity() | |
| def forward(self, x): | |
| input = x | |
| x = self.conv_dw(x) | |
| x = x.permute(0, 2, 3, 1) | |
| x = self.norm(x) | |
| x = self.mlp['fc1'](x) | |
| x = self.mlp['act'](x) | |
| x = self.mlp['fc2'](x) | |
| if self.gamma is not None: | |
| x = self.gamma * x | |
| x = x.permute(0, 3, 1, 2) | |
| x = input + self.drop_path(x) | |
| return x | |
| class ConvNeXt(nn.Module): | |
| def __init__(self, in_chans=3, depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024]): | |
| super().__init__() | |
| self.stem = nn.Sequential( | |
| nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), | |
| LayerNorm(dims[0], eps=1e-6, data_format="channels_first") | |
| ) | |
| self.stages = nn.ModuleList() | |
| for i in range(4): | |
| stage = nn.ModuleDict() | |
| if i > 0: | |
| stage['downsample'] = nn.Sequential( | |
| LayerNorm(dims[i-1], eps=1e-6, data_format="channels_first"), | |
| nn.Conv2d(dims[i-1], dims[i], kernel_size=2, stride=2), | |
| ) | |
| stage['blocks'] = nn.Sequential(*[Block(dim=dims[i]) for _ in range(depths[i])]) | |
| self.stages.append(stage) | |
| self.head = nn.ModuleDict({ | |
| "avgpool": nn.AdaptiveAvgPool2d((1, 1)), | |
| "norm": LayerNorm(dims[-1], eps=1e-6, data_format="channels_last") | |
| }) | |
| def forward_features(self, x): | |
| x = self.stem(x) | |
| for i in range(4): | |
| if 'downsample' in self.stages[i]: | |
| x = self.stages[i]['downsample'](x) | |
| x = self.stages[i]['blocks'](x) | |
| x = self.head['avgpool'](x) | |
| x = x.view(x.size(0), -1) | |
| return self.head['norm'](x) | |
| def forward(self, x): | |
| return self.forward_features(x) | |
| class DeepfakeModel(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.backbone = ConvNeXt() | |
| self.classifier = nn.Sequential( | |
| nn.Linear(1024, 512), | |
| nn.ReLU(), | |
| nn.Dropout(0.2), | |
| nn.Linear(512, 2) | |
| ) | |
| def forward(self, x): | |
| x = self.backbone(x) | |
| return self.classifier(x) | |
| class DeepfakeDetector: | |
| def __init__(self): | |
| self.model = DeepfakeModel() | |
| self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| self.model.to(self.device).eval() | |
| # Hugging Face Weights Loading - Stripping any hidden newlines/spaces | |
| repo_id = os.environ.get('HF_REPO_ID', '').strip() | |
| filename = os.environ.get('HF_FILENAME', 'convnext_video_fixed.pth').strip() | |
| weights_path = None | |
| if os.path.exists(filename): weights_path = filename | |
| elif repo_id: | |
| try: | |
| print(f"DEBUG: Downloading weights from {repo_id}...") | |
| weights_path = hf_hub_download(repo_id=repo_id, filename=filename) | |
| except Exception as e: | |
| print(f"HF Download Failed: {e}") | |
| if weights_path: | |
| state_dict = torch.load(weights_path, map_location=self.device) | |
| self.model.load_state_dict(state_dict) | |
| print("DEBUG: Model weights loaded successfully.") | |
| else: | |
| print("WARNING: Model weights NOT FOUND.") | |
| def predict(self, face_img): | |
| if face_img is None or face_img.size == 0: return 0.0 | |
| yuv = cv2.cvtColor(face_img, cv2.COLOR_BGR2YUV) | |
| yuv[:,:,0] = cv2.equalizeHist(yuv[:,:,0]) | |
| face_img = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR) | |
| img_rgb = cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB) | |
| img = cv2.resize(img_rgb, (224, 224)) | |
| img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 | |
| mean = torch.tensor([0.485, 0.456, 0.406], device=self.device).view(3, 1, 1) | |
| std = torch.tensor([0.229, 0.224, 0.225], device=self.device).view(3, 1, 1) | |
| img = (img - mean) / std | |
| img = img.unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| output = self.model(img) | |
| prob = torch.softmax(output, dim=1)[0, 1].item() * 100 | |
| return float(prob) | |
| # --- App Logic --- | |
| try: | |
| face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=True) | |
| except Exception as e: | |
| print(f"DEBUG: FaceMesh initialization failed: {e}") | |
| face_mesh = None | |
| detector = DeepfakeDetector() | |
| def process_media(input_file): | |
| if input_file is None: return "No file uploaded", None, None | |
| is_video = input_file.endswith(('.mp4', '.avi', '.mov', '.mkv')) | |
| if is_video: | |
| cap = cv2.VideoCapture(input_file) | |
| ret, frame = cap.read() | |
| cap.release() | |
| if not ret: return "Video error", None, None | |
| else: | |
| frame = cv2.imread(input_file) | |
| rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
| results = face_mesh.process(rgb_frame) if face_mesh else None | |
| if results and results.multi_face_landmarks: | |
| lms = results.multi_face_landmarks[0].landmark | |
| h, w, _ = frame.shape | |
| pts = np.array([[l.x * w, l.y * h] for l in lms]) | |
| x, y, fw, fh = cv2.boundingRect(pts.astype(np.int32)) | |
| face_crop = frame[max(0, y-10):min(h, y+fh+10), max(0, x-10):min(w, x+fw+10)] | |
| else: | |
| face_crop = frame | |
| risk = detector.predict(face_crop) | |
| verdict = "FAKE/MANIPULATED" if risk > 50 else "REAL/AUTHENTIC" | |
| output_text = f"## Verdict: {verdict}\n### Confidence: {risk:.2f}%" | |
| return output_text, Image.fromarray(cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB)), {"Neural Integrity": risk/100} | |
| # --- Gradio UI --- | |
| with gr.Blocks(theme=gr.themes.Soft(primary_hue="cyan")) as demo: | |
| gr.Markdown("# 🛡️ DeepShield AI: Stable Launch") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_media = gr.File(label="Upload Media") | |
| btn = gr.Button("🔍 Run Analysis", variant="primary") | |
| with gr.Column(): | |
| res_md = gr.Markdown("Analysis results will appear here...") | |
| prev_img = gr.Image(label="Face Analysis") | |
| label_out = gr.Label(label="Signals") | |
| btn.click(process_media, inputs=input_media, outputs=[res_md, prev_img, label_out]) | |
| if __name__ == "__main__": | |
| demo.launch() | |