DEV Community

monkeymore studio
monkeymore studio

Posted on

Browser-Based AI Image Colorization with DeOldify and ONNX Runtime

Introduction

In this article, we'll explore how to build a pure frontend image colorization system that runs entirely in the browser using AI. Unlike traditional image processing services that upload your photos to remote servers, this approach leverages WebAssembly and ONNX Runtime to perform neural network inference locally, ensuring complete privacy while delivering impressive results.

Why Browser-Based Colorization?

The Privacy Imperative

Traditional AI image processing services typically require uploading images to cloud servers:

❌ Your photos travel over the internet to unknown servers
❌ Processing happens on third-party infrastructure  
❌ Images may be stored, analyzed, or used for training
❌ Requires constant internet connection
❌ Upload/download delays for large images
Enter fullscreen mode Exit fullscreen mode

By implementing AI colorization directly in the browser, we achieve:

✅ Images never leave the user's device
✅ Zero network transmission of sensitive photos
✅ Works offline after initial model load
✅ Instant processing without upload delays
✅ Complete data sovereignty
Enter fullscreen mode Exit fullscreen mode

This is particularly important for:

  • Historical archives - Old family photos containing personal information
  • Medical imaging - HIPAA compliance requires data locality
  • Legal documents - Attorney-client privilege protection
  • Journalism - Source protection and operational security

Technical Architecture

Our implementation uses DeOldify, a deep learning model for image colorization, packaged as an ONNX model and executed via ONNX Runtime Web in the browser.

System Architecture

Core Data Structures

Image File Interface

Each image is tracked with comprehensive metadata:

interface ImageFile {
  id: string;                    // Unique identifier
  file: File;                    // Original file object
  previewUrl: string;            // Blob URL for preview
  colorizedUrl?: string;         // Result blob URL
  colorizedFileName?: string;    // Generated filename
  error?: string;                // Error message if failed
  processing?: boolean;          // Current processing state
}
Enter fullscreen mode Exit fullscreen mode

Model Configuration

const MODEL_URL = "https://raw.githubusercontent.com/linmingren/openmodels/main/models/deoldify/deoldify.quant.onnx";

// Model is quantized to reduce size from ~100MB to ~25MB
// while maintaining quality through INT8 quantization
Enter fullscreen mode Exit fullscreen mode

Implementation Deep Dive

Step 1: Service Worker Registration

The Service Worker acts as a caching layer for the AI model:

useEffect(() => {
  // Register Service Worker for model caching
  if (typeof window !== 'undefined' && 'serviceWorker' in navigator) {
    navigator.serviceWorker
      .register('/colorize-sw.js')
      .then((registration) => {
        console.log('[Colorize] Service Worker registered:', registration);
      })
      .catch((error) => {
        console.error('[Colorize] Service Worker registration failed:', error);
      });
  }
  // ... model loading
}, []);
Enter fullscreen mode Exit fullscreen mode

Why Service Worker?

  • Caches the 25MB model locally after first download
  • Enables offline usage on subsequent visits
  • Intercepts fetch requests to serve cached model
  • Survives page refreshes and navigation

Step 2: Service Worker Implementation

// colorize-sw.js
const CACHE_NAME = 'colorize-model-cache-v1';
const MODEL_URL = 'https://raw.githubusercontent.com/linmingren/openmodels/main/models/deoldify/deoldify.quant.onnx';

// Install: Cache model immediately
self.addEventListener('install', (event) => {
  event.waitUntil(
    caches.open(CACHE_NAME).then((cache) => {
      return cache.add(MODEL_URL);
    })
  );
  self.skipWaiting();
});

// Fetch: Serve from cache if available
self.addEventListener('fetch', (event) => {
  const url = new URL(event.request.url);

  if (url.href.includes('deoldify.quant.onnx')) {
    event.respondWith(
      caches.match(event.request).then((cachedResponse) => {
        if (cachedResponse) {
          return cachedResponse;  // Return cached model
        }
        return fetch(event.request);  // Fetch from network
      })
    );
  }
});
Enter fullscreen mode Exit fullscreen mode

Step 3: ONNX Runtime Loading


async function loadModel() {
  // Load ONNX Runtime if not present
  if (!window.ort) {
    const script = document.createElement("script");
    script.src = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.min.js";
    document.head.appendChild(script);

    await new Promise<void>((resolve) => {
      script.onload = () => resolve();
    });
  }

  // Fetch model with progress tracking
  const response = await fetch(MODEL_URL);
  const reader = response.body?.getReader();
  const chunks = [];
  let receivedLength = 0;

  // Stream download with progress
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    receivedLength += value.length;
    const percent = Math.round((receivedLength / (25 * 1024 * 1024)) * 100);
    setLoadProgress(percent);
  }

  // Assemble model buffer
  const modelBuffer = new Uint8Array(
    chunks.reduce((acc, val) => acc + val.length, 0)
  ).buffer;

  // Create inference session
  const session = await window.ort.InferenceSession.create(
    new Uint8Array(modelBuffer)
  );

  sessionRef.current = session;
}
Enter fullscreen mode Exit fullscreen mode

