Implement REST API endpoints and JS event handlers for all XAI/Forensics features

app.py

@@ -2,10 +2,10 @@ import os
 import base64
 from io import BytesIO
 from PIL import Image
+from pydantic import BaseModel
+from fastapi.responses import HTMLResponse
 import gradio as gr
 from gradio import Server
-from gradio.data_classes import FileData
-from fastapi.responses import HTMLResponse
 
 from xai_engine.model import get_model
 from xai_engine.methods import run_all_methods, run_single_method, run_per_head_attention
@@ -18,7 +18,6 @@ app = Server(title="DeepfakeDetection-XAI")
 demo = app
 
 
-
 def b64_to_pil(b64_str: str) -> Image.Image:
     """Helper to convert base64 image data URL to PIL Image."""
     if not b64_str or not isinstance(b64_str, str):
@@ -31,20 +30,64 @@ def b64_to_pil(b64_str: str) -> Image.Image:
     return Image.open(BytesIO(image_data)).convert("RGB")
 
 
-@app.api(name="analyze")
-def analyze_api(image_b64: str, alpha: float = 0.6, cmap: str = "jet") -> dict:
-    """
-    Run complete ViT Deepfake analysis + 6 attribution methods.
-    Returns prediction summary and list of attribution map overlays.
-    """
+# Pydantic Schemas for REST API Endpoints
+class AnalyzeReq(BaseModel):
+    image_b64: str
+    alpha: float = 0.6
+    cmap: str = "jet"
+
+
+class MethodDetailReq(BaseModel):
+    image_b64: str
+    method_name: str = "rollout"
+    steps: int = 20
+    patch_size: int = 32
+    stride: int = 16
+    target_layer_idx: int = 11
+    alpha: float = 0.6
+    cmap: str = "jet"
+
+
+class HeadsReq(BaseModel):
+    image_b64: str
+    layer_idx: int = 11
+    cmap: str = "jet"
+
+
+class ForensicReq(BaseModel):
+    image_b64: str
+    tool_name: str = "ela"
+    quality: int = 75
+    scale: int = 50
+    contrast: int = 20
+    gradient_intensity: int = 90
+    blue_mode: str = "Abs"
+    channel: str = "Luminance"
+    bit: int = 0
+    minmax_radius: int = 2
+
+
+class ReportReq(BaseModel):
+    image_b64: str
+
+
+class SubmitRunReq(BaseModel):
+    image_b64: str
+    ground_truth: str = "unknown"
+    opt_in_image: bool = True
+
+
+# REST API Routes mounted on FastAPI app
+@app.post("/api/analyze")
+def analyze_post(req: AnalyzeReq) -> dict:
     try:
-        image = b64_to_pil(image_b64)
+        image = b64_to_pil(req.image_b64)
     except Exception as e:
         return {"status": "error", "message": f"Image decoding failed: {e}"}
 
     model = get_model()
     pred_data = model.predict(image)
