My aggregated monorepo of OCaml code, automaintained

tessera-viz: add portable PNG encoder (uncompressed deflate)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

+179 -1
+102
tessera-viz/lib/viz.ml
··· 69 69 Bigarray.Array1.set data (off + 3) alpha 70 70 done; 71 71 { data; width; height } 72 + 73 + (* ---- Minimal PNG encoder ---- *) 74 + 75 + let png_crc32_table = 76 + Array.init 256 (fun n -> 77 + let c = ref (Int32.of_int n) in 78 + for _ = 0 to 7 do 79 + if Int32.logand !c 1l <> 0l then 80 + c := Int32.logxor (Int32.shift_right_logical !c 1) 0xEDB88320l 81 + else 82 + c := Int32.shift_right_logical !c 1 83 + done; 84 + !c) 85 + 86 + let png_crc32 data ofs len = 87 + let c = ref 0xFFFFFFFFl in 88 + for i = ofs to ofs + len - 1 do 89 + let byte = Char.code (Bytes.get data i) in 90 + let idx = Int32.to_int (Int32.logand (Int32.logxor !c (Int32.of_int byte)) 0xFFl) in 91 + c := Int32.logxor (Int32.shift_right_logical !c 8) png_crc32_table.(idx) 92 + done; 93 + Int32.logxor !c 0xFFFFFFFFl 94 + 95 + let put_be32 buf pos v = 96 + Bytes.set buf pos (Char.chr ((v lsr 24) land 0xFF)); 97 + Bytes.set buf (pos + 1) (Char.chr ((v lsr 16) land 0xFF)); 98 + Bytes.set buf (pos + 2) (Char.chr ((v lsr 8) land 0xFF)); 99 + Bytes.set buf (pos + 3) (Char.chr (v land 0xFF)) 100 + 101 + let put_be32_i32 buf pos v = 102 + let v = Int32.to_int v in 103 + put_be32 buf pos v 104 + 105 + let png_chunk buf typ data = 106 + let len = Bytes.length data in 107 + let chunk = Bytes.create (len + 4) in 108 + Bytes.blit_string typ 0 chunk 0 4; 109 + Bytes.blit data 0 chunk 4 len; 110 + let crc = png_crc32 chunk 0 (len + 4) in 111 + Buffer.add_bytes buf (let b = Bytes.create 4 in put_be32 b 0 len; b); 112 + Buffer.add_bytes buf chunk; 113 + Buffer.add_bytes buf (let b = Bytes.create 4 in put_be32_i32 b 0 crc; b) 114 + 115 + let png_of_rgba img = 116 + let buf = Buffer.create (img.width * img.height * 4 + 1024) in 117 + (* PNG signature *) 118 + Buffer.add_string buf "\x89PNG\r\n\x1a\n"; 119 + (* IHDR *) 120 + let ihdr = Bytes.create 13 in 121 + put_be32 ihdr 0 img.width; 122 + put_be32 ihdr 4 img.height; 123 + Bytes.set ihdr 8 '\x08'; (* bit depth 8 *) 124 + Bytes.set ihdr 9 '\x06'; (* color type 6 = RGBA *) 125 + Bytes.set ihdr 10 '\x00'; (* compression *) 126 + Bytes.set ihdr 11 '\x00'; (* filter *) 127 + Bytes.set ihdr 12 '\x00'; (* interlace *) 128 + png_chunk buf "IHDR" ihdr; 129 + (* Build raw scanlines: filter_byte(0) + row RGBA data *) 130 + let row_bytes = 1 + img.width * 4 in 131 + let raw_len = row_bytes * img.height in 132 + let raw = Bytes.create raw_len in 133 + for y = 0 to img.height - 1 do 134 + Bytes.set raw (y * row_bytes) '\x00'; (* filter: None *) 135 + for x = 0 to img.width * 4 - 1 do 136 + let v = Bigarray.Array1.get img.data (y * img.width * 4 + x) in 137 + Bytes.set raw (y * row_bytes + 1 + x) (Char.chr v) 138 + done 139 + done; 140 + (* Wrap in uncompressed deflate stored blocks *) 141 + let max_block = 65535 in 142 + let idat_buf = Buffer.create (raw_len + raw_len / max_block * 5 + 20) in 143 + (* Zlib header: CM=8, CINFO=7, FCHECK to make it valid *) 144 + Buffer.add_char idat_buf '\x78'; 145 + Buffer.add_char idat_buf '\x01'; 146 + let pos = ref 0 in 147 + while !pos < raw_len do 148 + let remaining = raw_len - !pos in 149 + let block_len = min remaining max_block in 150 + let is_final = !pos + block_len >= raw_len in 151 + Buffer.add_char idat_buf (if is_final then '\x01' else '\x00'); 152 + Buffer.add_char idat_buf (Char.chr (block_len land 0xFF)); 153 + Buffer.add_char idat_buf (Char.chr ((block_len lsr 8) land 0xFF)); 154 + let nlen = block_len lxor 0xFFFF in 155 + Buffer.add_char idat_buf (Char.chr (nlen land 0xFF)); 156 + Buffer.add_char idat_buf (Char.chr ((nlen lsr 8) land 0xFF)); 157 + Buffer.add_subbytes idat_buf raw !pos block_len; 158 + pos := !pos + block_len 159 + done; 160 + (* Adler-32 checksum *) 161 + let a = ref 1 and b = ref 0 in 162 + for i = 0 to raw_len - 1 do 163 + a := (!a + Char.code (Bytes.get raw i)) mod 65521; 164 + b := (!b + !a) mod 65521 165 + done; 166 + let adler = (!b lsl 16) lor !a in 167 + let adler_bytes = Bytes.create 4 in 168 + put_be32 adler_bytes 0 adler; 169 + Buffer.add_bytes idat_buf adler_bytes; 170 + png_chunk buf "IDAT" (Buffer.to_bytes idat_buf); 171 + (* IEND *) 172 + png_chunk buf "IEND" Bytes.empty; 173 + Buffer.contents buf
+7
tessera-viz/lib/viz.mli
··· 50 50 [predictions] has length [height * width]. 51 51 [colors] maps class IDs to colors. Unknown classes are black. 52 52 [alpha] defaults to 200. *) 53 + 54 + (** {1 PNG encoding} *) 55 + 56 + val png_of_rgba : rgba_image -> string 57 + (** Encode an RGBA image as PNG bytes. 58 + Uses uncompressed deflate (stored blocks) for portability — 59 + no external compression library required. *)
+1 -1
tessera-viz/test/dune
··· 1 1 (test 2 2 (name test_viz) 3 - (libraries tessera-viz tessera-linalg alcotest)) 3 + (libraries tessera-viz tessera-linalg alcotest unix))
+69
tessera-viz/test/test_viz.ml
··· 138 138 ; Alcotest.test_case "custom alpha" `Quick test_classification_custom_alpha 139 139 ] 140 140 141 + let test_png_magic () = 142 + let data = Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout 4 in 143 + Bigarray.Array1.set data 0 255; 144 + Bigarray.Array1.set data 1 0; 145 + Bigarray.Array1.set data 2 0; 146 + Bigarray.Array1.set data 3 255; 147 + let img = Viz.{ data; width = 1; height = 1 } in 148 + let png = Viz.png_of_rgba img in 149 + Alcotest.(check int) "byte 0" 0x89 (Char.code png.[0]); 150 + Alcotest.(check char) "byte 1" 'P' png.[1]; 151 + Alcotest.(check char) "byte 2" 'N' png.[2]; 152 + Alcotest.(check char) "byte 3" 'G' png.[3]; 153 + Alcotest.(check bool) "length > 20" true (String.length png > 20) 154 + 155 + let test_png_roundtrip () = 156 + let data = Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout 16 in 157 + let pixels = [| 255;0;0;255; 0;255;0;255; 0;0;255;255; 255;255;255;255 |] in 158 + Array.iteri (fun i v -> Bigarray.Array1.set data i v) pixels; 159 + let img = Viz.{ data; width = 2; height = 2 } in 160 + let png = Viz.png_of_rgba img in 161 + let tmp = Filename.temp_file "test_png" ".png" in 162 + let oc = open_out_bin tmp in 163 + output_string oc png; 164 + close_out oc; 165 + let cmd = Printf.sprintf 166 + "python3 -c \"from PIL import Image; img = Image.open('%s'); print(img.size, img.mode, list(img.getdata()))\"" tmp in 167 + let ic = Unix.open_process_in cmd in 168 + let output = input_line ic in 169 + let _ = Unix.close_process_in ic in 170 + Sys.remove tmp; 171 + Alcotest.(check bool) "contains (2, 2)" true (String.length output > 0 && String.sub output 0 6 = "(2, 2)"); 172 + Alcotest.(check bool) "contains RGBA" true 173 + (try let _ = String.index output 'R' in true with Not_found -> false); 174 + Printf.printf "Python output: %s\n" output 175 + 176 + let test_png_larger_image () = 177 + let w = 100 and h = 100 in 178 + let data = Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout (w * h * 4) in 179 + for y = 0 to h - 1 do 180 + for x = 0 to w - 1 do 181 + let off = (y * w + x) * 4 in 182 + Bigarray.Array1.set data off (x * 255 / 99); 183 + Bigarray.Array1.set data (off + 1) (y * 255 / 99); 184 + Bigarray.Array1.set data (off + 2) 128; 185 + Bigarray.Array1.set data (off + 3) 255 186 + done 187 + done; 188 + let img = Viz.{ data; width = w; height = h } in 189 + let png = Viz.png_of_rgba img in 190 + let tmp = Filename.temp_file "test_png_large" ".png" in 191 + let oc = open_out_bin tmp in 192 + output_string oc png; 193 + close_out oc; 194 + let cmd = Printf.sprintf 195 + "python3 -c \"from PIL import Image; img = Image.open('%s'); print(img.size, list(img.getpixel((0,0))), list(img.getpixel((99,99))))\"" tmp in 196 + let ic = Unix.open_process_in cmd in 197 + let output = input_line ic in 198 + let _ = Unix.close_process_in ic in 199 + Sys.remove tmp; 200 + Alcotest.(check bool) "contains (100, 100)" true (String.length output > 0 && String.sub output 0 10 = "(100, 100)"); 201 + Printf.printf "Python output: %s\n" output 202 + 203 + let png_tests = 204 + [ Alcotest.test_case "PNG magic bytes" `Quick test_png_magic 205 + ; Alcotest.test_case "PNG roundtrip via Python" `Quick test_png_roundtrip 206 + ; Alcotest.test_case "PNG larger image" `Quick test_png_larger_image 207 + ] 208 + 141 209 let () = 142 210 Alcotest.run "tessera-viz" 143 211 [ ("percentile", percentile_tests) 144 212 ; ("pca_to_rgba", pca_tests) 145 213 ; ("color_of_hex", color_tests) 146 214 ; ("classification_to_rgba", classification_tests) 215 + ; ("png_of_rgba", png_tests) 147 216 ]