diff options
Diffstat (limited to 'libpanto/src/image.zig')
| -rw-r--r-- | libpanto/src/image.zig | 358 |
1 files changed, 0 insertions, 358 deletions
diff --git a/libpanto/src/image.zig b/libpanto/src/image.zig deleted file mode 100644 index 26474d1..0000000 --- a/libpanto/src/image.zig +++ /dev/null @@ -1,358 +0,0 @@ -//! Native image processing for tool-returned attachments. -//! -//! Two responsibilities: -//! -//! 1. `detectCodec` — identify an attachment's codec from its leading -//! bytes (magic numbers), not its file extension. -//! 2. `maybeResize` — bound raster images to `max_dim` on each side so a -//! single screenshot can't blow out the model's context. PDFs and -//! already-small images pass through untouched. -//! -//! Raster codecs go through the vendored stb single-header trio -//! (decode -> Mitchell resize -> re-encode in the *same* codec). WEBP is -//! decode-only (jebp), so a resized WEBP is re-encoded as JPEG. - -const std = @import("std"); -const Allocator = std.mem.Allocator; - -const c = @cImport({ - @cInclude("stb_image.h"); - @cInclude("stb_image_resize2.h"); - @cInclude("stb_image_write.h"); - @cInclude("jebp.h"); -}); - -/// Longest side (px) allowed before we resize. pi uses 2000x2000. -pub const max_dim: u32 = 2000; -/// JPEG quality used when re-encoding (WEBP path, and JPEG inputs). -const jpeg_quality: c_int = 80; - -pub const Codec = enum { png, jpeg, gif, bmp, webp, pdf }; - -/// The MIME type string for a codec (static; do not free). -pub fn mediaTypeForCodec(codec: Codec) []const u8 { - return switch (codec) { - .png => "image/png", - .jpeg => "image/jpeg", - .gif => "image/gif", - .bmp => "image/bmp", - .webp => "image/webp", - .pdf => "application/pdf", - }; -} - -/// Detect a supported codec from leading bytes (magic numbers). -pub fn detectCodec(bytes: []const u8) ?Codec { - if (bytes.len >= 8 and std.mem.eql(u8, bytes[0..8], &.{ 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A })) - return .png; - if (bytes.len >= 3 and bytes[0] == 0xFF and bytes[1] == 0xD8 and bytes[2] == 0xFF) - return .jpeg; - if (bytes.len >= 6 and (std.mem.eql(u8, bytes[0..6], "GIF87a") or std.mem.eql(u8, bytes[0..6], "GIF89a"))) - return .gif; - if (bytes.len >= 2 and bytes[0] == 'B' and bytes[1] == 'M') - return .bmp; - // WEBP: "RIFF"????"WEBP" - if (bytes.len >= 12 and std.mem.eql(u8, bytes[0..4], "RIFF") and std.mem.eql(u8, bytes[8..12], "WEBP")) - return .webp; - if (bytes.len >= 5 and std.mem.eql(u8, bytes[0..5], "%PDF-")) - return .pdf; - return null; -} - -/// Map a media-type string back to a codec (for callers that already have -/// the MIME string). Returns null for unsupported types. -pub fn codecForMediaType(media_type: []const u8) ?Codec { - if (std.mem.eql(u8, media_type, "image/png")) return .png; - if (std.mem.eql(u8, media_type, "image/jpeg")) return .jpeg; - if (std.mem.eql(u8, media_type, "image/gif")) return .gif; - if (std.mem.eql(u8, media_type, "image/bmp")) return .bmp; - if (std.mem.eql(u8, media_type, "image/webp")) return .webp; - if (std.mem.eql(u8, media_type, "application/pdf")) return .pdf; - return null; -} - -/// The result of `maybeResize`. `media_type` is the MIME type of `data` -/// (may differ from the input when a WEBP was re-encoded as JPEG). -/// `data` is always an owned slice the caller must free. -pub const Processed = struct { - media_type: []const u8, // static string, do not free - data: []u8, // owned by `allocator` -}; - -/// Full attachment pipeline for a tool-returned media part: resolve the -/// media type (detecting from magic bytes when `hint` is null), then -/// resize. Returns owned raw bytes + the resolved media type. -/// -/// Errors `error.UnknownMediaType` when neither the hint nor magic-byte -/// detection recognizes the bytes — the caller decides how to surface -/// that (e.g. drop the attachment, or fall back to text). -pub fn process(allocator: Allocator, bytes: []const u8, hint: ?[]const u8) !Processed { - const codec = blk: { - if (hint) |h| { - if (codecForMediaType(h)) |hinted| break :blk hinted; - } - break :blk detectCodec(bytes) orelse return error.UnknownMediaType; - }; - return maybeResize(allocator, bytes, codec); -} - -/// Resize `bytes` so neither dimension exceeds `max_dim`, preserving the -/// input codec where possible. Returns an owned copy of the (possibly -/// unchanged) bytes plus the resulting media type. -/// -/// - PDF: returned verbatim (a copy), media type unchanged. -/// - raster <= max_dim on both sides: returned verbatim (a copy) — we -/// skip the decode/encode round-trip to avoid quality loss + CPU. -/// - stb-supported raster larger than max_dim: decode -> resize -> same -/// codec. -/// - WEBP larger than max_dim: jebp decode -> resize -> JPEG. -pub fn maybeResize(allocator: Allocator, bytes: []const u8, codec: Codec) !Processed { - const media_type = mediaTypeForCodec(codec); - - if (codec == .pdf) - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - - if (codec == .webp) return resizeWebp(allocator, bytes, media_type); - return resizeStb(allocator, bytes, media_type, codec); -} - -/// Probe a raster image's dimensions without fully decoding. Returns null -/// on failure. -fn probeDims(bytes: []const u8) ?struct { w: u32, h: u32 } { - var w: c_int = 0; - var h: c_int = 0; - var comp: c_int = 0; - const ok = c.stbi_info_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp); - if (ok == 0 or w <= 0 or h <= 0) return null; - return .{ .w = @intCast(w), .h = @intCast(h) }; -} - -/// Compute target dimensions that fit within `max_dim` x `max_dim` while -/// preserving aspect ratio. Returns null when no resize is needed. -fn targetDims(w: u32, h: u32) ?struct { w: u32, h: u32 } { - if (w <= max_dim and h <= max_dim) return null; - const wf: f64 = @floatFromInt(w); - const hf: f64 = @floatFromInt(h); - const scale = @min(@as(f64, @floatFromInt(max_dim)) / wf, @as(f64, @floatFromInt(max_dim)) / hf); - const nw: u32 = @max(1, @as(u32, @intFromFloat(@round(wf * scale)))); - const nh: u32 = @max(1, @as(u32, @intFromFloat(@round(hf * scale)))); - return .{ .w = nw, .h = nh }; -} - -const StbWriteCtx = struct { - list: *std.ArrayList(u8), - allocator: Allocator, - failed: bool = false, -}; - -fn stbWriteCb(ctx_opaque: ?*anyopaque, data: ?*anyopaque, size: c_int) callconv(.c) void { - const ctx: *StbWriteCtx = @ptrCast(@alignCast(ctx_opaque.?)); - if (ctx.failed or size <= 0) return; - const bytes: [*]const u8 = @ptrCast(data.?); - ctx.list.appendSlice(ctx.allocator, bytes[0..@intCast(size)]) catch { - ctx.failed = true; - }; -} - -fn resizeStb(allocator: Allocator, bytes: []const u8, media_type: []const u8, codec: Codec) !Processed { - const dims = probeDims(bytes) orelse - // Can't parse it; pass through rather than fail the read. - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - - const target = targetDims(dims.w, dims.h) orelse - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - - // Decode to RGBA (4 channels) for a uniform resize path. - var w: c_int = 0; - var h: c_int = 0; - var comp: c_int = 0; - const pixels = c.stbi_load_from_memory(bytes.ptr, @intCast(bytes.len), &w, &h, &comp, 4); - if (pixels == null) - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - defer c.stbi_image_free(pixels); - - const out_pixels = try allocator.alloc(u8, @as(usize, target.w) * @as(usize, target.h) * 4); - defer allocator.free(out_pixels); - - const res = c.stbir_resize_uint8_srgb( - pixels, - w, - h, - 0, - out_pixels.ptr, - @intCast(target.w), - @intCast(target.h), - 0, - c.STBIR_RGBA, - ); - if (res == null) - return error.ResizeFailed; - - var out: std.ArrayList(u8) = .empty; - errdefer out.deinit(allocator); - var ctx = StbWriteCtx{ .list = &out, .allocator = allocator }; - - const tw: c_int = @intCast(target.w); - const th: c_int = @intCast(target.h); - const ok = switch (codec) { - .png => c.stbi_write_png_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, tw * 4), - .bmp => c.stbi_write_bmp_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr), - // stb has no GIF encoder; re-encode resized GIFs as PNG (lossless). - .gif => c.stbi_write_png_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, tw * 4), - .jpeg => c.stbi_write_jpg_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, jpeg_quality), - else => unreachable, - }; - if (ok == 0 or ctx.failed) return error.EncodeFailed; - - const result_media: []const u8 = switch (codec) { - .gif => "image/png", // re-encoded - else => media_type, - }; - return .{ .media_type = result_media, .data = try out.toOwnedSlice(allocator) }; -} - -fn resizeWebp(allocator: Allocator, bytes: []const u8, media_type: []const u8) !Processed { - var img: c.jebp_image_t = std.mem.zeroes(c.jebp_image_t); - // Peek at the header first to learn dimensions cheaply. - if (c.jebp_decode_size(&img, bytes.len, bytes.ptr) != c.JEBP_OK) - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - - const w: u32 = @intCast(img.width); - const h: u32 = @intCast(img.height); - const target = targetDims(w, h) orelse - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - - // Full decode to RGBA (jebp_color_t is 4 contiguous bytes per pixel). - if (c.jebp_decode(&img, bytes.len, bytes.ptr) != c.JEBP_OK) - return .{ .media_type = media_type, .data = try allocator.dupe(u8, bytes) }; - defer c.jebp_free_image(&img); - - const src: [*]const u8 = @ptrCast(img.pixels); - const out_pixels = try allocator.alloc(u8, @as(usize, target.w) * @as(usize, target.h) * 4); - defer allocator.free(out_pixels); - - const res = c.stbir_resize_uint8_srgb( - src, - @intCast(w), - @intCast(h), - 0, - out_pixels.ptr, - @intCast(target.w), - @intCast(target.h), - 0, - c.STBIR_RGBA, - ); - if (res == null) return error.ResizeFailed; - - // TODO: when the source WEBP has an alpha layer, re-encoding to JPEG - // flattens transparency, which can look wrong for screenshots and - // diagrams. Consider re-encoding to PNG when alpha is present. For now - // we always emit JPEG: there is no small single-header WEBP encoder, - // and token size matters more than fidelity for LLM input. - var out: std.ArrayList(u8) = .empty; - errdefer out.deinit(allocator); - var ctx = StbWriteCtx{ .list = &out, .allocator = allocator }; - const tw: c_int = @intCast(target.w); - const th: c_int = @intCast(target.h); - const ok = c.stbi_write_jpg_to_func(stbWriteCb, &ctx, tw, th, 4, out_pixels.ptr, jpeg_quality); - if (ok == 0 or ctx.failed) return error.EncodeFailed; - - return .{ .media_type = "image/jpeg", .data = try out.toOwnedSlice(allocator) }; -} - -// ----------------------------------------------------------------------------- -// Tests -// ----------------------------------------------------------------------------- - -const testing = std.testing; - -test "detectCodec + mediaTypeForCodec - magic bytes" { - try testing.expectEqualStrings("image/png", mediaTypeForCodec(detectCodec(&.{ 0x89, 'P', 'N', 'G', 0x0D, 0x0A, 0x1A, 0x0A }).?)); - try testing.expectEqualStrings("image/jpeg", mediaTypeForCodec(detectCodec(&.{ 0xFF, 0xD8, 0xFF, 0xE0 }).?)); - try testing.expectEqualStrings("image/gif", mediaTypeForCodec(detectCodec("GIF89a....").?)); - try testing.expectEqualStrings("image/bmp", mediaTypeForCodec(detectCodec("BM....").?)); - try testing.expectEqualStrings("application/pdf", mediaTypeForCodec(detectCodec("%PDF-1.7").?)); - const webp = "RIFF" ++ &[_]u8{ 0, 0, 0, 0 } ++ "WEBP"; - try testing.expectEqualStrings("image/webp", mediaTypeForCodec(detectCodec(webp).?)); - try testing.expect(detectCodec("not an image") == null); - try testing.expect(detectCodec(&.{0x89}) == null); -} - -test "targetDims - skip when small, scale when large preserving aspect" { - try testing.expect(targetDims(100, 100) == null); - try testing.expect(targetDims(max_dim, max_dim) == null); - const t = targetDims(4000, 2000).?; - try testing.expectEqual(@as(u32, 2000), t.w); - try testing.expectEqual(@as(u32, 1000), t.h); - const t2 = targetDims(1000, 8000).?; - try testing.expectEqual(@as(u32, 250), t2.w); - try testing.expectEqual(@as(u32, 2000), t2.h); -} - -test "process - detects type from raw bytes when hint absent" { - const a = testing.allocator; - // A tiny PNG (header + IHDR enough for stbi_info) — but simplest is to - // round-trip an stb-encoded small PNG and feed it with no hint. - const w: c_int = 4; - const h: c_int = 4; - var px: [4 * 4 * 4]u8 = undefined; - @memset(&px, 0x40); - var png: std.ArrayList(u8) = .empty; - defer png.deinit(a); - var ctx = StbWriteCtx{ .list = &png, .allocator = a }; - try testing.expect(c.stbi_write_png_to_func(stbWriteCb, &ctx, w, h, 4, &px, w * 4) != 0); - - const out = try process(a, png.items, null); - defer a.free(out.data); - try testing.expectEqualStrings("image/png", out.media_type); - - // Unknown bytes -> error so the caller can fall back / drop. - try testing.expectError(error.UnknownMediaType, process(a, "totally not an image", null)); -} - -test "maybeResize - PDF passes through unchanged" { - const a = testing.allocator; - const pdf = "%PDF-1.7\nfake pdf body"; - const out = try maybeResize(a, pdf, .pdf); - defer a.free(out.data); - try testing.expectEqualStrings("application/pdf", out.media_type); - try testing.expectEqualStrings(pdf, out.data); -} - -test "maybeResize - small PNG round-trips, large PNG shrinks and stays PNG" { - const a = testing.allocator; - - // Build a small (8x8) RGBA PNG via stb and confirm pass-through. - const small_w: c_int = 8; - const small_h: c_int = 8; - var small_px: [8 * 8 * 4]u8 = undefined; - for (&small_px, 0..) |*b, i| b.* = @truncate(i); - var small_png: std.ArrayList(u8) = .empty; - defer small_png.deinit(a); - var sctx = StbWriteCtx{ .list = &small_png, .allocator = a }; - try testing.expect(c.stbi_write_png_to_func(stbWriteCb, &sctx, small_w, small_h, 4, &small_px, small_w * 4) != 0); - - const small_out = try maybeResize(a, small_png.items, .png); - defer a.free(small_out.data); - try testing.expectEqualStrings("image/png", small_out.media_type); - // Small image is returned verbatim (byte-identical copy). - try testing.expectEqualSlices(u8, small_png.items, small_out.data); - - // Build a large (2400x100) PNG and confirm it shrinks to <= max_dim. - const big_w: c_int = 2400; - const big_h: c_int = 100; - const big_px = try a.alloc(u8, @as(usize, @intCast(big_w * big_h * 4))); - defer a.free(big_px); - @memset(big_px, 0x7F); - var big_png: std.ArrayList(u8) = .empty; - defer big_png.deinit(a); - var bctx = StbWriteCtx{ .list = &big_png, .allocator = a }; - try testing.expect(c.stbi_write_png_to_func(stbWriteCb, &bctx, big_w, big_h, 4, big_px.ptr, big_w * 4) != 0); - - const big_out = try maybeResize(a, big_png.items, .png); - defer a.free(big_out.data); - try testing.expectEqualStrings("image/png", big_out.media_type); - const dims = probeDims(big_out.data).?; - try testing.expectEqual(@as(u32, 2000), dims.w); - try testing.expect(dims.w <= max_dim and dims.h <= max_dim); -} |
