1use super::lzw::LzwEncoder;
11use crate::error::{CodecError, CodecResult};
12use crate::frame::VideoFrame;
13use oximedia_core::PixelFormat;
14use std::collections::HashMap;
15use std::io::Write;
16
17const MAX_COLORS: usize = 256;
19
20const GIF89A_HEADER: &[u8] = b"GIF89a";
22
23const EXTENSION_INTRODUCER: u8 = 0x21;
25
26const IMAGE_SEPARATOR: u8 = 0x2C;
28
29const TRAILER: u8 = 0x3B;
31
32const GRAPHICS_CONTROL_LABEL: u8 = 0xF9;
34
35const APPLICATION_LABEL: u8 = 0xFF;
37
38#[allow(dead_code)]
40const DISPOSAL_NONE: u8 = 0;
41
42#[allow(dead_code)]
44const DISPOSAL_KEEP: u8 = 1;
45
46#[allow(dead_code)]
48const DISPOSAL_BACKGROUND: u8 = 2;
49
50#[allow(dead_code)]
52const DISPOSAL_PREVIOUS: u8 = 3;
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56pub enum DitheringMethod {
57 None,
59 FloydSteinberg,
61 Ordered,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub enum QuantizationMethod {
68 MedianCut,
70 Octree,
72}
73
74#[derive(Debug, Clone)]
76pub struct GifEncoderConfig {
77 pub colors: usize,
79 pub quantization: QuantizationMethod,
81 pub dithering: DitheringMethod,
83 pub transparent_index: Option<u8>,
85 pub loop_count: u16,
87}
88
89impl Default for GifEncoderConfig {
90 fn default() -> Self {
91 Self {
92 colors: 256,
93 quantization: QuantizationMethod::MedianCut,
94 dithering: DitheringMethod::None,
95 transparent_index: None,
96 loop_count: 0,
97 }
98 }
99}
100
101#[derive(Debug, Clone)]
103pub struct GifFrameConfig {
104 pub delay_time: u16,
106 pub disposal_method: u8,
108 pub left: u16,
110 pub top: u16,
112}
113
114impl Default for GifFrameConfig {
115 fn default() -> Self {
116 Self {
117 delay_time: 10, disposal_method: DISPOSAL_BACKGROUND,
119 left: 0,
120 top: 0,
121 }
122 }
123}
124
125pub struct GifEncoderState {
127 config: GifEncoderConfig,
129 width: u32,
131 height: u32,
133 output: Vec<u8>,
135 palette: Vec<u8>,
137}
138
139impl GifEncoderState {
140 pub fn new(width: u32, height: u32, config: GifEncoderConfig) -> CodecResult<Self> {
142 if width == 0 || height == 0 || width > 65535 || height > 65535 {
143 return Err(CodecError::InvalidParameter(format!(
144 "Invalid dimensions: {}x{}",
145 width, height
146 )));
147 }
148
149 if !(2..=256).contains(&config.colors) {
150 return Err(CodecError::InvalidParameter(format!(
151 "Invalid color count: {}",
152 config.colors
153 )));
154 }
155
156 Ok(Self {
157 config,
158 width,
159 height,
160 output: Vec::new(),
161 palette: Vec::new(),
162 })
163 }
164
165 pub fn encode(
171 &mut self,
172 frames: &[VideoFrame],
173 frame_configs: &[GifFrameConfig],
174 ) -> CodecResult<Vec<u8>> {
175 if frames.is_empty() {
176 return Err(CodecError::InvalidParameter("No frames to encode".into()));
177 }
178
179 if frames.len() != frame_configs.len() {
180 return Err(CodecError::InvalidParameter(
181 "Frame count mismatch with configs".into(),
182 ));
183 }
184
185 self.output.clear();
186
187 self.palette = self.generate_global_palette(frames)?;
189
190 self.write_header()?;
192
193 self.write_screen_descriptor()?;
195
196 let palette = self.palette.clone();
198 self.write_color_table(&palette)?;
199
200 if frames.len() > 1 {
202 self.write_netscape_extension()?;
203 }
204
205 for (frame, frame_config) in frames.iter().zip(frame_configs) {
207 self.write_frame(frame, frame_config)?;
208 }
209
210 self.output.write_all(&[TRAILER])?;
212
213 Ok(self.output.clone())
214 }
215
216 fn write_header(&mut self) -> CodecResult<()> {
218 self.output.write_all(GIF89A_HEADER)?;
219 Ok(())
220 }
221
222 fn write_screen_descriptor(&mut self) -> CodecResult<()> {
224 let width = self.width as u16;
225 let height = self.height as u16;
226
227 self.output.write_all(&width.to_le_bytes())?;
228 self.output.write_all(&height.to_le_bytes())?;
229
230 let global_color_table_flag = 1u8 << 7;
232 let color_resolution = 7u8 << 4; let sort_flag = 0u8;
234 let size = Self::color_table_size_field(self.palette.len() / 3);
235 let packed = global_color_table_flag | color_resolution | sort_flag | size;
236
237 self.output.write_all(&[packed])?;
238 self.output.write_all(&[0])?; self.output.write_all(&[0])?; Ok(())
242 }
243
244 fn write_color_table(&mut self, table: &[u8]) -> CodecResult<()> {
246 let size = Self::next_power_of_two(table.len() / 3) * 3;
247 self.output.write_all(table)?;
248
249 if table.len() < size {
251 self.output
252 .resize(self.output.len() + (size - table.len()), 0);
253 }
254
255 Ok(())
256 }
257
258 fn write_netscape_extension(&mut self) -> CodecResult<()> {
260 self.output.write_all(&[EXTENSION_INTRODUCER])?;
261 self.output.write_all(&[APPLICATION_LABEL])?;
262 self.output.write_all(&[11])?; self.output.write_all(b"NETSCAPE2.0")?;
264 self.output.write_all(&[3])?; self.output.write_all(&[1])?; self.output
267 .write_all(&self.config.loop_count.to_le_bytes())?;
268 self.output.write_all(&[0])?; Ok(())
271 }
272
273 fn write_frame(&mut self, frame: &VideoFrame, config: &GifFrameConfig) -> CodecResult<()> {
275 self.write_graphics_control_extension(config)?;
277
278 let rgba_data = self.frame_to_rgba(frame)?;
280 let indices = self.quantize_frame(&rgba_data)?;
281
282 self.write_image_descriptor(config)?;
284
285 self.write_image_data(&indices)?;
287
288 Ok(())
289 }
290
291 fn write_graphics_control_extension(&mut self, config: &GifFrameConfig) -> CodecResult<()> {
293 self.output.write_all(&[EXTENSION_INTRODUCER])?;
294 self.output.write_all(&[GRAPHICS_CONTROL_LABEL])?;
295 self.output.write_all(&[4])?; let disposal_method = (config.disposal_method & 0x07) << 2;
299 let user_input_flag = 0u8;
300 let transparency_flag = if self.config.transparent_index.is_some() {
301 1u8
302 } else {
303 0u8
304 };
305 let packed = disposal_method | user_input_flag | transparency_flag;
306
307 self.output.write_all(&[packed])?;
308 self.output.write_all(&config.delay_time.to_le_bytes())?;
309 self.output
310 .write_all(&[self.config.transparent_index.unwrap_or(0)])?;
311 self.output.write_all(&[0])?; Ok(())
314 }
315
316 fn write_image_descriptor(&mut self, config: &GifFrameConfig) -> CodecResult<()> {
318 self.output.write_all(&[IMAGE_SEPARATOR])?;
319 self.output.write_all(&config.left.to_le_bytes())?;
320 self.output.write_all(&config.top.to_le_bytes())?;
321 self.output.write_all(&(self.width as u16).to_le_bytes())?;
322 self.output.write_all(&(self.height as u16).to_le_bytes())?;
323
324 self.output.write_all(&[0])?;
326
327 Ok(())
328 }
329
330 fn write_image_data(&mut self, indices: &[u8]) -> CodecResult<()> {
332 let color_bits = Self::bits_needed(self.palette.len() / 3);
334 let min_code_size = color_bits.max(2);
335
336 self.output.write_all(&[min_code_size])?;
337
338 let mut encoder = LzwEncoder::new(min_code_size)?;
340 let compressed = encoder.compress(indices)?;
341
342 let mut offset = 0;
344 while offset < compressed.len() {
345 let block_size = (compressed.len() - offset).min(255);
346 self.output.write_all(&[block_size as u8])?;
347 self.output
348 .write_all(&compressed[offset..offset + block_size])?;
349 offset += block_size;
350 }
351
352 self.output.write_all(&[0])?;
354
355 Ok(())
356 }
357
358 fn frame_to_rgba(&self, frame: &VideoFrame) -> CodecResult<Vec<u8>> {
360 if frame.width != self.width || frame.height != self.height {
361 return Err(CodecError::InvalidParameter(format!(
362 "Frame size {}x{} doesn't match canvas {}x{}",
363 frame.width, frame.height, self.width, self.height
364 )));
365 }
366
367 match frame.format {
368 PixelFormat::Rgba32 => {
369 if frame.planes.is_empty() {
370 return Err(CodecError::InvalidData("Frame has no planes".into()));
371 }
372 Ok(frame.planes[0].data.to_vec())
373 }
374 PixelFormat::Rgb24 => {
375 if frame.planes.is_empty() {
376 return Err(CodecError::InvalidData("Frame has no planes".into()));
377 }
378 let rgb = &frame.planes[0].data;
379 let mut rgba = Vec::with_capacity((self.width * self.height * 4) as usize);
380 for chunk in rgb.chunks_exact(3) {
381 rgba.extend_from_slice(chunk);
382 rgba.push(255);
383 }
384 Ok(rgba)
385 }
386 _ => Err(CodecError::UnsupportedFeature(format!(
387 "Pixel format {} not supported for GIF encoding",
388 frame.format
389 ))),
390 }
391 }
392
393 fn generate_global_palette(&self, frames: &[VideoFrame]) -> CodecResult<Vec<u8>> {
395 let mut all_colors = Vec::new();
397
398 for frame in frames {
399 let rgba = self.frame_to_rgba(frame)?;
400 for chunk in rgba.chunks_exact(4) {
401 let color = [chunk[0], chunk[1], chunk[2]];
402 all_colors.push(color);
403 }
404 }
405
406 let palette = match self.config.quantization {
408 QuantizationMethod::MedianCut => self.median_cut_quantize(&all_colors)?,
409 QuantizationMethod::Octree => self.octree_quantize(&all_colors)?,
410 };
411
412 Ok(palette)
413 }
414
415 fn median_cut_quantize(&self, colors: &[[u8; 3]]) -> CodecResult<Vec<u8>> {
417 let target_colors = self.config.colors.min(MAX_COLORS);
418
419 let mut buckets = vec![colors.to_vec()];
421
422 while buckets.len() < target_colors {
424 let largest_idx = buckets
426 .iter()
427 .enumerate()
428 .max_by_key(|(_, b)| b.len())
429 .map(|(i, _)| i)
430 .expect("buckets is non-empty inside the while loop");
431
432 let bucket = buckets.remove(largest_idx);
433 if bucket.is_empty() {
434 break;
435 }
436
437 let (mut min_r, mut max_r) = (255, 0);
439 let (mut min_g, mut max_g) = (255, 0);
440 let (mut min_b, mut max_b) = (255, 0);
441
442 for color in &bucket {
443 min_r = min_r.min(color[0]);
444 max_r = max_r.max(color[0]);
445 min_g = min_g.min(color[1]);
446 max_g = max_g.max(color[1]);
447 min_b = min_b.min(color[2]);
448 max_b = max_b.max(color[2]);
449 }
450
451 let range_r = max_r - min_r;
452 let range_g = max_g - min_g;
453 let range_b = max_b - min_b;
454
455 let mut bucket = bucket;
457 if range_r >= range_g && range_r >= range_b {
458 bucket.sort_by_key(|c| c[0]);
459 } else if range_g >= range_r && range_g >= range_b {
460 bucket.sort_by_key(|c| c[1]);
461 } else {
462 bucket.sort_by_key(|c| c[2]);
463 }
464
465 let mid = bucket.len() / 2;
467 let (left, right) = bucket.split_at(mid);
468 buckets.push(left.to_vec());
469 buckets.push(right.to_vec());
470 }
471
472 let mut palette = Vec::with_capacity(target_colors * 3);
474 for bucket in buckets {
475 if bucket.is_empty() {
476 continue;
477 }
478
479 let mut sum_r = 0u32;
480 let mut sum_g = 0u32;
481 let mut sum_b = 0u32;
482
483 for color in &bucket {
484 sum_r += u32::from(color[0]);
485 sum_g += u32::from(color[1]);
486 sum_b += u32::from(color[2]);
487 }
488
489 let count = bucket.len() as u32;
490 palette.push((sum_r / count) as u8);
491 palette.push((sum_g / count) as u8);
492 palette.push((sum_b / count) as u8);
493 }
494
495 Ok(palette)
496 }
497
498 fn octree_quantize(&self, colors: &[[u8; 3]]) -> CodecResult<Vec<u8>> {
500 let mut tree = OctreeQuantizer::new(self.config.colors);
501
502 for &color in colors {
503 tree.add_color(color);
504 }
505
506 let palette = tree.get_palette();
507 Ok(palette)
508 }
509
510 fn quantize_frame(&self, rgba: &[u8]) -> CodecResult<Vec<u8>> {
512 let mut indices = Vec::with_capacity((self.width * self.height) as usize);
513
514 match self.config.dithering {
515 DitheringMethod::None => {
516 for chunk in rgba.chunks_exact(4) {
517 let color = [chunk[0], chunk[1], chunk[2]];
518 let index = self.find_closest_color(color);
519 indices.push(index);
520 }
521 }
522 DitheringMethod::FloydSteinberg => {
523 indices = self.floyd_steinberg_dither(rgba)?;
524 }
525 DitheringMethod::Ordered => {
526 indices = self.ordered_dither(rgba)?;
527 }
528 }
529
530 Ok(indices)
531 }
532
533 fn find_closest_color(&self, color: [u8; 3]) -> u8 {
535 let mut best_index = 0;
536 let mut best_distance = u32::MAX;
537
538 for i in 0..(self.palette.len() / 3) {
539 let pal_r = self.palette[i * 3];
540 let pal_g = self.palette[i * 3 + 1];
541 let pal_b = self.palette[i * 3 + 2];
542
543 let dr = i32::from(color[0]) - i32::from(pal_r);
544 let dg = i32::from(color[1]) - i32::from(pal_g);
545 let db = i32::from(color[2]) - i32::from(pal_b);
546
547 let distance = (dr * dr + dg * dg + db * db) as u32;
548
549 if distance < best_distance {
550 best_distance = distance;
551 best_index = i;
552 }
553 }
554
555 best_index as u8
556 }
557
558 #[allow(clippy::cast_possible_wrap)]
560 fn floyd_steinberg_dither(&self, rgba: &[u8]) -> CodecResult<Vec<u8>> {
561 let width = self.width as usize;
562 let height = self.height as usize;
563
564 let mut errors = vec![[0i16; 3]; width * height];
566 let mut indices = Vec::with_capacity(width * height);
567
568 for y in 0..height {
569 for x in 0..width {
570 let idx = (y * width + x) * 4;
571 let pixel_idx = y * width + x;
572
573 let r = (i16::from(rgba[idx]) + errors[pixel_idx][0]).clamp(0, 255) as u8;
575 let g = (i16::from(rgba[idx + 1]) + errors[pixel_idx][1]).clamp(0, 255) as u8;
576 let b = (i16::from(rgba[idx + 2]) + errors[pixel_idx][2]).clamp(0, 255) as u8;
577
578 let color = [r, g, b];
580 let index = self.find_closest_color(color);
581 indices.push(index);
582
583 let pal_r = self.palette[index as usize * 3];
585 let pal_g = self.palette[index as usize * 3 + 1];
586 let pal_b = self.palette[index as usize * 3 + 2];
587
588 let err_r = i16::from(r) - i16::from(pal_r);
589 let err_g = i16::from(g) - i16::from(pal_g);
590 let err_b = i16::from(b) - i16::from(pal_b);
591
592 if x + 1 < width {
594 let next_idx = pixel_idx + 1;
595 errors[next_idx][0] += err_r * 7 / 16;
596 errors[next_idx][1] += err_g * 7 / 16;
597 errors[next_idx][2] += err_b * 7 / 16;
598 }
599
600 if y + 1 < height {
601 if x > 0 {
602 let next_idx = pixel_idx + width - 1;
603 errors[next_idx][0] += err_r * 3 / 16;
604 errors[next_idx][1] += err_g * 3 / 16;
605 errors[next_idx][2] += err_b * 3 / 16;
606 }
607
608 let next_idx = pixel_idx + width;
609 errors[next_idx][0] += err_r * 5 / 16;
610 errors[next_idx][1] += err_g * 5 / 16;
611 errors[next_idx][2] += err_b * 5 / 16;
612
613 if x + 1 < width {
614 let next_idx = pixel_idx + width + 1;
615 errors[next_idx][0] += err_r / 16;
616 errors[next_idx][1] += err_g / 16;
617 errors[next_idx][2] += err_b / 16;
618 }
619 }
620 }
621 }
622
623 Ok(indices)
624 }
625
626 fn ordered_dither(&self, rgba: &[u8]) -> CodecResult<Vec<u8>> {
628 #[rustfmt::skip]
630 const BAYER_MATRIX: [[i16; 4]; 4] = [
631 [ 0, 8, 2, 10 ],
632 [ 12, 4, 14, 6 ],
633 [ 3, 11, 1, 9 ],
634 [ 15, 7, 13, 5 ],
635 ];
636
637 let width = self.width as usize;
638 let height = self.height as usize;
639 let mut indices = Vec::with_capacity(width * height);
640
641 for y in 0..height {
642 for x in 0..width {
643 let idx = (y * width + x) * 4;
644
645 let threshold = BAYER_MATRIX[y % 4][x % 4] * 16 - 128;
647
648 let r = (i16::from(rgba[idx]) + threshold).clamp(0, 255) as u8;
649 let g = (i16::from(rgba[idx + 1]) + threshold).clamp(0, 255) as u8;
650 let b = (i16::from(rgba[idx + 2]) + threshold).clamp(0, 255) as u8;
651
652 let color = [r, g, b];
653 let index = self.find_closest_color(color);
654 indices.push(index);
655 }
656 }
657
658 Ok(indices)
659 }
660
661 fn color_table_size_field(colors: usize) -> u8 {
663 let size = Self::next_power_of_two(colors);
664 let bits = Self::bits_needed(size);
665 bits.saturating_sub(1)
666 }
667
668 fn next_power_of_two(n: usize) -> usize {
670 let mut power = 2;
671 while power < n {
672 power *= 2;
673 }
674 power
675 }
676
677 fn bits_needed(n: usize) -> u8 {
679 if n <= 2 {
680 1
681 } else {
682 (n as f64).log2().ceil() as u8
683 }
684 }
685}
686
687struct OctreeNode {
689 children: [Option<Box<OctreeNode>>; 8],
690 color_sum: [u32; 3],
691 pixel_count: u32,
692 is_leaf: bool,
693}
694
695impl OctreeNode {
696 fn new() -> Self {
697 Self {
698 children: Default::default(),
699 color_sum: [0, 0, 0],
700 pixel_count: 0,
701 is_leaf: false,
702 }
703 }
704}
705
706struct OctreeQuantizer {
708 root: OctreeNode,
709 #[allow(dead_code)]
710 max_colors: usize,
711 leaf_count: usize,
712}
713
714impl OctreeQuantizer {
715 fn new(max_colors: usize) -> Self {
716 Self {
717 root: OctreeNode::new(),
718 max_colors,
719 leaf_count: 0,
720 }
721 }
722
723 fn add_color(&mut self, color: [u8; 3]) {
724 Self::add_color_recursive(&mut self.root, color, 0, &mut self.leaf_count);
725 }
726
727 fn add_color_recursive(
728 node: &mut OctreeNode,
729 color: [u8; 3],
730 depth: u8,
731 leaf_count: &mut usize,
732 ) {
733 if depth >= 8 || node.is_leaf {
734 node.color_sum[0] += u32::from(color[0]);
735 node.color_sum[1] += u32::from(color[1]);
736 node.color_sum[2] += u32::from(color[2]);
737 node.pixel_count += 1;
738 if !node.is_leaf {
739 node.is_leaf = true;
740 *leaf_count += 1;
741 }
742 return;
743 }
744
745 let index = Self::get_child_index(color, depth);
746
747 if node.children[index].is_none() {
748 node.children[index] = Some(Box::new(OctreeNode::new()));
749 }
750
751 if let Some(child) = &mut node.children[index] {
752 Self::add_color_recursive(child, color, depth + 1, leaf_count);
753 }
754 }
755
756 fn get_child_index(color: [u8; 3], depth: u8) -> usize {
757 let shift = 7 - depth;
758 let r_bit = ((color[0] >> shift) & 1) as usize;
759 let g_bit = ((color[1] >> shift) & 1) as usize;
760 let b_bit = ((color[2] >> shift) & 1) as usize;
761 (r_bit << 2) | (g_bit << 1) | b_bit
762 }
763
764 fn get_palette(&self) -> Vec<u8> {
765 let mut palette = Vec::new();
766 self.collect_colors(&self.root, &mut palette);
767 palette
768 }
769
770 fn collect_colors(&self, node: &OctreeNode, palette: &mut Vec<u8>) {
771 if node.is_leaf && node.pixel_count > 0 {
772 let r = (node.color_sum[0] / node.pixel_count) as u8;
773 let g = (node.color_sum[1] / node.pixel_count) as u8;
774 let b = (node.color_sum[2] / node.pixel_count) as u8;
775 palette.extend_from_slice(&[r, g, b]);
776 return;
777 }
778
779 for child in &node.children {
780 if let Some(child) = child {
781 self.collect_colors(child, palette);
782 }
783 }
784 }
785}
786
787#[cfg(test)]
788mod tests {
789 use super::*;
790
791 #[test]
792 fn test_color_table_size_field() {
793 assert_eq!(GifEncoderState::color_table_size_field(2), 0);
794 assert_eq!(GifEncoderState::color_table_size_field(4), 1);
795 assert_eq!(GifEncoderState::color_table_size_field(256), 7);
796 }
797
798 #[test]
799 fn test_next_power_of_two() {
800 assert_eq!(GifEncoderState::next_power_of_two(1), 2);
801 assert_eq!(GifEncoderState::next_power_of_two(3), 4);
802 assert_eq!(GifEncoderState::next_power_of_two(100), 128);
803 }
804
805 #[test]
806 fn test_bits_needed() {
807 assert_eq!(GifEncoderState::bits_needed(2), 1);
808 assert_eq!(GifEncoderState::bits_needed(4), 2);
809 assert_eq!(GifEncoderState::bits_needed(256), 8);
810 }
811}