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 use crate::image::*;
16 use crate::*;
17 use std::fs::File;
18 use std::io::prelude::*;
19 
20 #[derive(Default)]
21 pub struct Y4MWriter {
22     pub filename: Option<String>,
23     header_written: bool,
24     file: Option<File>,
25     write_alpha: bool,
26 }
27 
28 impl Y4MWriter {
create(filename: &str) -> Self29     pub fn create(filename: &str) -> Self {
30         Self {
31             filename: Some(filename.to_owned()),
32             ..Self::default()
33         }
34     }
35 
create_from_file(file: File) -> Self36     pub fn create_from_file(file: File) -> Self {
37         Self {
38             file: Some(file),
39             ..Self::default()
40         }
41     }
42 
write_header(&mut self, image: &Image) -> bool43     fn write_header(&mut self, image: &Image) -> bool {
44         if self.header_written {
45             return true;
46         }
47         self.write_alpha = false;
48 
49         if image.alpha_present && (image.depth != 8 || image.yuv_format != PixelFormat::Yuv444) {
50             println!("WARNING: writing alpha is currently only supported in 8bpc YUV444, ignoring alpha channel");
51         }
52 
53         let y4m_format = match image.depth {
54             8 => match image.yuv_format {
55                 PixelFormat::None
56                 | PixelFormat::AndroidP010
57                 | PixelFormat::AndroidNv12
58                 | PixelFormat::AndroidNv21 => "",
59                 PixelFormat::Yuv444 => {
60                     if image.alpha_present {
61                         self.write_alpha = true;
62                         "C444alpha XYSCSS=444"
63                     } else {
64                         "C444 XYSCSS=444"
65                     }
66                 }
67                 PixelFormat::Yuv422 => "C422 XYSCSS=422",
68                 PixelFormat::Yuv420 => "C420jpeg XYSCSS=420JPEG",
69                 PixelFormat::Yuv400 => "Cmono XYSCSS=400",
70             },
71             10 => match image.yuv_format {
72                 PixelFormat::None
73                 | PixelFormat::AndroidP010
74                 | PixelFormat::AndroidNv12
75                 | PixelFormat::AndroidNv21 => "",
76                 PixelFormat::Yuv444 => "C444p10 XYSCSS=444P10",
77                 PixelFormat::Yuv422 => "C422p10 XYSCSS=422P10",
78                 PixelFormat::Yuv420 => "C420p10 XYSCSS=420P10",
79                 PixelFormat::Yuv400 => "Cmono10 XYSCSS=400",
80             },
81             12 => match image.yuv_format {
82                 PixelFormat::None
83                 | PixelFormat::AndroidP010
84                 | PixelFormat::AndroidNv12
85                 | PixelFormat::AndroidNv21 => "",
86                 PixelFormat::Yuv444 => "C444p12 XYSCSS=444P12",
87                 PixelFormat::Yuv422 => "C422p12 XYSCSS=422P12",
88                 PixelFormat::Yuv420 => "C420p12 XYSCSS=420P12",
89                 PixelFormat::Yuv400 => "Cmono12 XYSCSS=400",
90             },
91             _ => {
92                 return false;
93             }
94         };
95         let y4m_color_range = if image.yuv_range == YuvRange::Limited {
96             "XCOLORRANGE=LIMITED"
97         } else {
98             "XCOLORRANGE=FULL"
99         };
100         let header = format!(
101             "YUV4MPEG2 W{} H{} F25:1 Ip A0:0 {y4m_format} {y4m_color_range}\n",
102             image.width, image.height
103         );
104         if self.file.is_none() {
105             assert!(self.filename.is_some());
106             let file = File::create(self.filename.unwrap_ref());
107             if file.is_err() {
108                 return false;
109             }
110             self.file = Some(file.unwrap());
111         }
112         if self.file.unwrap_ref().write_all(header.as_bytes()).is_err() {
113             return false;
114         }
115         self.header_written = true;
116         true
117     }
118 
write_frame(&mut self, image: &Image) -> bool119     pub fn write_frame(&mut self, image: &Image) -> bool {
120         if !self.write_header(image) {
121             return false;
122         }
123         let frame_marker = "FRAME\n";
124         if self
125             .file
126             .unwrap_ref()
127             .write_all(frame_marker.as_bytes())
128             .is_err()
129         {
130             return false;
131         }
132         let planes: &[Plane] = if self.write_alpha { &ALL_PLANES } else { &YUV_PLANES };
133         for plane in planes {
134             let plane = *plane;
135             if !image.has_plane(plane) {
136                 continue;
137             }
138             if image.depth == 8 {
139                 for y in 0..image.height(plane) {
140                     let row = if let Ok(row) = image.row(plane, y as u32) {
141                         row
142                     } else {
143                         return false;
144                     };
145                     let pixels = &row[..image.width(plane)];
146                     if self.file.unwrap_ref().write_all(pixels).is_err() {
147                         return false;
148                     }
149                 }
150             } else {
151                 for y in 0..image.height(plane) {
152                     let row16 = if let Ok(row16) = image.row16(plane, y as u32) {
153                         row16
154                     } else {
155                         return false;
156                     };
157                     let pixels16 = &row16[..image.width(plane)];
158                     let mut pixels: Vec<u8> = Vec::new();
159                     // y4m is always little endian.
160                     for &pixel16 in pixels16 {
161                         pixels.extend_from_slice(&pixel16.to_le_bytes());
162                     }
163                     if self.file.unwrap_ref().write_all(&pixels[..]).is_err() {
164                         return false;
165                     }
166                 }
167             }
168         }
169         true
170     }
171 }
172