Update index.html

index.html

@@ -77,6 +77,7 @@
       // ===== 1) Config =====
       const MODEL_ID = "onnx-community/ijepa_vith14_1k"; // <-- I-JEPA ViT-H/14, ImageNet-1k
       const EXAMPLE_IMAGE_URL = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png";
+      const SUPPORTED_RESOLUTIONS = [224, 336, 448]; // Specific resolutions I-JEPA ONNX might expect
 
       // DOM
       const imageLoader = document.getElementById("imageLoader");
@@ -104,6 +105,7 @@
       let animationFrameId = null;
       let lastMouseEvent = null;
       let maxPixels = null;
+      let imageCropParams = null; // To store crop parameters for consistent redraws
 
       function updateStatus(text, isLoading=false){
         statusText.textContent = text;
@@ -124,9 +126,7 @@
 
         try{
           extractor = await pipeline("image-feature-extraction", MODEL_ID, { device, dtype });
-          // Try to fetch config-provided patch size if present
           patchSize = extractor?.model?.config?.patch_size ?? patchSize;
-          // Avoid internal resizes — we control canvas dims
           if (extractor?.processor?.image_processor) extractor.processor.image_processor.do_resize = false;
           updateStatus("Ready. Please select an image.");
         }catch(e){
@@ -154,15 +154,9 @@
           const res = await fetch(EXAMPLE_IMAGE_URL);
           const blob = await res.blob();
           loadImageOntoCanvas(URL.createObjectURL(blob));
-        }catch(e){
-          console.error(e);
-          updateStatus("Failed to load example image.");
-        }
-      }
-      function handleImageUpload(e){
-        const f = e.target.files?.[0];
-        if (f) loadImageOntoCanvas(URL.createObjectURL(f));
+        }catch(e){ console.error(e); updateStatus("Failed to load example image."); }
       }
