diff --git a/Project.toml b/Project.toml index 232241b..961a944 100644 --- a/Project.toml +++ b/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" diff --git a/src/ColorQuantization.jl b/src/ColorQuantization.jl index 3c59bfd..3896568 100644 --- a/src/ColorQuantization.jl +++ b/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 diff --git a/src/octree.jl b/src/octree.jl new file mode 100644 index 0000000..695b22e --- /dev/null +++ b/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 +