Step 4: Image Preprocessing

The DeOldify model expects specific input format (NCHW - Batch, Channel, Height, Width):

const preprocess = (
  imageData: ImageData, 
  width: number, 
  height: number
): Float32Array => {
  const floatArr = new Float32Array(width * height * 3);
  const floatArr2 = new Float32Array(width * height * 3);

  // Convert RGBA to RGB float array
  let j = 0;
  for (let i = 1; i < imageData.data.length + 1; i++) {
    if (i % 4 !== 0) {  // Skip alpha channel
      floatArr[j] = imageData.data[i - 1];
      j += 1;
    }
  }

  // Reorganize to NCHW format (Channel first)
  let k = 0, l, m;

  // Red channel
  for (let i = 0; i < floatArr.length; i += 3) {
    floatArr2[k] = floatArr[i];
    k += 1;
  }

  // Green channel
  l = k;
  for (let i = 1; i < floatArr.length; i += 3) {
    floatArr2[l] = floatArr[i];
    l += 1;
  }

  // Blue channel
  m = l;
  for (let i = 2; i < floatArr.length; i += 3) {
    floatArr2[m] = floatArr[i];
    m += 1;
  }

  return floatArr2;
};
Enter fullscreen mode Exit fullscreen mode

Step 5: Colorization Pipeline


const colorizeImage = async (
  imageFile: ImageFile
): Promise<{ blob: Blob; fileName: string }> => {
  if (!sessionRef.current) {
    throw new Error("Model not loaded");
  }

  // Step 1: Load image
  const img = await createImageBitmap(imageFile.file);

  // Step 2: Resize to model input size (256x256)
  const size = 256;
  const canvas = new OffscreenCanvas(size, size);
  const ctx = canvas.getContext("2d")!;
  ctx.drawImage(img, 0, 0, size, size);

  // Step 3: Get pixel data
  const inputImageData = ctx.getImageData(0, 0, size, size);

  // Step 4: Preprocess for model
  const processed = preprocess(inputImageData, size, size);
  const input = new window.ort.Tensor(
    new Float32Array(processed), 
    [1, 3, size, size]  // Batch=1, Channels=3, Height=256, Width=256
  );

  // Step 5: Run inference
  const result = await sessionRef.current.run({ input });
  const output = result["out"] as import("onnxruntime-web").Tensor;

  // Step 6: Postprocess output
  const processedImageData = postprocess(output);

  // Step 7: Restore original dimensions
  const outputCanvas = new OffscreenCanvas(img.width, img.height);
  const outputCtx = outputCanvas.getContext("2d")!;
  outputCtx.drawImage(processedCanvas, 0, 0, img.width, img.height);

  // Step 8: Export as PNG
  const blob = await outputCanvas.convertToBlob({ type: "image/png" });

  return { 
    blob, 
    fileName: `${baseName}_colorized.png` 
  };
};
Enter fullscreen mode Exit fullscreen mode

Step 6: Postprocessing

Convert model output back to displayable image:

const postprocess = (tensor: import("onnxruntime-web").Tensor): ImageData => {
  const channels = tensor.dims[1];  // 3 (RGB)
  const height = tensor.dims[2];    // 256
  const width = tensor.dims[3];     // 256

  const imageData = new ImageData(width, height);
  const data = imageData.data;
  const tensorData = tensor.data as Float32Array;

  // Convert NCHW back to RGBA
  for (let h = 0; h < height; h++) {
    for (let w = 0; w < width; w++) {
      const rgb = [];

      // Extract RGB values for this pixel
      for (let c = 0; c < channels; c++) {
        const tensorIndex = (c * height + h) * width + w;
        const value = tensorData[tensorIndex];
        // Clamp to valid pixel range
        rgb.push(Math.round(Math.max(0, Math.min(255, value))));
      }

      // Set RGBA values
      const pixelIndex = (h * width + w) * 4;
      data[pixelIndex] = rgb[0];      // R
      data[pixelIndex + 1] = rgb[1];  // G
      data[pixelIndex + 2] = rgb[2];  // B
      data[pixelIndex + 3] = 255;     // A (fully opaque)
    }
  }

  return imageData;
};
Enter fullscreen mode Exit fullscreen mode

Step 7: Batch Processing

Handle multiple images with progress tracking:

