Skip to main content

comiconv/
lib.rs

1//! Library for comic book conversion.
2//!
3//! You can convert locally or on a server running comiconv-server.
4
5use cra::{ArcEntry, ArcError, ArcReader, ArcWriter};
6use image::{
7    codecs::{
8        jpeg::JpegEncoder,
9        png::{CompressionType, FilterType, PngEncoder},
10        webp::WebPEncoder,
11    },
12    ColorType, DynamicImage, ImageError, ImageReader,
13};
14use indicatif::{style::TemplateError, ProgressBar, ProgressStyle};
15use infer::image::is_jxl;
16use jxl_oxide::integration::JxlDecoder;
17use libavif_image::{is_avif, read as read_avif, save as save_avif, Error as AvifError};
18use rayon::prelude::{IntoParallelIterator, ParallelIterator};
19use sha2::{Digest, Sha256};
20use std::{
21    fmt::Display,
22    fs::{rename, File},
23    io::{self, Cursor, Read, Write},
24    net::TcpStream,
25    str::FromStr,
26    sync::{Arc, Mutex},
27    time::Duration,
28};
29use thiserror::Error;
30use zune_core::{bit_depth::BitDepth, colorspace::ColorSpace, options::EncoderOptions};
31use zune_jpegxl::{JxlEncodeErrors, JxlSimpleEncoder};
32
33/// This is the main error type for the library
34#[derive(Error, Debug)]
35#[error(transparent)]
36pub enum ConvError {
37    ArcError(#[from] ArcError),
38    IoError(#[from] io::Error),
39    TemplateError(#[from] TemplateError),
40    AvifError(#[from] AvifError),
41    ImageError(#[from] ImageError),
42    #[error("{0:?}")]
43    JxlEncodeError(JxlEncodeErrors),
44    #[error("Invalid server response")]
45    InvalidResponse,
46    #[error("Hash mismatch")]
47    HashMismatch,
48}
49
50pub type ConvResult<T> = Result<T, ConvError>;
51
52/// Enum representing all supported target image formats
53#[derive(Clone, Copy, Debug)]
54pub enum Format {
55    Jpeg,
56    JpegXL,
57    Png,
58    Webp,
59    Avif,
60}
61
62impl Display for Format {
63    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64        f.write_str(match self {
65            Format::Jpeg => "jpg",
66            Format::JpegXL => "jxl",
67            Format::Png => "png",
68            Format::Webp => "webp",
69            Format::Avif => "avif",
70        })
71    }
72}
73
74impl FromStr for Format {
75    type Err = String;
76
77    fn from_str(s: &str) -> Result<Self, Self::Err> {
78        match s.to_ascii_lowercase().as_str() {
79            "avif" => Ok(Format::Avif),
80            "jpeg" | "jpg" => Ok(Format::Jpeg),
81            "jxl" => Ok(Format::JpegXL),
82            "webp" => Ok(Format::Webp),
83            "png" => Ok(Format::Png),
84            _ => Err(format!("Invalid format: {s}")),
85        }
86    }
87}
88
89/// This is the main struct for converting
90/// `quality` is ignored for webp
91#[derive(Clone, Copy, Debug)]
92pub struct Converter {
93    pub quality: u8,
94    pub speed: u8,
95    pub format: Format,
96    pub backup: bool,
97    pub quiet: bool,
98}
99
100impl Default for Converter {
101    fn default() -> Self {
102        Self {
103            quality: 30,
104            speed: 3,
105            format: Format::Avif,
106            backup: false,
107            quiet: false,
108        }
109    }
110}
111
112impl Converter {
113    /// Takes a path to a file and converts it
114    pub fn convert_file(self, file: &str) -> ConvResult<()> {
115        let buf = {
116            let mut buf = Vec::new();
117            File::open(file)?.read_to_end(&mut buf)?;
118            buf
119        };
120        if !self.quiet {
121            println!("Converting {}...", file);
122        }
123        let data = self.convert(&buf, None)?;
124        if self.backup {
125            rename(file, format!("{}.bak", file))?;
126        }
127        File::create(file)?.write_all(&data)?;
128        Ok(())
129    }
130
131    /// Takes a path to a file, a tcp connection and converts the file using a server
132    pub fn convert_file_online(self, file: &str, stream: &mut TcpStream) -> ConvResult<()> {
133        let buf = {
134            let mut buf = Vec::new();
135            File::open(file)?.read_to_end(&mut buf)?;
136            buf
137        };
138        if !self.quiet {
139            println!("Converting {}...", file);
140        }
141        let data = self.convert_online(&buf, stream)?;
142        if self.backup {
143            rename(file, format!("{}.bak", file))?;
144        }
145        File::create(file)?.write_all(&data)?;
146        Ok(())
147    }
148
149    /// Takes contents of a file as a slice of bytes and return the new converted file as bytes.
150    /// Optionally takes a stream to write progress information for the client (only really used by the server).
151    pub fn convert(
152        mut self,
153        buf: &[u8],
154        mut status_stream: Option<&mut TcpStream>,
155    ) -> ConvResult<Vec<u8>> {
156        self.speed = self.speed.clamp(0, 10);
157        self.quality = self.quality.clamp(0, 100);
158        let mut archive = ArcReader::new(buf)?;
159        let mut writer = ArcWriter::new(archive.format());
160        let file_count = archive
161            .by_ref()
162            .filter(|entry| !matches!(entry, ArcEntry::Directory(_)))
163            .count();
164        if let Some(ref mut stream) = status_stream {
165            stream.write_all(&(file_count as u32).to_be_bytes())?;
166        }
167        let mut bar = if self.quiet {
168            ProgressBar::hidden()
169        } else {
170            ProgressBar::new(file_count as u64)
171        };
172        bar.set_style(
173            ProgressStyle::default_bar()
174                .template("Convert  [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})")?
175                .progress_chars("=>-"),
176        );
177        let status_stream = status_stream.map(|stream| Arc::new(Mutex::new(stream)));
178        let pb = Arc::new(Mutex::new(&mut bar));
179        writer.extend(
180            &archive
181                .entries()
182                .into_par_iter()
183                .map(|entry| {
184                    Ok(match entry {
185                        ArcEntry::File(name, data) => {
186                            let data = self.convert_image(data)?;
187                            if let Some(stream) = status_stream.clone() {
188                                stream.lock().unwrap().write_all(b"plus")?
189                            }
190                            pb.clone().lock().unwrap().inc(1);
191                            ArcEntry::File(
192                                format!(
193                                    "{}.{}",
194                                    name.rsplit_once('.').unwrap_or((name, "")).0,
195                                    self.format
196                                ),
197                                data,
198                            )
199                        }
200                        other => other.clone(),
201                    })
202                })
203                .collect::<ConvResult<Vec<ArcEntry>>>()?,
204        );
205        bar.finish();
206        Ok(writer.archive()?)
207    }
208
209    /// Takes file contents as a slice and a tcp connection to the server
210    pub fn convert_online(mut self, buf: &[u8], stream: &mut TcpStream) -> ConvResult<Vec<u8>> {
211        self.speed = self.speed.clamp(0, 10);
212        self.quality = self.quality.clamp(0, 100);
213        stream.set_nodelay(true)?;
214        stream.set_read_timeout(Some(Duration::from_secs(10)))?;
215        stream.write_all(b"comi")?;
216        {
217            let mut buf = [0; 4];
218            stream.read_exact(&mut buf)?;
219            if &buf != b"conv" {
220                return Err(ConvError::InvalidResponse);
221            }
222        }
223        let format = match self.format {
224            Format::Avif => b'A',
225            Format::Webp => b'W',
226            Format::Png => b'P',
227            Format::Jpeg => b'J',
228            Format::JpegXL => todo!(),
229        };
230        let mut left = buf.len();
231        {
232            let mut buf = [0; 8];
233            buf[0] = format;
234            buf[1] = self.speed;
235            buf[2] = self.quality;
236            buf[4..].copy_from_slice(&(left as u32).to_be_bytes());
237            stream.write_all(&buf)?;
238        }
239        let hash = {
240            let mut hasher = Sha256::new();
241            hasher.update(buf);
242            hasher.finalize()
243        };
244        stream.write_all(&hash)?;
245        let mut sent = 0;
246        let pb = if self.quiet {
247            ProgressBar::hidden()
248        } else {
249            ProgressBar::new(left as u64)
250        };
251        pb.set_style(
252            ProgressStyle::default_bar()
253                .template("Upload   [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")?
254                .progress_chars("=>-"),
255        );
256        while left > 0 {
257            let size = left.min(1024 * 1024);
258            stream.write_all(&buf[sent..sent + size])?;
259            let mut buf = [0; 2];
260            stream.read_exact(&mut buf)?;
261            if &buf != b"ok" {
262                return Err(ConvError::InvalidResponse);
263            }
264            sent += size;
265            left = left.saturating_sub(size);
266            pb.inc(size as u64);
267        }
268        pb.finish();
269        let mut left = {
270            let mut buf = [0; 4];
271            stream.read_exact(&mut buf)?;
272            u32::from_be_bytes(buf)
273        };
274        let pb = if self.quiet {
275            ProgressBar::hidden()
276        } else {
277            ProgressBar::new(left as u64)
278        };
279        pb.set_style(
280            ProgressStyle::default_bar()
281                .template("Convert  [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})")?
282                .progress_chars("=>-"),
283        );
284        while left > 0 {
285            let response = {
286                let mut buf = [0; 4];
287                stream.read_exact(&mut buf)?;
288                buf
289            };
290            if &response != b"plus" {
291                return Err(ConvError::InvalidResponse);
292            }
293            pb.inc(1);
294            left -= 1;
295        }
296        pb.finish();
297        let mut left = {
298            let mut buf = [0; 4];
299            stream.read_exact(&mut buf)?;
300            u32::from_be_bytes(buf)
301        };
302        let hash = {
303            let mut buf = [0; 32];
304            stream.read_exact(&mut buf)?;
305            buf
306        };
307        let pb = if self.quiet {
308            ProgressBar::hidden()
309        } else {
310            ProgressBar::new(left as u64)
311        };
312        pb.set_style(
313            ProgressStyle::default_bar()
314                .template("Download [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")?
315                .progress_chars("=>-"),
316        );
317        let mut data = Vec::new();
318        while left > 0 {
319            let mut buf = [0; 1024 * 1024];
320            let read = stream.read(&mut buf)?;
321            pb.inc(read as u64);
322            data.extend_from_slice(&buf[..read]);
323            left = left.saturating_sub(read as u32);
324        }
325        pb.finish();
326        let mut hasher = Sha256::new();
327        hasher.update(&data);
328        if hasher.finalize() != hash.into() {
329            return Err(ConvError::HashMismatch);
330        }
331        Ok(data)
332    }
333
334    fn convert_image(self, buf: &[u8]) -> ConvResult<Vec<u8>> {
335        let image = if is_avif(buf) {
336            read_avif(buf)?
337        } else if is_jxl(buf) {
338            DynamicImage::from_decoder(JxlDecoder::new(buf)?)?
339        } else {
340            ImageReader::new(Cursor::new(buf))
341                .with_guessed_format()?
342                .decode()?
343        };
344        let mut data = Vec::new();
345        match self.format {
346            Format::Avif => {
347                data = save_avif(&image)?.to_vec();
348            }
349            Format::Webp => image.write_with_encoder(WebPEncoder::new_lossless(&mut data))?,
350            Format::Png => image.write_with_encoder(PngEncoder::new_with_quality(
351                &mut data,
352                match self.speed.clamp(0, 2) {
353                    0 => CompressionType::Fast,
354                    1 => CompressionType::Default,
355                    2 => CompressionType::Best,
356                    _ => unreachable!(),
357                },
358                FilterType::Adaptive,
359            ))?,
360            Format::Jpeg => image
361                .into_rgb8()
362                .write_with_encoder(JpegEncoder::new_with_quality(&mut data, self.quality))?,
363            Format::JpegXL => {
364                let (color, depth) = image_to_zune_colot_type(&image);
365                data = JxlSimpleEncoder::new(
366                    image.as_bytes(),
367                    EncoderOptions::new(image.width() as _, image.height() as _, color, depth),
368                )
369                .encode()
370                .map_err(ConvError::JxlEncodeError)?;
371            }
372        }
373        Ok(data)
374    }
375}
376
377fn image_to_zune_colot_type(image: &DynamicImage) -> (ColorSpace, BitDepth) {
378    match image.color() {
379        ColorType::L8 => (ColorSpace::Luma, BitDepth::Eight),
380        ColorType::La16 => (ColorSpace::LumaA, BitDepth::Sixteen),
381        ColorType::Rgb16 => (ColorSpace::RGB, BitDepth::Sixteen),
382        ColorType::Rgba16 => (ColorSpace::RGBA, BitDepth::Sixteen),
383        ColorType::Rgb32F => (ColorSpace::RGB, BitDepth::Float32),
384        ColorType::Rgba32F => (ColorSpace::RGBA, BitDepth::Float32),
385        ColorType::Rgba8 => (ColorSpace::RGBA, BitDepth::Eight),
386        ColorType::L16 => (ColorSpace::Luma, BitDepth::Sixteen),
387        ColorType::La8 => (ColorSpace::LumaA, BitDepth::Eight),
388        ColorType::Rgb8 => (ColorSpace::RGB, BitDepth::Eight),
389        _ => unimplemented!(),
390    }
391}