+      function handleImageUpload(e){ if (e.target.files?.[0]) loadImageOntoCanvas(URL.createObjectURL(e.target.files[0])); }
       function handleDragOver(e){ e.preventDefault(); dropZone.classList.add("border-blue-500","bg-gray-800"); }
       function handleDragLeave(e){ e.preventDefault(); dropZone.classList.remove("border-blue-500","bg-gray-800"); }
       function handleDrop(e){
@@ -180,6 +174,21 @@
       function handleSliderInput(e){ imageScale = parseFloat(e.target.value); scaleValue.textContent = `${imageScale.toFixed(2)}x`; }
       function handleSliderChange(){ if (currentImageUrl) loadImageOntoCanvas(currentImageUrl); }
 
+      function findClosestSupportedResolution(targetDim) {
+          return SUPPORTED_RESOLUTIONS.reduce((prev, curr) =>
+              Math.abs(curr - targetDim) < Math.abs(prev - targetDim) ? curr : prev
+          );
+      }
+      
+      function redrawOriginalImage() {
+        if (!originalImage || !imageCropParams) return;
+        ctx.drawImage(
+            originalImage,
+            imageCropParams.sx, imageCropParams.sy, imageCropParams.sWidth, imageCropParams.sHeight,
+            0, 0, imageCanvas.width, imageCanvas.height
+        );
+      }
+
       function loadImageOntoCanvas(url){
         currentImageUrl = url;
         originalImage = new Image();
@@ -187,63 +196,68 @@
           if (!patchSize){ updateStatus("Error: patch size unknown."); return; }
           canvasPlaceholder.style.display = "none";
           imageCanvas.style.display = "block";
-
-          let newW = originalImage.naturalWidth * imageScale;
-          let newH = originalImage.naturalHeight * imageScale;
-
-          const px = newW * newH;
-          if (px > maxPixels){
-            const r = Math.sqrt(maxPixels / px);
-            newW *= r; newH *= r;
+          
+          const { naturalWidth: w, naturalHeight: h } = originalImage;
+          
+          // --- NEW CENTER-CROP LOGIC ---
+          const cropSize = Math.min(w, h);
+          const sx = (w - cropSize) / 2;
+          const sy = (h - cropSize) / 2;
+          imageCropParams = { sx, sy, sWidth: cropSize, sHeight: cropSize };
+          
+          // Determine target canvas resolution
+          let scaledCropSize = cropSize * imageScale;
+          if (scaledCropSize * scaledCropSize > maxPixels) {
+              scaledCropSize = Math.sqrt(maxPixels);
           }
-
-          const croppedW = Math.floor(newW / patchSize) * patchSize;
-          const croppedH = Math.floor(newH / patchSize) * patchSize;
-          if (croppedW < patchSize || croppedH < patchSize){
-            updateStatus("Scaled image is too small to process.");
-            imageCanvas.style.display = "none";
-            canvasPlaceholder.style.display = "block";
-            canvasPlaceholder.textContent = "Scaled image is too small.";
-            return;
+          let chosenResolution = findClosestSupportedResolution(scaledCropSize);
+          
+          // Ensure chosen resolution is at least one patch size
+          if (chosenResolution < patchSize) {
+              updateStatus("Scaled image is too small to process.");
+              imageCanvas.style.display = "none";
+              canvasPlaceholder.style.display = "block";
+              canvasPlaceholder.textContent = "Scaled image is too small.";
+              return;
           }
+          
+          imageCanvas.width = chosenResolution;
+          imageCanvas.height = chosenResolution;
 
-          imageCanvas.width = croppedW;
-          imageCanvas.height = croppedH;
-          ctx.drawImage(originalImage, 0, 0, croppedW, croppedH);
-          await processImage();
+          redrawOriginalImage(); // Initial draw with center-crop and resize
+          
+          await processImage(chosenResolution);
           setTimeout(() => { canvasContainer.scrollIntoView({ behavior: "smooth", block: "center" }); }, 100);
         };
         originalImage.onerror = () => { updateStatus("Failed to load the selected image."); canvasPlaceholder.style.display = "block"; imageCanvas.style.display = "none"; };
         originalImage.src = url;
       }
 
-      async function processImage(){
+      async function processImage(chosenResolution){
         if (!extractor) return;
         updateStatus("Analyzing with I‑JEPA... 🧠", true);
         similarityScores = null; lastHoverData = null;
         try{
           const imageData = await RawImage.fromCanvas(imageCanvas);
-          // No pooling: we want per-token outputs
           const features = await extractor(imageData, { pooling: "none" }); // [1, T, D]
 
-          // Compute how many tokens are patches vs special tokens robustly.
           const totalTokens = features.dims?.[1] ?? features.shape?.[1] ?? features.size?.[1];
           const nPatches = (imageCanvas.width / patchSize) * (imageCanvas.height / patchSize);
           const specialTokens = Math.max(0, totalTokens - nPatches);
 
-          const patchFeatures = features.slice(null, [specialTokens, nPatches]); // [1, nPatches, D]
+          const patchFeatures = features.slice(null, [specialTokens, nPatches]);
           const normalized = patchFeatures.normalize(2, -1);
-          const sims = await matmul(normalized, normalized.permute(0,2,1)); // [1, nPatches, nPatches]
+          const sims = await matmul(normalized, normalized.permute(0,2,1));
           similarityScores = (await sims.tolist())[0];
 
-          updateStatus(`Image processed (${imageCanvas.width}×${imageCanvas.height}). Hover to explore. ✨`);
+          updateStatus(`Image processed at ${chosenResolution}×${chosenResolution}. Hover to explore. ✨`);
         }catch(err){
           console.error("I‑JEPA processing error:", err);
-          updateStatus("An error occurred during image processing.");
+          updateStatus("An error occurred during processing. The image size might be unsupported.");
         }
       }
 
-      function handleTouchMove(e){ e.preventDefault(); if (e.touches.length>0) handleMouseMove(e.touches[0]); }
+      function handleTouchMove(e){ e.preventDefault(); if (e.touches.length > 0) handleMouseMove(e.touches[0]); }
       function handleMouseMove(e){ lastMouseEvent = e; if (!animationFrameId) animationFrameId = requestAnimationFrame(drawLoop); }
 
       function drawLoop(){
@@ -253,7 +267,7 @@
         const scaleY = imageCanvas.height / rect.height;
         const x = (lastMouseEvent.clientX - rect.left) * scaleX;
         const y = (lastMouseEvent.clientY - rect.top) * scaleY;
-        if (x<0||x>=imageCanvas.width||y<0||y>=imageCanvas.height){ animationFrameId = null; return; }
+        if (x < 0 || x >= imageCanvas.width || y < 0 || y >= imageCanvas.height){ animationFrameId = null; return; }
 
         const patchesPerRow = imageCanvas.width / patchSize;
         const patchX = Math.floor(x / patchSize);
@@ -267,36 +281,26 @@
         animationFrameId = null;
       }
 
-      const INFERNO_COLORMAP = [
-        [0.0,[0,0,4]],[0.1,[39,12,69]],[0.2,[84,15,104]],[0.3,[128,31,103]],[0.4,[170,48,88]],
-        [0.5,[209,70,68]],[0.6,[240,97,47]],[0.7,[253,138,28]],[0.8,[252,185,26]],[0.9,[240,231,56]],[1.0,[252,255,160]]
-      ];
-      function getInfernoColor(t){
-        for (let i=1;i<INFERNO_COLORMAP.length;i++){
-          const [tp,cp]=INFERNO_COLORMAP[i-1]; const [tc,cc]=INFERNO_COLORMAP[i];
-          if (t<=tc){ const a=(t-tp)/(tc-tp); const r=cp[0]+a*(cc[0]-cp[0]); const g=cp[1]+a*(cc[1]-cp[1]); const b=cp[2]+a*(cc[2]-cp[2]); return `rgb(${Math.round(r)}, ${Math.round(g)}, ${Math.round(b)})`; }
-        }
-        const last=INFERNO_COLORMAP[INFERNO_COLORMAP.length-1][1];
-        return `rgb(${last.join(",")})`;
-      }
+      const INFERNO_COLORMAP = [ [0.0,[0,0,4]],[0.1,[39,12,69]],[0.2,[84,15,104]],[0.3,[128,31,103]],[0.4,[170,48,88]], [0.5,[209,70,68]],[0.6,[240,97,47]],[0.7,[253,138,28]],[0.8,[252,185,26]],[0.9,[240,231,56]],[1.0,[252,255,160]] ];
+      function getInfernoColor(t){ for (let i=1;i<INFERNO_COLORMAP.length;i++){ const [tp,cp]=INFERNO_COLORMAP[i-1]; const [tc,cc]=INFERNO_COLORMAP[i]; if (t<=tc){ const a=(t-tp)/(tc-tp); const r=cp[0]+a*(cc[0]-cp[0]); const g=cp[1]+a*(cc[1]-cp[1]); const b=cp[2]+a*(cc[2]-cp[2]); return `rgb(${Math.round(r)}, ${Math.round(g)}, ${Math.round(b)})`; } } const last=INFERNO_COLORMAP[INFERNO_COLORMAP.length-1][1]; return `rgb(${last.join(",")})`; }
 
       function drawHighlights(queryIndex, allPatches){
         const patchesPerRow = imageCanvas.width / patchSize;
         if (isOverlayMode){
-          ctx.drawImage(originalImage, 0, 0, imageCanvas.width, imageCanvas.height);
+          redrawOriginalImage();
           ctx.fillStyle = "rgba(0,0,0,0.6)"; ctx.fillRect(0,0,imageCanvas.width,imageCanvas.height);
         } else {
           ctx.fillStyle = getInfernoColor(0); ctx.fillRect(0,0,imageCanvas.width,imageCanvas.height);
         }
-        if (allPatches.length>0){
-          const scores = allPatches.map(p=>p.score);
+        if (allPatches.length > 0){
+          const scores = allPatches.map(p => p.score);
           const minS = Math.min(...scores); const maxS = Math.max(...scores); const rng = maxS - minS;
           for (const p of allPatches){
             if (p.index === queryIndex) continue;
-            const t = rng > 1e-4 ? (p.score - minS)/rng : 1;
+            const t = rng > 1e-4 ? (p.score - minS) / rng : 1;
             const py = Math.floor(p.index / patchesPerRow);
             const px = p.index % patchesPerRow;
-            if (isOverlayMode){ const a = Math.pow(t,2)*0.8; ctx.fillStyle = `rgba(255,255,255,${a})`; }
+            if (isOverlayMode){ const a = Math.pow(t, 2) * 0.8; ctx.fillStyle = `rgba(255,255,255,${a})`; }
             else { ctx.fillStyle = getInfernoColor(t); }
             ctx.fillRect(px*patchSize, py*patchSize, patchSize, patchSize);
           }
@@ -310,7 +314,7 @@
       function clearHighlights(){
         if (animationFrameId){ cancelAnimationFrame(animationFrameId); animationFrameId = null; }
         lastMouseEvent = null; lastHoverData = null;
-        if (originalImage){ ctx.drawImage(originalImage, 0, 0, imageCanvas.width, imageCanvas.height); }
+        if (originalImage) redrawOriginalImage();
       }
 
       initialize();