-    results = run_all_methods(model, image, alpha=alpha, cmap=cmap)
+    results = run_all_methods(model, image, alpha=req.alpha, cmap=req.cmap)
 
     res_dicts = [
         {
@@ -65,29 +108,19 @@ def analyze_api(image_b64: str, alpha: float = 0.6, cmap: str = "jet") -> dict:
     }
 
 
-@app.api(name="method_detail")
-def method_detail_api(
-    image_b64: str,
-    method_name: str,
-    steps: int = 20,
-    patch_size: int = 32,
-    stride: int = 16,
-    target_layer_idx: int = 11,
-    alpha: float = 0.6,
-    cmap: str = "jet",
-) -> dict:
-    """Deep-dive analysis for a single specified XAI method with custom parameters."""
+@app.post("/api/method_detail")
+def method_detail_post(req: MethodDetailReq) -> dict:
     try:
-        image = b64_to_pil(image_b64)
+        image = b64_to_pil(req.image_b64)
     except Exception as e:
         return {"status": "error", "message": f"Image decoding failed: {e}"}
 
     model = get_model()
     res = run_single_method(
-        model, image, method_name,
-        alpha=alpha, cmap=cmap, steps=steps,
-        patch_size=patch_size, stride=stride,
-        target_layer_idx=target_layer_idx,
+        model, image, req.method_name,
+        alpha=req.alpha, cmap=req.cmap, steps=req.steps,
+        patch_size=req.patch_size, stride=req.stride,
+        target_layer_idx=req.target_layer_idx,
     )
 
     return {
@@ -103,76 +136,59 @@ def method_detail_api(
     }
 
 
-@app.api(name="heads")
-def heads_api(image_b64: str, layer_idx: int = 11, cmap: str = "jet") -> dict:
-    """Extract and visualize all 6 attention heads for a selected transformer layer."""
+@app.post("/api/heads")
+def heads_post(req: HeadsReq) -> dict:
     try:
-        image = b64_to_pil(image_b64)
+        image = b64_to_pil(req.image_b64)
     except Exception as e:
         return {"status": "error", "message": f"Image decoding failed: {e}"}
 
     model = get_model()
-    head_data = run_per_head_attention(model, image, layer_idx=layer_idx, cmap=cmap)
+    head_data = run_per_head_attention(model, image, layer_idx=req.layer_idx, cmap=req.cmap)
     head_data["status"] = "success"
     return head_data
 
 
-@app.api(name="forensic")
-def forensic_api(
-    image_b64: str,
-    tool_name: str,
-    quality: int = 75,
-    scale: int = 50,
-    contrast: int = 20,
-    gradient_intensity: int = 90,
-    blue_mode: str = "Abs",
-    channel: str = "Luminance",
-    bit: int = 0,
-    minmax_radius: int = 2,
-) -> dict:
-    """Run classic digital forensic analysis (ELA, Gradient, BitPlane, MinMax, Wavelet)."""
+@app.post("/api/forensic")
+def forensic_post(req: ForensicReq) -> dict:
     try:
-        image = b64_to_pil(image_b64)
+        image = b64_to_pil(req.image_b64)
     except Exception as e:
         return {"status": "error", "message": f"Image decoding failed: {e}"}
 
     res = run_forensic_tool(
-        image, tool_name,
-        quality=quality, scale=scale, contrast=contrast,
-        gradient_intensity=gradient_intensity, blue_mode=blue_mode,
-        channel=channel, bit=bit, minmax_radius=minmax_radius,
+        image, req.tool_name,
+        quality=req.quality, scale=req.scale, contrast=req.contrast,
+        gradient_intensity=req.gradient_intensity, blue_mode=req.blue_mode,
+        channel=req.channel, bit=req.bit, minmax_radius=req.minmax_radius,
     )
     res["status"] = "success"
     return res
 
 
-@app.api(name="report")
-def report_api(image_b64: str) -> FileData:
-    """Generate and return downloadable multi-panel forensic report PNG."""
-    image = b64_to_pil(image_b64)
+@app.post("/api/report")
+def report_post(req: ReportReq) -> dict:
+    image = b64_to_pil(req.image_b64)
     model = get_model()
     pred_data = model.predict(image)
     results = run_all_methods(model, image)
     report_path = generate_forensic_report(image, pred_data, results)
-    return FileData(path=report_path)
+    with open(report_path, "rb") as f:
+        rep_b64 = "data:image/png;base64," + base64.b64encode(f.read()).decode("utf-8")
+    return {"status": "success", "report_b64": rep_b64}
 
 
-@app.api(name="submit_run")
-def submit_run_api(
-    image_b64: str,
-    ground_truth: str = "unknown",
-    opt_in_image: bool = True,
-) -> dict:
-    """Log analysis run, prediction, ground truth feedback, and XAI composite heatmaps to dataset."""
+@app.post("/api/submit_run")
+def submit_run_post(req: SubmitRunReq) -> dict:
     try:
-        image = b64_to_pil(image_b64)
+        image = b64_to_pil(req.image_b64)
     except Exception as e:
         return {"status": "error", "message": f"Image decoding failed: {e}"}
 
     model = get_model()
     pred_data = model.predict(image)
     results = run_all_methods(model, image)
-    run_id = log_run_dataset(image, pred_data, results, ground_truth=ground_truth, opt_in_image=opt_in_image)
+    run_id = log_run_dataset(image, pred_data, results, ground_truth=req.ground_truth, opt_in_image=req.opt_in_image)
 
     return {
         "status": "success",
@@ -212,5 +228,8 @@ async def homepage():
         return f.read()
 
 
+app.setup()
+
+
 if __name__ == "__main__":
     app.launch(show_error=True)

index.html

@@ -422,6 +422,12 @@
       document.getElementById('tab-' + tabId).classList.add('border-teal-400', 'text-teal-400');
       document.getElementById('tab-' + tabId).classList.remove('border-transparent', 'text-slate-400');
       document.getElementById('section-' + tabId).classList.remove('hidden');
+
+      if (currentImageB64) {
+        if (tabId === 'explorer') runExplorerMethod();
+        if (tabId === 'heads') runPerHeadInspector();
+        if (tabId === 'forensics') runForensicAnalysis();
+      }
     }
 
     function handleFileSelect(evt) {
@@ -448,13 +454,12 @@
     }
 
     async function apiCall(endpoint, payload) {
-      const resp = await fetch('/gradio_api/' + endpoint, {
+      const resp = await fetch('/api/' + endpoint, {
         method: 'POST',
         headers: { 'Content-Type': 'application/json' },
-        body: JSON.stringify({ data: [payload] })
+        body: JSON.stringify(payload)
       });
-      const data = await resp.json();
-      return data.data ? data.data[0] : data;
+      return await resp.json();
     }
 
     async function runAnalysis() {
@@ -464,11 +469,9 @@
       document.getElementById('btn-analyze').disabled = true;
 
       try {
-        const resp = await fetch('/api/predict_all', { method: 'POST' }); // Fallback direct fetch
-        const res = await apiCall('analyze', currentImageB64);
+        const res = await apiCall('analyze', { image_b64: currentImageB64, alpha: 0.6, cmap: 'jet' });
         
         if (res.status === 'success') {
-          // Update Prediction Badge
           const pred = res.prediction;
           const isFake = pred.prediction === 'FAKE';
           const badge = document.getElementById('prediction-badge');
@@ -478,7 +481,6 @@
           document.getElementById('confidence-val').innerText = pred.confidence_pct.toFixed(2) + '%';
           document.getElementById('confidence-box').classList.remove('hidden');
 
-          // Render Grid
           const grid = document.getElementById('results-grid');
           grid.innerHTML = '';
           res.results.forEach(r => {
@@ -498,24 +500,144 @@
 
           document.getElementById('submit-panel').classList.remove('hidden');
           document.getElementById('btn-report').disabled = false;
+        } else {
+          alert('Analysis error: ' + (res.message || 'Unknown error'));
         }
       } catch (err) {
         console.error(err);
+        alert('API Request Failed: ' + err);
       } finally {
         document.getElementById('loading-spinner').classList.add('hidden');
         document.getElementById('btn-analyze').disabled = false;
       }
     }
 
+    async function runExplorerMethod() {
+      if (!currentImageB64) return;
+      const method = document.getElementById('exp-method').value;
+      const cmap = document.getElementById('exp-cmap').value;
+      const alpha = parseFloat(document.getElementById('exp-alpha').value);
+      const steps = parseInt(document.getElementById('exp-steps').value);
+      const patch_size = parseInt(document.getElementById('exp-patch').value);
+      const target_layer_idx = parseInt(document.getElementById('exp-layer').value);
+
+      const ctrlIG = document.getElementById('ctrl-integrated');
+      if (ctrlIG) ctrlIG.classList.toggle('hidden', method !== 'integrated');
+      const ctrlOcc = document.getElementById('ctrl-occlusion');
+      if (ctrlOcc) ctrlOcc.classList.toggle('hidden', method !== 'occlusion');
+      const ctrlGC = document.getElementById('ctrl-gradcam');
+      if (ctrlGC) ctrlGC.classList.toggle('hidden', method !== 'gradcam');
+
+      const res = await apiCall('method_detail', {
+        image_b64: currentImageB64,
+        method_name: method,
+        steps: steps,
+        patch_size: patch_size,
+        target_layer_idx: target_layer_idx,
+        alpha: alpha,
+        cmap: cmap
+      });
+
+      if (res.status === 'success' && res.result) {
+        document.getElementById('exp-img-overlay').src = res.result.overlay_b64;
+        document.getElementById('exp-img-raw').src = res.result.raw_heatmap_b64;
+      }
+    }
+
+    async function runPerHeadInspector() {
+      if (!currentImageB64) return;
+      const layer_idx = parseInt(document.getElementById('head-layer-input').value);
+      document.getElementById('head-layer-val').innerText = layer_idx;
+
+      const res = await apiCall('heads', {
+        image_b64: currentImageB64,
+        layer_idx: layer_idx,
+        cmap: 'jet'
+      });
+
+      if (res.status === 'success' && res.head_overlays) {
+        const grid = document.getElementById('heads-grid');
+        grid.innerHTML = '';
+        res.head_overlays.forEach((b64, idx) => {
+          const card = document.createElement('div');
+          card.className = 'glass-panel rounded-xl p-2 text-center hover:border-teal-500/50 transition cursor-pointer';
+          card.onclick = () => openModal(`Layer ${layer_idx} — Head ${idx + 1}`, b64);
+          card.innerHTML = `
+            <img src="${b64}" class="w-full h-24 object-cover rounded-lg mb-1" />
+            <span class="text-[10px] font-mono text-slate-300">Head ${idx + 1}</span>
+          `;
+          grid.appendChild(card);
+        });
+      }
+    }
+
+    async function runForensicAnalysis() {
+      if (!currentImageB64) return;
+      const tool = document.getElementById('for-tool').value;
+      const quality = parseInt(document.getElementById('for-quality').value);
+      const bit = parseInt(document.getElementById('for-bit').value);
+
+      const ctrlELA = document.getElementById('ctrl-ela');
+      if (ctrlELA) ctrlELA.classList.toggle('hidden', tool !== 'ela');
+      const ctrlBit = document.getElementById('ctrl-bitplane');
+      if (ctrlBit) ctrlBit.classList.toggle('hidden', tool !== 'bitplane');
+
+      const res = await apiCall('forensic', {
+        image_b64: currentImageB64,
+        tool_name: tool,
+        quality: quality,
+        bit: bit
+      });
+
+      if (res.status === 'success' && res.result_b64) {
+        const img = document.getElementById('for-img-output');
+        img.src = res.result_b64;
+        img.classList.remove('hidden');
+        document.getElementById('for-placeholder').classList.add('hidden');
+        document.getElementById('for-info').innerText = `${res.tool_name} analysis complete (${res.compute_time_ms} ms).`;
+      }
+    }
+
+    async function generateReport() {
+      if (!currentImageB64) return;
+      const btn = document.getElementById('btn-report');
+      btn.disabled = true;
+      btn.innerHTML = '<i class="fa-solid fa-spinner fa-spin"></i> Generating...';
+
+      try {
+        const res = await apiCall('report', { image_b64: currentImageB64 });
+        if (res.status === 'success' && res.report_b64) {
+          openModal('Forensic Report', res.report_b64);
+          const a = document.createElement('a');
+          a.href = res.report_b64;
+          a.download = `forensic_report_${Date.now()}.png`;
+          a.click();
+        }
+      } catch (err) {
+        console.error(err);
+      } finally {
+        btn.disabled = false;
+        btn.innerHTML = '<i class="fa-solid fa-file-pdf"></i> Export Forensic Report';
+      }
+    }
+
     async function submitDatasetRun() {
       if (!currentImageB64) return;
-      const gt = document.querySelector('input[name="ground_truth"]:checked').value;
+      const gtEl = document.querySelector('input[name="ground_truth"]:checked');
+      const gt = gtEl ? gtEl.value : 'unknown';
       const optIn = document.getElementById('opt-in-image').checked;
 
-      const res = await apiCall('submit_run', currentImageB64);
-      const statusP = document.getElementById('submit-status');
-      statusP.innerText = `Run ${res.run_id} logged to dataset buildborderless/deepfake-explainability!`;
-      statusP.classList.remove('hidden');
+      const res = await apiCall('submit_run', {
+        image_b64: currentImageB64,
+        ground_truth: gt,
+        opt_in_image: optIn
+      });
+
+      if (res.status === 'success') {
+        const statusP = document.getElementById('submit-status');
+        statusP.innerText = `Run ${res.run_id} logged to dataset buildborderless/deepfake-explainability!`;
+        statusP.classList.remove('hidden');
+      }
     }
 
     function openModal(title, b64) {
@@ -530,3 +652,4 @@
   </script>
 </body>
 </html>
+