Building a Browser-Based AI Image Colorization Tool

Introduction

In this article, we'll explore how to implement an AI-powered image colorization tool that runs entirely in the browser. This approach offers significant advantages in terms of privacy, cost, and user experience.

Why Browser-Based Implementation?

There are several compelling reasons to implement image colorization in the browser:

1. Privacy Protection

When users colorize images in the browser, their photos never leave their device. This is particularly important for:

• Old family photos that may contain sensitive information
• Business documents
• Medical images
• Any content users want to keep private

2. No Server Costs

Traditional server-based image processing requires:

• Powerful GPU servers for AI inference
• Bandwidth for uploading/downloading images
• Storage for temporary files

By running the model in the browser, we eliminate all these costs.

3. Offline Capability

Once the model is loaded, users can colorize images without an internet connection. This is especially valuable for:

• Users in areas with limited connectivity
• Mobile users who want to save data
• Privacy-conscious users who prefer air-gapped processing

4. Zero Latency

No network round-trip means faster processing. Users see results immediately after clicking the colorize button.

Technical Architecture

Core Implementation

1. Data Structures

interface ImageFile {
  id: string;
  file: File;
  previewUrl: string;
  colorizedUrl?: string;
  colorizedFileName?: string;
  error?: string;
  processing?: boolean;
}

2. Loading ONNX Runtime and Model

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

async function loadModel() {
  // Load ONNX Runtime from CDN
  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, reject) => {
    script.onload = () => resolve();
    script.onerror = () => reject(new Error("Failed to load ONNX Runtime"));
  });

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

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    // Update progress...
  }

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

  return session;
}

3. Image Preprocessing

The model expects input in a specific format: a 256x256 RGB image as a Float32Array with shape [1, 3, 256, 256].

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

  // Extract RGB channels from RGBA ImageData
  let j = 0;
  for (let i = 1; i < imageData.data.length + 1; i++) {
    if (i % 4 !== 0) {
      floatArr[j] = imageData.data[i - 1];
      j += 1;
    }
  }

  // Reorganize to [R..., G..., B...] format
  const floatArr2 = new Float32Array(width * height * 3);
  // ... channel reorganization logic ...

  return floatArr2;
}

4. Running Inference

async function colorizeImage(imageFile: ImageFile, session) {
  const img = await createImageBitmap(imageFile.file);
  const size = 256;

  // Resize to model's expected input size
  const canvas = new OffscreenCanvas(size, size);
  const ctx = canvas.getContext("2d");
  ctx.drawImage(img, 0, 0, size, size);

  const inputImageData = ctx.getImageData(0, 0, size, size);
  const processed = preprocess(inputImageData, size, size);

  // Create tensor and run inference
  const input = new window.ort.Tensor(new Float32Array(processed), [1, 3, size, size]);
  const result = await session.run({ input });

  const output = result["out"];
  return postprocess(output);
}

5. Postprocessing

function postprocess(tensor): ImageData {
  const channels = tensor.dims[1];
  const height = tensor.dims[2];
  const width = tensor.dims[3];

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

  for (let h = 0; h < height; h++) {
    for (let w = 0; w < width; w++) {
      const rgb = [];
      for (let c = 0; c < channels; c++) {
        const tensorIndex = (c * height + h) * width + w;
        const value = tensorData[tensorIndex];
        rgb.push(Math.round(Math.max(0, Math.min(255, value))));
      }

      const pixelIndex = (h * width + w) * 4;
      imageData.data[pixelIndex] = rgb[0];
      imageData.data[pixelIndex + 1] = rgb[1];
      imageData.data[pixelIndex + 2] = rgb[2];
      imageData.data[pixelIndex + 3] = 255;
    }
  }

  return imageData;
}

Service Worker for Model Caching

The model is approximately 25MB, so we use a Service Worker to cache it for faster subsequent loads:

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

self.addEventListener('fetch', (event) => {
  if (url.href.includes('deoldify.quant.onnx')) {
    event.respondWith(
      caches.match(event.request).then((cachedResponse) => {
        if (cachedResponse) {
          return cachedResponse;
        }
        return fetch(event.request).then((networkResponse) => {
          const responseToCache = networkResponse.clone();
          caches.open(CACHE_NAME).then((cache) => {
            cache.put(event.request, responseToCache);
          });
          return networkResponse;
        });
      })
    );
  }
});

Processing Flow

Key Technologies Used

Technology	Purpose
ONNX Runtime Web	Run ONNX models in browser
DeOldify Model	Pre-trained colorization model
OffscreenCanvas	Process images in background
Service Worker	Cache AI model for faster load
Web Workers	(Optional) Further parallelization

Performance Considerations

1. Model Size: The quantized DeOldify model is ~25MB, which is reasonable for browser caching
2. Processing Time: Colorization takes 1-3 seconds depending on device
3. Memory Usage: Peak memory is approximately 3x the image size
4. WebGL Backend: ONNX Runtime can use WebGL for GPU acceleration when available

Conclusion

Browser-based AI image colorization is a practical approach that delivers excellent user experience while protecting privacy and reducing costs. The combination of ONNX Runtime Web, WebGL acceleration, and Service Worker caching makes it possible to run sophisticated AI models directly in the user's browser.

Try it yourself at Free Image Tools

Note: When sharing or linking to our tool, you can help us understand your traffic source by linking from your site. This helps us improve our services! When you link to us, please include a Referer header pointing to your website so we can track where our visitors come from.

Experience the power of browser-based AI image colorization. No upload required - your images stay on your device!