1 // Copyright 2024 Google LLC 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 #[allow(unused_imports)] 16 use crate::image::*; 17 use crate::*; 18 19 use std::fs::File; 20 use std::io::prelude::*; 21 22 use super::Writer; 23 24 #[derive(Default)] 25 pub(crate) struct Y4MWriter { 26 header_written: bool, 27 write_alpha: bool, 28 skip_headers: bool, 29 } 30 31 impl Y4MWriter { 32 #[allow(unused)] create(skip_headers: bool) -> Self33 pub(crate) fn create(skip_headers: bool) -> Self { 34 Self { 35 skip_headers, 36 ..Default::default() 37 } 38 } 39 write_header(&mut self, file: &mut File, image: &Image) -> AvifResult<()>40 fn write_header(&mut self, file: &mut File, image: &Image) -> AvifResult<()> { 41 if self.header_written { 42 return Ok(()); 43 } 44 self.write_alpha = false; 45 46 if image.alpha_present && (image.depth != 8 || image.yuv_format != PixelFormat::Yuv444) { 47 println!("WARNING: writing alpha is currently only supported in 8bpc YUV444, ignoring alpha channel"); 48 } 49 50 let y4m_format = match image.depth { 51 8 => match image.yuv_format { 52 PixelFormat::None 53 | PixelFormat::AndroidP010 54 | PixelFormat::AndroidNv12 55 | PixelFormat::AndroidNv21 => "", 56 PixelFormat::Yuv444 => { 57 if image.alpha_present { 58 self.write_alpha = true; 59 "C444alpha XYSCSS=444" 60 } else { 61 "C444 XYSCSS=444" 62 } 63 } 64 PixelFormat::Yuv422 => "C422 XYSCSS=422", 65 PixelFormat::Yuv420 => "C420jpeg XYSCSS=420JPEG", 66 PixelFormat::Yuv400 => "Cmono XYSCSS=400", 67 }, 68 10 => match image.yuv_format { 69 PixelFormat::None 70 | PixelFormat::AndroidP010 71 | PixelFormat::AndroidNv12 72 | PixelFormat::AndroidNv21 => "", 73 PixelFormat::Yuv444 => "C444p10 XYSCSS=444P10", 74 PixelFormat::Yuv422 => "C422p10 XYSCSS=422P10", 75 PixelFormat::Yuv420 => "C420p10 XYSCSS=420P10", 76 PixelFormat::Yuv400 => "Cmono10 XYSCSS=400", 77 }, 78 12 => match image.yuv_format { 79 PixelFormat::None 80 | PixelFormat::AndroidP010 81 | PixelFormat::AndroidNv12 82 | PixelFormat::AndroidNv21 => "", 83 PixelFormat::Yuv444 => "C444p12 XYSCSS=444P12", 84 PixelFormat::Yuv422 => "C422p12 XYSCSS=422P12", 85 PixelFormat::Yuv420 => "C420p12 XYSCSS=420P12", 86 PixelFormat::Yuv400 => "Cmono12 XYSCSS=400", 87 }, 88 _ => { 89 return Err(AvifError::NotImplemented); 90 } 91 }; 92 let y4m_color_range = if image.yuv_range == YuvRange::Limited { 93 "XCOLORRANGE=LIMITED" 94 } else { 95 "XCOLORRANGE=FULL" 96 }; 97 let header = format!( 98 "YUV4MPEG2 W{} H{} F25:1 Ip A0:0 {y4m_format} {y4m_color_range}\n", 99 image.width, image.height 100 ); 101 file.write_all(header.as_bytes()) 102 .or(Err(AvifError::IoError))?; 103 self.header_written = true; 104 Ok(()) 105 } 106 } 107 108 impl Writer for Y4MWriter { write_frame(&mut self, file: &mut File, image: &Image) -> AvifResult<()>109 fn write_frame(&mut self, file: &mut File, image: &Image) -> AvifResult<()> { 110 if !self.skip_headers { 111 self.write_header(file, image)?; 112 let frame_marker = "FRAME\n"; 113 file.write_all(frame_marker.as_bytes()) 114 .or(Err(AvifError::IoError))?; 115 } 116 let planes: &[Plane] = if self.write_alpha { &ALL_PLANES } else { &YUV_PLANES }; 117 for plane in planes { 118 let plane = *plane; 119 if !image.has_plane(plane) { 120 continue; 121 } 122 if image.depth == 8 { 123 for y in 0..image.height(plane) { 124 let row = image.row(plane, y as u32)?; 125 let pixels = &row[..image.width(plane)]; 126 file.write_all(pixels).or(Err(AvifError::IoError))?; 127 } 128 } else { 129 for y in 0..image.height(plane) { 130 let row16 = image.row16(plane, y as u32)?; 131 let pixels16 = &row16[..image.width(plane)]; 132 let mut pixels: Vec<u8> = Vec::new(); 133 // y4m is always little endian. 134 for &pixel16 in pixels16 { 135 pixels.extend_from_slice(&pixel16.to_le_bytes()); 136 } 137 file.write_all(&pixels[..]).or(Err(AvifError::IoError))?; 138 } 139 } 140 } 141 Ok(()) 142 } 143 } 144