Spaces:
Runtime error
Runtime error
| import importlib.util | |
| import os | |
| import sys | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from safetensors.torch import load_file | |
| from transformers import ViTImageProcessor | |
| # ----------------------------- | |
| # CONFIGURATION | |
| # ----------------------------- | |
| MODEL_REPO = "SARVM/Refined_ViT" | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print("=" * 60) | |
| print("FORESIGHT - MODEL INITIALIZATION") | |
| print("=" * 60) | |
| processor = ViTImageProcessor( | |
| do_resize=True, | |
| size={"height": 224, "width": 224}, | |
| do_rescale=True, | |
| rescale_factor=1 / 255, | |
| do_normalize=True, | |
| image_mean=[0.485, 0.456, 0.406], | |
| image_std=[0.229, 0.224, 0.225], | |
| ) | |
| config_path = hf_hub_download(repo_id=MODEL_REPO, filename="configuration_vit_refiner.py", token=HF_TOKEN) | |
| model_path = hf_hub_download(repo_id=MODEL_REPO, filename="modeling_vit_refiner.py", token=HF_TOKEN) | |
| config_spec = importlib.util.spec_from_file_location("configuration_vit_refiner", config_path) | |
| config_module = importlib.util.module_from_spec(config_spec) | |
| sys.modules["configuration_vit_refiner"] = config_module | |
| config_spec.loader.exec_module(config_module) | |
| model_spec = importlib.util.spec_from_file_location("modeling_vit_refiner", model_path) | |
| model_module = importlib.util.module_from_spec(model_spec) | |
| sys.modules["modeling_vit_refiner"] = model_module | |
| model_spec.loader.exec_module(model_module) | |
| ViTRefinerConfig = config_module.ViTRefinerConfig | |
| ViTRefinerModel = model_module.ViTRefinerModel | |
| config = ViTRefinerConfig.from_pretrained(MODEL_REPO, token=HF_TOKEN) | |
| model = ViTRefinerModel(config) | |
| # Force 'eager' attention mode to enable explicit attention weights output | |
| if hasattr(model, "vit"): | |
| if hasattr(model.vit, "config"): | |
| model.vit.config._attn_implementation = "eager" | |
| model.vit.config.output_attentions = True | |
| if hasattr(model.vit, "vit") and hasattr(model.vit.vit, "config"): | |
| model.vit.vit.config._attn_implementation = "eager" | |
| model.vit.vit.config.output_attentions = True | |
| weights_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.safetensors", token=HF_TOKEN) | |
| state_dict = load_file(weights_path) | |
| model.load_state_dict(state_dict, strict=False) | |
| model.to(DEVICE) | |
| model.eval() | |
| model.config.id2label = {0: "FAKE", 1: "REAL"} | |
| model.config.label2id = {"FAKE": 0, "REAL": 1} | |
| # ----------------------------- | |
| # ATTENTION HOOK EXTRACTION | |
| # ----------------------------- | |
| captured_attentions = [] | |
| def attention_hook(module, input, output): | |
| if isinstance(output, tuple): | |
| if len(output) > 1 and output[1] is not None: | |
| captured_attentions.append(output[1].detach().cpu()) | |
| elif isinstance(output[0], torch.Tensor) and output[0].ndim == 4: | |
| captured_attentions.append(output[0].detach().cpu()) | |
| # Register forward hooks on all ViTSelfAttention blocks | |
| for module in model.modules(): | |
| if module.__class__.__name__ == "ViTSelfAttention": | |
| module.register_forward_hook(attention_hook) | |
| # ----------------------------- | |
| # ATTENTION ROLLOUT | |
| # ----------------------------- | |
| def compute_attention_rollout(attentions): | |
| if not attentions: | |
| return None | |
| att_mat = torch.stack(attentions) | |
| if att_mat.dim() == 5: | |
| att_mat = att_mat[:, 0] | |
| att_mat = att_mat.mean(dim=1) | |
| residual_att = torch.eye(att_mat.size(-1), device=att_mat.device) | |
| aug_att_mat = att_mat + residual_att | |
| aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1, keepdim=True) | |
| joint_attentions = torch.zeros_like(aug_att_mat) | |
| joint_attentions[0] = aug_att_mat[0] | |
| for n in range(1, aug_att_mat.size(0)): | |
| joint_attentions[n] = aug_att_mat[n] @ joint_attentions[n - 1] | |
| return joint_attentions[-1] | |
| # ----------------------------- | |
| # PREDICTION FUNCTION | |
| # ----------------------------- | |
| def predict(image): | |
| global captured_attentions | |
| if image is None: | |
| return None, None, None | |
| if not isinstance(image, Image.Image): | |
| image = Image.fromarray(image) | |
| image = image.convert("RGB") | |
| inputs = processor(images=image, return_tensors="pt") | |
| inputs = {key: value.to(DEVICE) for key, value in inputs.items()} | |
| captured_attentions.clear() | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits | |
| probs = torch.nn.functional.softmax(logits, dim=-1) | |
| confidence, predicted_class_idx = torch.max(probs, dim=-1) | |
| prediction = model.config.id2label[predicted_class_idx.item()] | |
| confidence_pct = round(confidence.item() * 100, 2) | |
| rollout = compute_attention_rollout(captured_attentions) | |
| if rollout is not None: | |
| mask = rollout[0, 1:] | |
| size = int(mask.shape[0] ** 0.5) | |
| mask = mask.reshape(size, size).numpy() | |
| mask = cv2.resize(mask, image.size) | |
| mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8) | |
| heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET) | |
| heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) | |
| overlay = cv2.addWeighted(np.array(image), 0.6, heatmap, 0.4, 0) | |
| else: | |
| print("Warning: No attention matrices captured.") | |
| overlay = np.array(image) | |
| return prediction, f"{confidence_pct}%", overlay | |
| # ----------------------------- | |
| # UI DESIGN | |
| # ----------------------------- | |
| custom_css = """ | |
| :root { --primary-blue: #2563eb; --hero-text: #0f172a; } | |
| .dark { --hero-text: #f8fafc; } | |
| body { background-color: var(--background-fill-primary); } | |
| .hero { text-align: center; font-family: 'Inter', sans-serif; font-size: 48px; font-weight: 800; letter-spacing: -0.04em; margin-top: 50px; color: var(--hero-text) !important; } | |
| .sub { text-align: center; opacity: 0.7; font-size: 14px; font-weight: 600; letter-spacing: 0.1em; text-transform: uppercase; margin-bottom: 40px; color: var(--body-text-color); } | |
| .glass { background: var(--block-background-fill) !important; border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 24px !important; box-shadow: var(--block-shadow); transition: all 0.2s ease; } | |
| .glass:hover { border-color: var(--primary-blue) !important; box-shadow: 0 4px 20px rgba(37, 99, 235, 0.1); } | |
| button.primary { background: var(--primary-blue) !important; color: white !important; border: none !important; font-weight: 600 !important; padding: 12px 24px !important; border-radius: 8px !important; box-shadow: 0 4px 12px rgba(37, 99, 235, 0.2) !important; } | |
| button.primary:hover { background: #1d4ed8 !important; transform: translateY(-1px); box-shadow: 0 6px 16px rgba(37, 99, 235, 0.3) !important; } | |
| .gr-label { font-weight: 600 !important; font-size: 12px !important; text-transform: uppercase; color: var(--primary-blue) !important; } | |
| """ | |
| with gr.Blocks() as demo: | |
| gr.Markdown("<div class='hero'>FORESIGHT<span style='color:#3b82f6'>.</span></div>") | |
| gr.Markdown("<div class='sub'>Deep Intelligence Neural Analysis</div>") | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Group(elem_classes="glass"): | |
| input_image = gr.Image(type="pil", label="Source Input") | |
| run_btn = gr.Button("RUN DIAGNOSTIC", variant="primary") | |
| with gr.Column(): | |
| with gr.Group(elem_classes="glass"): | |
| output_label = gr.Label(label="Classification Verdict") | |
| output_conf = gr.Textbox(label="Confidence Rating", interactive=False) | |
| heatmap_output = gr.Image(label="Vulnerability Visualization") | |
| run_btn.click( | |
| fn=predict, | |
| inputs=input_image, | |
| outputs=[output_label, output_conf, heatmap_output] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.getenv("PORT", "7860")), | |
| css=custom_css, | |
| theme=gr.themes.Soft( | |
| primary_hue="blue", | |
| font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"] | |
| ) | |
| ) |