sixel: proper color quantization

just use an octree. works fine afaict, though obviously somewhat slower
than the static method (encoding is 2-pass now) & still has banding
issues with many colors (will need dithering)

also, fixed a bug that caused initial masks of bands to get misplaced

1 files changed, 215 insertions, 37 deletions

diff --git a/adapter/img/sixel.nim b/adapter/img/sixel.nim
index 02bb356e..7f4514d9 100644
--- a/adapter/img/sixel.nim
+++ b/adapter/img/sixel.nim
@@ -12,6 +12,7 @@
 #
 # This way, the image can be vertically cropped in ~constant time.
 
+import std/algorithm
 import std/options
 import std/os
 import std/strutils
@@ -23,38 +24,40 @@ import utils/twtstr
 const DCSSTART = "\eP"
 const ST = "\e\\"
 
+type SixelBand = object
+ c: int
+ data: seq[uint8]
+
 # data is binary 0..63; the output is the final ASCII form.
-proc compressSixel(data: openArray[uint8]; c: uint8): string =
-  var outs = newStringOfCap(data.len div 4 + 3)
+proc compressSixel(band: SixelBand): string =
+  var outs = newStringOfCap(band.data.len div 4 + 3)
   outs &= '#'
-  outs &= $c
+  outs &= $band.c
   var n = 0
   var c = char(0)
-  for u in data:
+  for u in band.data:
     let cc = char(u + 0x3F)
     if c != cc:
       if n > 3:
         outs &= '!' & $n & c
       else: # for char(0) n is also 0, so it is ignored.
-        outs &= c.repeat(n)
+        for i in 0 ..< n:
+          outs &= c
       c = cc
       n = 0
     inc n
   if n > 3:
     outs &= '!' & $n & c
   else:
-    outs &= c.repeat(n)
+    for i in 0 ..< n:
+      outs &= c
   return outs
 
-type SixelBand = object
- c: uint8
- data: seq[uint8]
-
-func find(bands: seq[SixelBand]; c: uint8): int =
-  for i, band in bands:
-    if band.c == c:
+func find(bands: seq[SixelBand]; c: int): int =
+  for i in 0 ..< bands.len:
+    if bands[i].c == c:
       return i
-  -1
+  return -1
 
 proc setU32BE(s: var string; n: uint32) =
   s[0] = char(n and 0xFF)
@@ -68,6 +71,190 @@ proc putU32BE(s: var string; n: uint32) =
   s &= char((n shr 16) and 0xFF)
   s &= char((n shr 24) and 0xFF)
 
