Octree quantization

Project.toml

@@ -9,10 +9,14 @@ Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 ImageBase = "c817782e-172a-44cc-b673-b171935fbb9e"
 LazyModules = "8cdb02fc-e678-4876-92c5-9defec4f444e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RegionTrees = "dee08c22-ab7f-5625-9660-a9af2021b33f"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 Clustering = "0.14.3"
 Colors = "0.12"
 ImageBase = "0.1"
 LazyModules = "0.3"
+RegionTrees = "0.3.2"
+StaticArrays = "0.12, 1"
 julia = "1.6"

src/ColorQuantization.jl

@@ -1,10 +1,13 @@
 module ColorQuantization
 
 using Colors
-using ImageBase: FixedPoint, floattype, FixedPointNumbers.rawtype
+using ImageBase: FixedPoint, floattype, FixedPointNumbers.rawtype, FixedPointNumbers.N0f8
 using ImageBase: channelview, colorview, restrict
 using Random: AbstractRNG, GLOBAL_RNG
 using LazyModules: @lazy
+using RegionTrees
+using StaticArrays
+
 #! format: off
 @lazy import Clustering = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5"
 #! format: on
@@ -15,8 +18,9 @@ include("api.jl")
 include("utils.jl")
 include("uniform.jl")
 include("clustering.jl") # lazily loaded
+include("octree.jl")
 
 export AbstractColorQuantizer, quantize
-export UniformQuantization, KMeansQuantization
+export UniformQuantization, KMeansQuantization, OctreeQuantization
 
 end

src/octree.jl

@@ -0,0 +1,105 @@
+
+struct OctreeQuantization <: AbstractColorQuantizer
+    numcolors::Int
+    function OctreeQuantization(numcolors::Int=256; kwargs...)
+        return new(numcolors)
+    end
+end
+
+function (alg::OctreeQuantization)(img::AbstractArray)
+    return octreequantization!(img; numcolors=alg.numcolors)
+end
+
+function octreequantization!(img; numcolors = 256, precheck::Bool = false)
+    # ensure the img is in RGB colorspace
+    if (eltype(img) != RGB{N0f8})
+        error("Octree Algorithm requires img to be in RGB colorspace")
+    end
+
+    # checks if image has more colors than in numcolors
+    if precheck == true
+        unumcolors = length(unique(img))
+        # @show unumcolors
+        if unumcolors <= numcolors
+            @debug "Image has $unumcolors unique colors"
+            return unique(img)
+        end
+    end
+
+    # step 1: creating the octree
+    root = Cell(SVector(0.0, 0.0, 0.0), SVector(1.0, 1.0, 1.0), [0, [], RGB{N0f8}.(0.0, 0.0, 0.0), 0])
+    cases = map(p -> [0, Vector{Int}([]), RGB{N0f8}.(0.0, 0.0, 0.0), 1], 1:8)
+    split!(root, cases)
+    inds = collect(1:length(img))
+
+    function putin(c::Cell, i::Int, level::Int)
+        if (level == 8)
+            push!(c.data[2], i)
+            c.data[3] = img[i]
+            return
+        else
+            if (isleaf(c) == true && level <= 8)
+                cadata = map(p -> [0, Vector{Int}([]), RGB{N0f8}.(0.0, 0.0, 0.0), level + 1], 1:8)
+                split!(c, cadata)
+            end
+            r, g, b = map(p -> bitstring(UInt8(p * 255)), channelview([img[i]]))
+            rgb = r[level] * g[level] * b[level]
+            rgb = parse(Int, rgb, base=2) + 1
+            c.children[rgb].data[1] += 1
+            putin(c.children[rgb], i, level + 1)
+        end
+    end
+
+    level = 1
+    root.data[1] = length(inds)
+
+    # build the tree
+    for i in inds
+        r, g, b = map(p -> bitstring(UInt8(p * 255)), channelview([img[i]]))
+        rgb = r[level] * g[level] * b[level]
+        rgb = parse(Int, rgb, base=2) + 1
+        root.children[rgb].data[1] += 1
+        putin(root.children[rgb], i, level)
+    end
+
+    # step 2: reducing tree to a certain number of colors
+    # there is scope for improvements in allleaves as it's found again n again
+    leafs = [p for p in allleaves(root)]
+    filter!(p -> !iszero(p.data[1]), leafs)
+    tobe_reduced = leafs[1]
+
+    while (length(leafs) > numcolors)
+        parents = unique([parent(p) for p in leafs])
+        parents = sort(parents; by = c -> c.data[1])
+        tobe_reduced = parents[1]
+
+        for i = 1:8
+            append!(tobe_reduced.data[2], tobe_reduced.children[i].data[2])
+            tobe_reduced.data[3] +=
+                tobe_reduced.children[i].data[3] * tobe_reduced.children[i].data[1]
+        end
+
+        tobe_reduced.data[3] /= tobe_reduced.data[1]
+        filter!(!in(tobe_reduced.children),leafs)
+        push!(leafs, tobe_reduced)
+        tobe_reduced.children = nothing
+    end
+
+    # step 3: palette formation  and quantisation now
+    da = [p.data for p in leafs]
+    for i in da
+        for j in i[2]
+            img[j] = i[3]
+        end
+    end
+
+    colors = [p[3] for p in da]
+    return colors
+end
+
+function octreequantization(img; kwargs...)
+    img_copy = deepcopy(img)
+    palette = octreequantization!(img_copy; kwargs...)
+    return img_copy, palette
+end
+