main.rs

#![allow(non_snake_case)]

#[macro_use]
extern crate ndarray;

use std::{
    env,
    fs::File,
    io::{BufRead, BufReader, BufWriter, Read, Write},
};

use ndarray::prelude::*;

/// Ivan Borgia Dardi - 101987
///
/// Uso: ./dct arquivo_entrada.pgm arquivo_saida.pgm
///
fn main() -> Result<(), Box<std::error::Error>> {
    // Leitura de argumentos da linha de comando

    let args: Vec<String> = env::args().collect();

    // Inicializacao de matrizes

    let DCT: Array2<f64> = Array::from_shape_fn((8, 8), |(r, c)| {
        if r == 0 {
            1.0 / 8f64.sqrt()
        } else {
            0.25f64.sqrt() * ((2.0 * c as f64 + 1.0) * r as f64 * std::f64::consts::PI / 16.0).cos()
        }
    });

    let Q50: Array2<f64> = array![
        [16, 11, 10, 16, 24, 40, 51, 61],
        [12, 12, 14, 19, 26, 58, 60, 55],
        [14, 13, 16, 24, 40, 57, 69, 56],
        [14, 17, 22, 29, 51, 87, 80, 62],
        [18, 22, 37, 56, 68, 109, 103, 77],
        [24, 35, 55, 64, 81, 104, 113, 92],
        [49, 64, 78, 87, 103, 121, 120, 101],
        [72, 92, 95, 98, 112, 100, 103, 99],
    ].mapv(|x| x as f64);

    // Leitura da imagem

    let mut buffered = BufReader::new(File::open(&args[1])?);

    let mut magic = String::new();
    buffered.read_line(&mut magic)?;
    assert_eq!("P2\n", magic);

    let mut dims = String::new();
    buffered.read_line(&mut dims)?;
    let (width, height) = {
        let dims: Vec<usize> = dims
            .split_whitespace()
            .map(|x| x.parse().unwrap())
            .collect();
        (dims[0], dims[1])
    };

    let mut depth = String::new();
    buffered.read_line(&mut depth)?;
    let depth = depth.trim().parse::<usize>().unwrap();

    let mut data = String::new();
    buffered.read_to_string(&mut data)?;
    let mut M: Array2<f64> = Array::from_iter(
        data.split_whitespace()
            .map(|x| x.parse::<u8>().unwrap() as f64),
    ).into_shape((width, height))?;

    // Aplicacao da DCT e da quantificacao

    for mut chunk in M.exact_chunks_mut((8, 8)) {
        chunk -= 128.0;
        let D = &DCT.dot(&chunk.view()).dot(&DCT.t());
        chunk.assign(D);
        chunk /= &Q50;
        chunk.mapv_inplace(|x| x.round());
    }

    // Processo reverso

    for mut chunk in M.exact_chunks_mut((8, 8)) {
        chunk *= &Q50;
        let D = &DCT.t().dot(&chunk.view()).dot(&DCT);
        chunk.assign(D);
        chunk.mapv_inplace(|x| x.round());
        chunk += 128.0;
    }

    let N = M.mapv(|x| x as u8);

    // Escrita da imagem

    let mut output = BufWriter::new(File::create(&args[2])?);
    writeln!(&mut output, "P2")?;
    writeln!(&mut output, "{} {}", width, height)?;
    writeln!(&mut output, "{}", depth)?;
    let mut line = String::new();
    let mut len = 0;
    for x in N.iter() {
        let x = {
            let mut x = x.to_string();
            x.push_str(" ");
            x
        };
        let x_len = x.len();
        len += x_len;
        if len > 70 {
            writeln!(&mut output, "{}", line)?;
            line = x;
            len = x_len;
        } else {
            line.push_str(&x);
        }
    }
    writeln!(&mut output, "{}", line)?;

    Ok(())
}