+type Node {.acyclic.} = ref object
+  leaf: bool
+  c: RGBColor
+  n: uint32
+  r: uint32
+  g: uint32
+  b: uint32
+  children: array[8, Node]
+
+proc getIdx(c: RGBColor; level: int): uint8 {.inline.} =
+  let sl = 7 - level
+  let idx = (((c.r shr sl) and 1) shl 2) or
+    (((c.g shr sl) and 1) shl 1) or
+    (c.b shr sl) and 1
+  return idx
+
+type TrimMap = array[7, seq[Node]]
+
+# Insert a node into the octree.
+# Returns true if a new leaf was inserted, false otherwise.
+proc insert(parent: Node; c: RGBColor; trimMap: var TrimMap; level = 0;
+    n = 1u32): bool =
+  # max level is 7, because we only have ~6.5 bits (0..100, inclusive)
+  # (it *is* 0-indexed, but one extra level is needed for the final leaves)
+  assert not parent.leaf and level < 8
+  let idx = c.getIdx(level)
+  let old = parent.children[idx]
+  if old == nil:
+    if level == 7:
+      parent.children[idx] = Node(
+        leaf: true,
+        c: c,
+        n: n,
+        r: uint32(c.r) * n,
+        g: uint32(c.g) * n,
+        b: uint32(c.b) * n
+      )
+      return true
+    else:
+      let container = Node(leaf: false)
+      parent.children[idx] = container
+      trimMap[level].add(container)
+      return container.insert(c, trimMap, level + 1, n)
+  elif old.leaf:
+    if old.c == c:
+      old.n += n
+      old.r += uint32(c.r) * n
+      old.g += uint32(c.g) * n
+      old.b += uint32(c.b) * n
+      return false
+    else:
+      let container = Node(leaf: false)
+      parent.children[idx] = container
+      let nlevel = level + 1
+      container.children[old.c.getIdx(nlevel)] = old # skip an alloc :)
+      trimMap[level].add(container)
+      return container.insert(c, trimMap, nlevel, n)
+  else:
+    return old.insert(c, trimMap, level + 1, n)
+
+proc trim(trimMap: var TrimMap; K: var int) =
+  var node: Node = nil
+  for i in countdown(trimMap.high, 0):
+    if trimMap[i].len > 0:
+      node = trimMap[i].pop()
+      break
+  assert node != nil
+  var r = 0u32
+  var g = 0u32
+  var b = 0u32
+  var n = 0u32
+  var k = K + 1
+  for child in node.children.mitems:
+    if child != nil:
+      assert child.leaf
+      r += child.r
+      g += child.g
+      b += child.b
+      n += child.n
+      child = nil
+      dec k
+  node.leaf = true
+  node.c = rgb(uint8(r div n), uint8(g div n), uint8(b div n))
+  node.r = r
+  node.g = g
+  node.b = b
+  node.n = n
+  K = k
+
+proc getPixel(s: string; m: int; bgcolor: ARGBColor): RGBColor {.inline.} =
+  let r = uint8(s[m])
+  let g = uint8(s[m + 1])
+  let b = uint8(s[m + 2])
+  let a = uint8(s[m + 3])
+  var c0 = RGBAColorBE(r: r, g: g, b: b, a: a)
+  if c0.a != 255:
+    let c1 = bgcolor.blend(c0)
+    return RGBColor(uint32(rgb(c1.r, c1.g, c1.b)).fastmul(100))
+  return RGBColor(uint32(rgb(c0.r, c0.g, c0.b)).fastmul(100))
+
+proc quantize(s: string; bgcolor: ARGBColor; palette: int): Node =
+  let root = Node(leaf: false)
+  # number of leaves
+  var K = 0
+  # map of non-leaves for each level.
+  # (note: somewhat confusingly, this actually starts at level 1.)
+  var trimMap: array[7, seq[Node]]
+  # batch together insertions of color runs
+  var pc0 = RGBColor(0)
+  var pcs = 0u32
+  for i in 0 ..< s.len div 4:
+    let m = i * 4
+    let c0 = s.getPixel(m, bgcolor)
+    inc pcs
+    if pc0 != c0:
+      K += int(root.insert(c0, trimMap, n = pcs))
+      pcs = 0
+    while K > palette:
+      # trim the tree.
+      trimMap.trim(K)
+  if pcs > 0:
+    K += int(root.insert(pc0, trimMap, n = pcs))
+    while K > palette:
+      # trim the tree.
+      trimMap.trim(K)
+  return root
+
+type
+  QuantMap = array[4096, seq[tuple[idx: int; c: RGBColor]]]
+
+  ColorPair = tuple[c: RGBColor; n: uint32]
+
+proc flatten(node: Node; map: var QuantMap; cols: var seq[ColorPair]) =
+  if node.leaf:
+    cols.add((node.c, node.n))
+  else:
+    for child in node.children:
+      if child != nil:
+        child.flatten(map, cols)
+
+proc flatten(node: Node; outs: var string; palette: int): QuantMap =
+  var map: QuantMap
+  var cols = newSeqOfCap[ColorPair](palette)
+  node.flatten(map, cols)
+  # try to set the most common colors as the smallest numbers (so we write less)
+  cols.sort(proc(a, b: ColorPair): int = cmp(a.n, b.n), order = Descending)
+  for n, it in cols:
+    let n = n + 1
+    let c = it.c
+    # 2 is RGB
+    outs &= '#' & $n & ";2;" & $c.r & ';' & $c.g & ';' & $c.b
+    let i = (int(c.r shr 4) shl 8) or
+      (int(c.g shr 4) shl 4) or
+      (int(c.b shr 4))
+    map[i].add((n, c))
+  return map
+
+proc getColor(map: QuantMap; c: RGBColor): int =
+  var i = (int(c.r shr 4) shl 8) or
+    (int(c.g shr 4) shl 4) or
+    int(c.b shr 4)
+  if map[i].len == 0:
+    var j = i
+    while true:
+      dec i
+      inc j
+      if i >= 0 and map[i].len > 0:
+        break
+      if j < map.len and map[j].len > 0:
+        i = j
+        break
+      assert i >= 0 or j < map.len # assuming map isn't empty...
+  var minDist = uint32.high
+  var resIdx = -1
+  for (idx, ic) in map[i]:
+    var d = uint32(abs(int32(c.r) - int32(ic.r))) +
+      uint32(abs(int32(c.g) - int32(ic.g))) +
+      uint32(abs(int32(c.b) - int32(ic.b)))
+    if d < minDist:
+      minDist = d
+      resIdx = idx
+  assert resIdx != -1
+  return resIdx
+
 proc encode(s: string; width, height, offx, offy, cropw: int; halfdump: bool;
     bgcolor: ARGBColor; palette: int) =
   if width == 0 or height == 0:
