Something went wrong. Try again.
fork of zigimg/zigimg, ported to zig 0.17
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229const color = @import("color.zig");const std = @import("std");
const OctTreeQuantizer = @This();
const MaxDepth = 8;
pub const Error = std.mem.Allocator.Error || error{ InvalidColorIndex, ColorNotFound };
root_node: Node = .{},levels: [MaxDepth]?*Node = @splat(null),area_allocator: std.heap.ArenaAllocator,
pub fn init(allocator: std.mem.Allocator) OctTreeQuantizer { var result = OctTreeQuantizer{ .area_allocator = std.heap.ArenaAllocator.init(allocator), }; result.root_node.init(0, &result); return result;}
pub fn deinit(self: *OctTreeQuantizer) void { self.area_allocator.deinit();}
pub fn allocateNode(self: *OctTreeQuantizer) Error!*Node { return try self.area_allocator.allocator().create(Node);}
pub fn addLevelNode(self: *OctTreeQuantizer, level: i32, node: *Node) void { node.level_next = self.levels[@intCast(level)]; self.levels[@intCast(level)] = node;}
pub fn addColor(self: *OctTreeQuantizer, color_value: anytype) Error!void { try self.root_node.addColor(color_value, 0, self);}
pub fn getPaletteIndex(self: OctTreeQuantizer, color_value: anytype) Error!usize { return try self.root_node.getPaletteIndex(color_value, 0);}
pub fn makePalette(self: *OctTreeQuantizer, color_count: u32, palette: []color.Rgba32) []color.Rgba32 { var leaf_count = self.root_node.countLeafNodes();
var level: u8 = MaxDepth - 1; while (level > 0) : (level -= 1) { var node_it = self.levels[level];
while (node_it) |node| { leaf_count -= @intCast(node.removeLeaves()); if (leaf_count <= color_count) { break; } node_it = node.level_next; }
if (leaf_count <= color_count) { break; } }
var make_palette_context = MakePaletteContext{ .palette = palette, .color_count = color_count }; self.root_node.makePalette(&make_palette_context);
return palette[0..make_palette_context.palette_index];}
fn anyColorToRgb24(color_value: anytype) color.Rgb24 { const T = @TypeOf(color_value);
if (T == color.Rgb24) { return color_value; }
const has_alpha_type = @hasField(T, "a"); if (has_alpha_type) { const premultiplied_alpha = color_value.to.premultipliedAlpha();
return color.Rgb24.from.u32Rgba(premultiplied_alpha.to.u32Rgba()); } else { return color.Rgb24.from.u32Rgb(color_value.to.u32Rgb()); }}
const MakePaletteContext = struct { palette: []color.Rgba32, palette_index: u32 = 0, color_count: u32 = 0,};
const Node = struct { red: u32 = 0, green: u32 = 0, blue: u32 = 0, reference_count: u32 = 0, palette_index: u32 = 0, children: [8]?*Node = @splat(null), level_next: ?*Node = null,
pub fn init(self: *Node, level: i32, parent: *OctTreeQuantizer) void { self.* = Node{};
if (level < (MaxDepth - 1)) { parent.addLevelNode(level, self); } }
pub fn isLeaf(self: Node) bool { return self.reference_count > 0; }
pub fn getColor(self: Node) color.Rgba32 { return color.Rgba32.from.rgb(@intCast(self.red / self.reference_count), @intCast(self.green / self.reference_count), @intCast(self.blue / self.reference_count)); }
pub fn addColor(self: *Node, source_color: anytype, level: i32, parent: *OctTreeQuantizer) Error!void { if (level >= MaxDepth) { const color_value = anyColorToRgb24(source_color);
self.red += color_value.r; self.green += color_value.g; self.blue += color_value.b; self.reference_count += 1; return; }
const index = getColorIndex(source_color, level); if (index >= self.children.len) { return Error.InvalidColorIndex; }
if (self.children[index]) |child| { try child.addColor(source_color, level + 1, parent); } else { var new_node = try parent.allocateNode(); new_node.init(level, parent); try new_node.addColor(source_color, level + 1, parent); self.children[index] = new_node; } }
pub fn getPaletteIndex(self: Node, source_color: anytype, level: i32) Error!usize { if (self.isLeaf()) { return self.palette_index; }
const index = getColorIndex(source_color, level);
if (self.children[index]) |child| { return try child.getPaletteIndex(source_color, level + 1); } else { for (self.children) |child_opt| { if (child_opt) |child| { return try child.getPaletteIndex(source_color, level + 1); } } }
return Error.ColorNotFound; }
pub fn countLeafNodes(self: Node) usize { if (self.isLeaf()) { return 1; }
var count: usize = 0; for (self.children) |child_opt| { if (child_opt) |child| { count += child.countLeafNodes(); } }
return count; }
pub fn makePalette(self: *Node, context: *MakePaletteContext) void { if (self.isLeaf()) { if (context.palette_index >= context.color_count) { return; }
context.palette[context.palette_index] = self.getColor(); self.palette_index = context.palette_index; context.palette_index += 1; }
for (self.children) |child_opt| { if (child_opt) |child| { child.makePalette(context); } } }
pub fn removeLeaves(self: *Node) i32 { var result: i32 = 0; for (self.children, 0..) |child_opt, index| { if (child_opt) |child| { self.red +%= child.red; self.green +%= child.green; self.blue +%= child.blue; self.reference_count +%= child.reference_count; result += 1; self.children[index] = null; } } return result - 1; }
inline fn getColorIndex(source_color: anytype, level: i32) usize { const color_value = anyColorToRgb24(source_color);
var index: usize = 0; const mask = @as(u8, 0b10000000) >> @intCast(level); if (color_value.r & mask != 0) { index |= 0b100; } if (color_value.g & mask != 0) { index |= 0b010; } if (color_value.b & mask != 0) { index |= 0b001; } return index; }};