Update index.html

index.html

@@ -35,14 +35,14 @@ small{color:#93a3b1}
 </div>
 </main>
 <script type="module">
-import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3';
+import { AutoProcessor, AutoModelForImageClassification, RawImage } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3';
 const $ = id => document.getElementById(id);
 const MODEL = 'buildborderless/CommunityForensics-DeepfakeDet-ViT';
-let clf;
+let processor, model;
 async function load(){
-  if(clf) return clf;
-  clf = await pipeline('image-classification', MODEL, { dtype: 'fp32' });
-  return clf;
+  if(model) return;
+  processor = await AutoProcessor.from_pretrained(MODEL);
+  model = await AutoModelForImageClassification.from_pretrained(MODEL, { dtype: 'fp32' });
 }
 const TOOLS = [[/midjourney/i,'Midjourney'],[/dall[\s\-·]?e|openai/i,'OpenAI / DALL·E'],[/stable\s?diffusion|sdxl|automatic1111/i,'Stable Diffusion'],[/firefly|generative ?fill/i,'Adobe Firefly'],[/imagen|gemini|synthid/i,'Google AI'],[/leonardo\.?ai/i,'Leonardo AI'],[/comfyui/i,'ComfyUI'],[/ideogram/i,'Ideogram']];
 function markers(bytes){
@@ -62,20 +62,23 @@ $('f').onchange = async e => {
   $('pct').textContent = ''; $('b').style.width = '0'; $('m').innerHTML = ''; $('raw').textContent = '';
   try {
     const mk = markers(new Uint8Array(await file.arrayBuffer()));
-    const c = await load();
+    await load();
     $('v').textContent = 'Checking...';
-    const r = await c(url, { top_k: 5 });
-    const hit = r.find(x => /fake|ai|synth|generat|artificial/i.test(x.label));
-    let ai = hit ? hit.score : null;
-    if(ai === null && r.length === 2){ ai = r[0].score < 0.5 ? r[0].score : 1 - r[1].score; }
-    $('raw').textContent = 'Model output: ' + r.map(x => x.label + ' ' + Math.round(x.score*100) + '%').join(', ');
+    const image = await RawImage.read(url);
+    const inputs = await processor(image);
+    const { logits } = await model(inputs);
+    const z = Array.from(logits.data);
+    let ai;
+    if(z.length === 1){ ai = 1 / (1 + Math.exp(-z[0])); }
+    else { const m = Math.max(...z), ex = z.map(v => Math.exp(v - m)), s = ex.reduce((a,b) => a + b, 0); ai = ex[1] / s; }
+    $('raw').textContent = 'Raw model score: ' + z.map(v => v.toFixed(2)).join(', ');
     mk.forEach(t => { const li = document.createElement('li'); li.textContent = t; $('m').appendChild(li); });
     if(mk.length) $('v').textContent = 'AI markers found';
-    else if(ai === null) $('v').textContent = 'Result below';
     else if(ai >= 0.85) $('v').textContent = 'Likely AI-generated';
     else if(ai >= 0.5) $('v').textContent = 'Possibly AI';
     else $('v').textContent = 'Likely a real photo';
-    if(ai !== null){ $('pct').textContent = 'Chance it is AI: ' + Math.round(ai*100) + '%'; $('b').style.width = Math.round(ai*100) + '%'; }
+    $('pct').textContent = 'Chance it is AI: ' + Math.round(ai*100) + '%';
+    $('b').style.width = Math.round(ai*100) + '%';
   } catch(err){
     $('v').textContent = 'Could not analyze this image.';
     $('raw').textContent = String(err);