@@ -80,14 +267,9 @@ proc encode(s: string; width, height, offx, offy, cropw: int; halfdump: bool;
     outs &= DCSSTART & 'q'
     # set raster attributes
     outs &= "\"1;1;" & $width & ';' & $height
-  for b in 16 ..< 256:
-    # laziest possible color register allocation scheme
-    #TODO obviously this produces sub-optimal results
-    let rgb = EightBitColor(b).toRGB()
-    let rgbq = RGBColor(uint32(rgb).fastmul(100))
-    let n = b - 15
-    # 2 is RGB
-    outs &= '#' & $n & ";2;" & $rgbq.r & ';' & $rgbq.g & ';' & $rgbq.b
+  # reserve one entry for empty lines
+  let node = s.quantize(bgcolor, palette - 1)
+  let map = node.flatten(outs, palette)
   if halfdump:
     # prepend prelude size
     let L = outs.len - 4 # subtract length field
@@ -109,25 +291,21 @@ proc encode(s: string; width, height, offx, offy, cropw: int; halfdump: bool;
       let realw = cropw - offx
       for j in 0 ..< realw:
         let m = n + (j + offx) * 4
-        let r = uint8(s[m])
-        let g = uint8(s[m + 1])
-        let b = uint8(s[m + 2])
-        let a = uint8(s[m + 3])
-        var c0 = RGBAColorBE(r: r, g: g, b: b, a: a)
-        if c0.a != 255:
-          let c1 = bgcolor.blend(c0)
-          c0 = RGBAColorBE(r: c1.r, g: c1.g, b: c1.b, a: c1.a)
-        let c = uint8(c0.toEightBit())
-        if (let k = bands.find(c); k != -1):
-          bands[k].data[j] = bands[k].data[j] or mask
-        else:
+        let c0 = s.getPixel(m, bgcolor)
+        let c = map.getColor(c0)
+        #TODO this could be optimized a lot more, by squashing together bands
+        # with empty runs at different places.
+        var k = bands.find(c)
+        if k == -1:
           bands.add(SixelBand(c: c, data: newSeq[uint8](realw)))
-          bands[^1].data[^1] = mask
+          k = bands.high
+        bands[k].data[j] = bands[k].data[j] or mask
       n += W
     outs.setLen(0)
     for i in 0 ..< bands.high:
-      outs &= bands[i].data.compressSixel(bands[i].c - 15) & '$'
-    outs &= bands[^1].data.compressSixel(bands[^1].c - 15)
+      outs &= bands[i].compressSixel()
+      outs &= '$'
+    outs &= bands[^1].compressSixel()
     if n >= H:
       outs &= ST
     else: