gif-encoder-2/src/OctreeQuant.js

GIF generator for Node.js

/*
  Authors
  Dmitry Alimov (Python version) https://github.com/delimitry/octree_color_quantizer
  Tom MacWright (JavaScript version) https://observablehq.com/@tmcw/octree-color-quantization
*/

const MAX_DEPTH = 8

class OctreeQuant {
  constructor() {
    this.levels = Array.from({ length: MAX_DEPTH }, () => [])
    this.root = new Node(0, this)
  }

  addColor(color) {
    this.root.addColor(color, 0, this)
  }

  makePalette(colorCount) {
    let palette = []
    let paletteIndex = 0
    let leafCount = this.leafNodes.length
    for (let level = MAX_DEPTH - 1; level > -1; level -= 1) {
      if (this.levels[level]) {
        for (let node of this.levels[level]) {
          leafCount -= node.removeLeaves()
          if (leafCount <= colorCount) break
        }
        if (leafCount <= colorCount) break
        this.levels[level] = []
      }
    }
    for (let node of this.leafNodes) {
      if (paletteIndex >= colorCount) break
      if (node.isLeaf) palette.push(node.color)
      node.paletteIndex = paletteIndex
      paletteIndex++
    }
    return palette
  }

  *makePaletteIncremental(colorCount) {
    let palette = []
    let paletteIndex = 0
    let leafCount = this.leafNodes.length
    for (let level = MAX_DEPTH - 1; level > -1; level -= 1) {
      if (this.levels[level]) {
        for (let node of this.levels[level]) {
          leafCount -= node.removeLeaves()
          if (leafCount <= colorCount) break
        }
        if (leafCount <= colorCount) break
        this.levels[level] = []
      }
      yield
    }
    for (let node of this.leafNodes) {
      if (paletteIndex >= colorCount) break
      if (node.isLeaf) palette.push(node.color)
      node.paletteIndex = paletteIndex
      paletteIndex++
    }
    yield
    return palette
  }

  get leafNodes() {
    return this.root.leafNodes
  }

  addLevelNode(level, node) {
    this.levels[level].push(node)
  }

  getPaletteIndex(color) {
    return this.root.getPaletteIndex(color, 0)
  }
}

class Node {
  constructor(level, parent) {
    this._color = new Color(0, 0, 0)
    this.pixelCount = 0
    this.paletteIndex = 0
    this.children = []
    this._debugColor
    if (level < MAX_DEPTH - 1) parent.addLevelNode(level, this)
  }

  get isLeaf() {
    return this.pixelCount > 0
  }

  get leafNodes() {
    let leafNodes = []

    for (let node of this.children) {
      if (!node) continue
      if (node.isLeaf) {
        leafNodes.push(node)
      } else {
        leafNodes.push(...node.leafNodes)
      }
    }

    return leafNodes
  }

  addColor(color, level, parent) {
    if (level >= MAX_DEPTH) {
      this._color.add(color)
      this.pixelCount++
      return
    }
    let index = getColorIndex(color, level)
    if (!this.children[index]) {
      this.children[index] = new Node(level, parent)
    }
    this.children[index].addColor(color, level + 1, parent)
  }

  getPaletteIndex(color, level) {
    if (this.isLeaf) {
      return this.paletteIndex
    }
    let index = getColorIndex(color, level)
    if (this.children[index]) {
      return this.children[index].getPaletteIndex(color, level + 1)
    } else {
      for (let node of this.children) {
        if (node) {
          return node.getPaletteIndex(color, level + 1)
        }
      }
    }
  }

  removeLeaves() {
    let result = 0
    for (let node of this.children) {
      if (!node) continue
      this._color.add(node._color)
      this.pixelCount += node.pixelCount
      result++
    }
    this.children = []
    return result - 1
  }

  get debugColor() {
    if (this._debugColor) return this._debugColor
    if (this.isLeaf) return this.color

    let c = new Color()
    let count = 0

    function traverse(node) {
      for (let child of node.children) {
        if (child.isLeaf) {
          c.add(child._color)
          count++
        } else {
          traverse(child)
        }
      }
    }

    traverse(this)
    return c.normalized(count)
  }

  get color() {
    return this._color.normalized(this.pixelCount)
  }
}

class Color {
  constructor(red = 0, green = 0, blue = 0) {
    this.red = red
    this.green = green
    this.blue = blue
  }

  clone() {
    return new Color(this.red, this.green, this.blue)
  }

  get array() {
    return [this.red, this.green, this.blue, this.red + this.green + this.blue]
  }

  toString() {
    return [this.red, this.green, this.blue].join(',')
  }

  toCSS() {
    return `rgb(${[this.red, this.green, this.blue].map(n => Math.floor(n)).join(',')})`
  }

  normalized(pixelCount) {
    return new Color(this.red / pixelCount, this.green / pixelCount, this.blue / pixelCount)
  }

  add(color) {
    this.red += color.red
    this.green += color.green
    this.blue += color.blue
  }
}

function getColorIndex(color, level) {
  let index = 0
  let mask = 0b10000000 >> level
  if (color.red & mask) index |= 0b100
  if (color.green & mask) index |= 0b010
  if (color.blue & mask) index |= 0b001
  return index
}

module.exports = { OctreeQuant, Node, Color }