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::Image;
16 use crate::image::ALL_PLANES;
17 use crate::reformat::rgb;
18 use crate::OptionExtension;
19 
20 use std::fs::File;
21 use std::io::prelude::*;
22 
23 #[derive(Default)]
24 pub struct RawWriter {
25     pub filename: Option<String>,
26     pub rgb: bool,
27     file: Option<File>,
28 }
29 
30 impl RawWriter {
create(filename: &str) -> Self31     pub fn create(filename: &str) -> Self {
32         Self {
33             filename: Some(filename.to_owned()),
34             ..Self::default()
35         }
36     }
37 
write_header(&mut self) -> bool38     fn write_header(&mut self) -> bool {
39         if self.file.is_none() {
40             assert!(self.filename.is_some());
41             let file = File::create(self.filename.unwrap_ref());
42             if file.is_err() {
43                 return false;
44             }
45             self.file = Some(file.unwrap());
46         }
47         true
48     }
49 
write_frame(&mut self, image: &Image) -> bool50     pub fn write_frame(&mut self, image: &Image) -> bool {
51         if !self.write_header() {
52             return false;
53         }
54         if self.rgb {
55             let mut rgb = rgb::Image::create_from_yuv(image);
56             rgb.format = rgb::Format::Rgba;
57             rgb.depth = 16;
58             //rgb.depth = 8;
59             rgb.premultiply_alpha = true;
60             rgb.is_float = true;
61             if rgb.allocate().is_err() || rgb.convert_from_yuv(image).is_err() {
62                 return false;
63             }
64             for y in 0..rgb.height {
65                 if rgb.depth == 8 {
66                     let row = rgb.row(y).unwrap();
67                     if self.file.unwrap_ref().write_all(row).is_err() {
68                         return false;
69                     }
70                 } else {
71                     let row = rgb.row16(y).unwrap();
72                     let mut row16: Vec<u8> = Vec::new();
73                     for &pixel in row {
74                         row16.extend_from_slice(&pixel.to_be_bytes());
75                     }
76                     if self.file.unwrap_ref().write_all(&row16[..]).is_err() {
77                         return false;
78                     }
79                 }
80             }
81             return true;
82         }
83         for plane in ALL_PLANES {
84             let plane_data = image.plane_data(plane);
85             if plane_data.is_none() {
86                 continue;
87             }
88             let plane_data = plane_data.unwrap();
89             for y in 0..plane_data.height {
90                 if image.depth == 8 {
91                     let row = image.row(plane, y);
92                     if row.is_err() {
93                         return false;
94                     }
95                     let row = &row.unwrap()[..plane_data.width as usize];
96                     if self.file.unwrap_ref().write_all(row).is_err() {
97                         return false;
98                     }
99                 } else {
100                     let row = image.row16(plane, y);
101                     if row.is_err() {
102                         return false;
103                     }
104                     let row = &row.unwrap()[..plane_data.width as usize];
105                     let mut row16: Vec<u8> = Vec::new();
106                     for &pixel in row {
107                         row16.extend_from_slice(&pixel.to_le_bytes());
108                     }
109                     if self.file.unwrap_ref().write_all(&row16[..]).is_err() {
110                         return false;
111                     }
112                 }
113             }
114         }
115         true
116     }
117 }
118