const colorizeAllImages = useCallback(async () => {
  if (selectedFiles.length === 0 || !sessionRef.current) return;

  setIsColorizing(true);
  const uncolorizedFiles = selectedFiles.filter(
    f => !f.colorizedUrl && !f.error
  );

  for (let i = 0; i < uncolorizedFiles.length; i++) {
    const imageFile = uncolorizedFiles[i];

    // Mark as processing
    setSelectedFiles(prev => prev.map(f => 
      f.id === imageFile.id ? { ...f, processing: true } : f
    ));

    // Small delay to allow UI update
    await new Promise(r => setTimeout(r, 50));

    try {
      const { blob, fileName } = await colorizeImage(imageFile);
      const colorizedUrl = URL.createObjectURL(blob);

      // Update with result
      setSelectedFiles(prev => prev.map(f => 
        f.id === imageFile.id ? {
          ...f,
          colorizedUrl,
          colorizedFileName: fileName,
          processing: false,
        } : f
      ));
    } catch (err) {
      setSelectedFiles(prev => prev.map(f => 
        f.id === imageFile.id ? {
          ...f,
          error: `Failed to colorize ${imageFile.file.name}`,
          processing: false,
        } : f
      ));
    }
  }

  setIsColorizing(false);
}, [selectedFiles]);
Enter fullscreen mode Exit fullscreen mode

Memory Management

Proper cleanup of blob URLs prevents memory leaks:

const removeImage = useCallback((id: string) => {
  setSelectedFiles(prev => {
    const file = prev.find(f => f.id === id);
    if (file) {
      // Revoke blob URLs to free memory
      URL.revokeObjectURL(file.previewUrl);
      if (file.colorizedUrl) {
        URL.revokeObjectURL(file.colorizedUrl);
      }
    }
    return prev.filter(f => f.id !== id);
  });
}, []);

const reset = useCallback(() => {
  selectedFiles.forEach(file => {
    URL.revokeObjectURL(file.previewUrl);
    if (file.colorizedUrl) URL.revokeObjectURL(file.colorizedUrl);
  });
  setSelectedFiles([]);
}, [selectedFiles]);
Enter fullscreen mode Exit fullscreen mode

Supported Formats

The system accepts common web image formats:

const validTypes = ['image/jpeg', 'image/png', 'image/webp'];

const newImages = Array.from(files)
  .filter(file => validTypes.includes(file.type))
  .map(file => ({
    id: `${file.name}-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`,
    file,
    previewUrl: URL.createObjectURL(file),
  }));
Enter fullscreen mode Exit fullscreen mode

Technical Stack

Component Technology Purpose
Framework React 19 UI components
Build Tool Next.js 16 SSR and static generation
Styling Tailwind CSS 4 Utility-first CSS
AI Runtime onnxruntime-web@1.24.2 Neural network inference
AI Model DeOldify (Quantized) Image colorization
Storage Service Worker Cache Model persistence
Graphics OffscreenCanvas Image processing
Icons Lucide React UI icons

Performance Characteristics

Metric Value
Model Size ~25 MB (quantized from ~100 MB)
Input Size 256×256 pixels (model constraint)
Output Size Matches original image dimensions
Processing Time 2-5 seconds per image (depends on device)
Memory Usage ~200-400 MB during inference
First Load 25 MB model download
Subsequent Loads Instant (cached)

Browser Compatibility

Requirements:

  • WebAssembly: All modern browsers (Chrome 57+, Firefox 52+, Safari 11+, Edge 16+)
  • Service Worker: Chrome 40+, Firefox 44+, Safari 11.1+, Edge 17+
  • OffscreenCanvas: Chrome 69+, Firefox 79+, Safari 16.4+, Edge 79+
  • WebGL (for ONNX): Universal support

Note: The model requires ~400MB RAM during inference. Mobile devices with limited memory may experience slower performance.

Why This Architecture?

Why Not Use Web Workers?

While Web Workers can offload processing from the main thread, in this implementation:

  1. ONNX Runtime Web runs in WASM, which already executes outside the main JavaScript thread
  2. Canvas operations (resize, draw) must happen on the main thread for DOM access
  3. Memory overhead of transferring large image buffers to workers is significant
  4. Simpler debugging with single-threaded code

For production with heavy usage, consider a Worker-based architecture for the image preprocessing steps.

Why Service Worker vs LocalStorage?

  • Size: LocalStorage limited to 5-10 MB; we need 25 MB
  • Binary data: Service Worker cache handles ArrayBuffers natively
  • Offline capability: Service Worker intercepts fetch requests automatically
  • Background sync: Can cache during idle time

Try It Yourself

Ready to bring your black and white photos to life? Try our browser-based colorization tool:

👉 Colorize Your Images

Your photos are processed entirely in your browser - they never leave your device, ensuring complete privacy.

Conclusion

Browser-based AI image processing represents a paradigm shift in web development. By leveraging WebAssembly and modern browser APIs, we can run sophisticated neural networks like DeOldify directly on the client side.

Key advantages:

  • Privacy-first: No image uploads required
  • Offline capable: Works without internet after first load
  • Cost-effective: No server GPU infrastructure needed
  • Instant results: No network latency

This architecture is ideal for any application handling sensitive visual data where privacy and speed are paramount.

Further Reading


Want to add AI-powered colorization to your own web applications? The combination of ONNX Runtime Web and browser-based model caching makes it remarkably straightforward to deploy deep learning models directly to users.

Top comments (0)