Testing of the @doc-json output
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
+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
+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
tessera-viz/test/dune
+69
tessera-viz/test/test_viz.ml
+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
]