feat: hilbert curves

Author: AshishKumar4

dataset_preprocessing.ipynb

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+"""
+Demo script for visualizing Hilbert curve patching in Vision Transformers.
+
+This script demonstrates:
+1. How a Hilbert curve maps through an image grid
+2. How patching/unpatching with Hilbert ordering works
+3. Visual comparison between row-major and Hilbert curve patch ordering
+
+Usage:
+    python demo_hilbert_curve.py [--image IMAGE_PATH] [--patch_size PATCH_SIZE]
+
+Options:
+    --image: Path to an image file (default: will use a sample image)
+    --patch_size: Size of patches (default: 16)
+"""
+import os
+os.environ["JAX_PLATFORMS"] = "cpu"
+import argparse
+import numpy as np
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+from PIL import Image
+import requests
+from io import BytesIO
+import os
+import cv2
+from flaxdiff.models.hilbert import (
+    visualize_hilbert_curve, 
+    demo_hilbert_patching,
+    hilbert_patchify,
+    hilbert_unpatchify,
+    hilbert_indices,
+    inverse_permutation,
+    patchify
+)
+
+def load_sample_image():
+    """Load a sample image if no image path is provided."""
+    print("Downloading a sample image...")
+    # Use a relatively small but detailed image
+    url = 'https://www.caledoniaplay.com/wp-content/uploads/2016/01/EDU-PRODUCT-DESCRIPTION-gallery-image-OUTDOOR-SEATING-RUSTIC-LOG-BENCH-1-555x462.jpg'
+    response = requests.get(url)
+    img = Image.open(BytesIO(response.content))
+    return np.array(img) / 255.0  # Normalize to [0, 1]
+
+def load_image(path):
+    """Load an image from the given path."""
+    img = Image.open(path)
+    # Convert to RGB if needed
+    if img.mode != 'RGB':
+        img = img.convert('RGB')
+    # Resize to ensure dimensions are divisible by patch_size
+    w, h = img.size
+    print(f"Loaded image of size: {img.size}")
+    img = np.array(img) / 255.0  # Normalize to [0, 1]
+    return img
+
+def main():
+    parser = argparse.ArgumentParser(description='Demonstrate Hilbert curve patching for ViTs')
+    parser.add_argument('--image', type=str, default=None, help='Path to input image')
+    parser.add_argument('--patch_size', type=int, default=16, help='Patch size')
+    args = parser.parse_args()
+    
+    # Load image
+    if args.image and os.path.exists(args.image):
+        print(f"Loading image from {args.image}...")
+        image = load_image(args.image)
+    else:
+        image = load_sample_image()
+    
+    print(f"Original image shape: {image.shape}")
+    image = cv2.resize(image, (512, 512))  # Resize to a fixed size for demo
+    print(f"Image shape: {image.shape}")
+    # Ensure image dimensions are divisible by patch_size
+    h, w = image.shape[:2]
+    patch_size = args.patch_size
+    
+    # Crop to make dimensions divisible by patch_size
+    new_h = (h // patch_size) * patch_size
+    new_w = (w // patch_size) * patch_size
+    if new_h != h or new_w != w:
+        print(f"Cropping image from {h}x{w} to {new_h}x{new_w} to make divisible by patch size {patch_size}")
+        image = image[:new_h, :new_w]
+    
+    # 1. Visualize the Hilbert curve mapping
+    print("\n1. Visualizing Hilbert curve mapping...")
+    fig_map = visualize_hilbert_curve(new_h, new_w, patch_size)
+    
+    # 2. Demonstrate the patching process
+    print("\n2. Demonstrating Hilbert curve patching...")
+    fig_demo, fig_recon = demo_hilbert_patching(image, patch_size)
+    
+    # 3. Additional example: Process through a simulated transformer block
+    print("\n3. Simulating how patches would flow through a transformer...")
+    
+    # Convert to JAX array and add batch dimension
+    jax_img = jnp.array(image)[None, ...]  # [1, H, W, C]
+    
+    # Get Hilbert curve patches and inverse indices
+    patches, inv_idx = hilbert_patchify(jax_img, patch_size)
+    
+    print(f"Original image shape: {jax_img.shape}")
+    print(f"Patches shape: {patches.shape}")
+    
+    # Simulate a transformer block that operates on the patch sequence
+    def simulate_transformer_block(patches):
+        """
+        Simulate a transformer block by applying a simple operation to patches.
+        For demonstration purposes, we'll just multiply by a learned weight matrix.
+        """
+        batch, n_patches, patch_dim = patches.shape
+        
+        # Simulate learned weights (identity + small random values)
+        key = jax.random.PRNGKey(42)
+        weights = jnp.eye(patch_dim) + jax.random.normal(key, (patch_dim, patch_dim)) * 0.05
+        
+        # Apply "attention" (just a matrix multiply for demo)
+        return jnp.matmul(patches, weights)
+    
+    # Process patches as if through a transformer
+    processed_patches = simulate_transformer_block(patches)
+    
+    # Unpatchify back to image space
+    h, w, c = jax_img.shape[1:]
+    reconstructed = hilbert_unpatchify(processed_patches, inv_idx, patch_size, h, w, c)
+    
+    # Visualize the processed result
+    fig_processed, ax = plt.subplots(1, 2, figsize=(12, 5))
+    ax[0].imshow(np.array(jax_img[0]))
+    ax[0].set_title("Original Image")
+    ax[0].axis('off')
+    
+    ax[1].imshow(np.clip(np.array(reconstructed[0]), 0, 1))
+    ax[1].set_title("After Simulated Transformer Processing")
+    ax[1].axis('off')
+    plt.tight_layout()
+    
+    # Save all figures
+    print("\nSaving visualization figures...")
+    fig_map.savefig("hilbert_curve_mapping.png")
+    fig_demo.savefig("hilbert_patch_demo.png")
+    fig_recon.savefig("hilbert_patch_reconstruction.png")
+    fig_processed.savefig("hilbert_transformer_simulation.png")
+    
+    print("\nDone! Check the following output files:")
+    print("- hilbert_curve_mapping.png - Visualizes how Hilbert curve maps through a grid")
+    print("- hilbert_patch_demo.png - Shows patch ordering comparison")
+    print("- hilbert_patch_reconstruction.png - Shows original vs reconstructed image")
+    print("- hilbert_transformer_simulation.png - Shows a simple simulated transformer effect")
+    
+    # Display plots if running in interactive environment
+    plt.show()
+
+if __name__ == "__main__":
+    main()

demo_hilbert_curve.py

@@ -335,73 +335,4 @@ def __call__(self, x:jax.Array, temb:jax.Array, textemb:jax.Array=None, extra_fe
 
         out = jnp.concatenate([out, extra_features], axis=-1) if extra_features is not None else out
 
-        return out
-
-# Convert Hilbert index d to 2D coordinates (x, y) for an n x n grid
-def _d2xy(n, d):
-    x = 0
-    y = 0
-    t = d
-    s = 1
-    while s < n:
-        rx = (t // 2) & 1
-        ry = (t ^ rx) & 1
-        if ry == 0:
-            if rx == 1:
-                x = n - 1 - x
-                y = n - 1 - y
-            x, y = y, x
-        x += s * rx
-        y += s * ry
-        t //= 4
-        s *= 2
-    return x, y
-
-# Hilbert index mapping for a rectangular grid of patches H_P x W_P
-
-def hilbert_indices(H_P, W_P):
-    size = max(H_P, W_P)
-    order = math.ceil(math.log2(size))
-    n = 1 << order
-    coords = []
-    for d in range(n * n):
-        x, y = _d2xy(n, d)
-        # x is column index, y is row index
-        if x < W_P and y < H_P:
-            coords.append((y, x))  # (row, col)
-            if len(coords) == H_P * W_P:
-                break
-    # Convert (row, col) to linear indices row-major
-    indices = [r * W_P + c for r, c in coords]
-    return jnp.array(indices, dtype=jnp.int32)
-
-# Inverse permutation: given idx where idx[i] = new position of element i, return inv such that inv[idx[i]] = i
-
-def inverse_permutation(idx):
-    inv = jnp.zeros_like(idx)
-    inv = inv.at[idx].set(jnp.arange(idx.shape[0], dtype=idx.dtype))
-    return inv
-
-# Patchify using Hilbert ordering: extract patches and reorder sequence
-
-def hilbert_patchify(x, patch_size):
-    B, H, W, C = x.shape
-    H_P = H // patch_size
-    W_P = W // patch_size
-    # Extract patches in row-major
-    patches = rearrange(x, 'b (h p1) (w p2) c -> b (h w) (p1 p2 c)', p1=patch_size, p2=patch_size)
-    idx = hilbert_indices(H_P, W_P)
-    return patches[:, idx, :]
-
-# Unpatchify from Hilbert ordering: reorder sequence back and reconstruct image
-
-def hilbert_unpatchify(patches, patch_size, H, W, C):
-    B, N, D = patches.shape
-    H_P = H // patch_size
-    W_P = W // patch_size
-    inv = inverse_permutation(hilbert_indices(H_P, W_P))
-    # Reorder back to row-major
-    linear = patches[:, inv, :]
-    # Reconstruct image
-    x = rearrange(linear, 'b (h w) (p1 p2 c) -> b (h p1) (w p2) c', h=H_P, w=W_P, p1=patch_size, p2=patch_size, c=C)
-    return x
+        return out