camera_model.rs

use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;
use std::str::FromStr;

use crate::Error;

pub struct CamModel {
    pub tstw: Vec<f32>,
    pub num_ctrl_pts: usize,
    pub ctrl_pts: Vec<f32>,
    pub weights: Vec<f32>,
    pub coefs: Vec<f32>,
    pub tonemap: Vec<f32>,
}

pub fn load_cam_model<P: AsRef<Path>>(cam_model: P, chans: usize) -> Result<CamModel, Error> {
    let (num_ctrl_pts, ctrl_pts) = load_ctrl_pts(&cam_model, chans)?;
    let (weights, coefs) = load_weights_coefs(&cam_model, chans, num_ctrl_pts)?;

    Ok(CamModel {
        tstw: load_tstw(&cam_model, chans)?,
        num_ctrl_pts,
        ctrl_pts,
        weights,
        coefs,
        tonemap: load_tonemap(&cam_model, chans)?,
    })
}

fn load_tstw<P: AsRef<Path>>(cam_model: &P, chans: usize) -> Result<Vec<f32>, Error> {
    let f = File::open(Path::join(cam_model.as_ref(), "raw2jpg_transform.txt"))?;
    let f = BufReader::new(f);

    // TODO: I don't understand this file and what this computation is doing, but I can't find any
    // source even going back to the original repo
    // NOTE: I've replaced the "wb_index" with the value it's provided in the HPVM repo (6)
    let wb_base = 5 + 5 * (6 - 1);

    let lines = f.lines().skip(wb_base).take(chans);

    let mut tstw = vec![0.0; chans * chans];
    let mut tstw_chunked = tstw.chunks_mut(chans).collect::<Vec<_>>();
    let tstw_ref: &mut [&mut [f32]] = tstw_chunked.as_mut_slice();

    for (i, line) in lines.enumerate() {
        let line = line?;
        let words = line.split_whitespace();
        let mut count = 0;
        for (j, word) in words.enumerate() {
            // The HPVM version computes this data structure and then transposes it separately, but
            // doing the transposition here since it makes more sense
            tstw_ref[j][i] = f32::from_str(word)?;
            count += 1;
        }
        assert_eq!(count, chans, "Missing tstw channels");
    }

    Ok(tstw)
}

fn load_ctrl_pts<P: AsRef<Path>>(cam_model: &P, chans: usize) -> Result<(usize, Vec<f32>), Error> {
    let f = File::open(Path::join(cam_model.as_ref(), "raw2jpg_ctrlPoints.txt"))?;
    let mut f = BufReader::new(f);

    let mut first_line = String::new();
    let _ = f.read_line(&mut first_line)?;
    let num_ctrl_pts = usize::from_str(first_line.trim())?;

    let mut ctrl_pts = vec![0.0; num_ctrl_pts * chans];
    let mut ctrl_pts_chunked = ctrl_pts.chunks_mut(chans).collect::<Vec<_>>();
    let result: &mut [&mut [f32]] = ctrl_pts_chunked.as_mut_slice();

    for (i, line) in f.lines().enumerate() {
        if i >= num_ctrl_pts {
            break;
        }

        let line = line?;
        let words = line.split_whitespace();
        let mut count = 0;
        for (j, word) in words.enumerate() {
            result[i][j] = f32::from_str(word)?;
            count += 1;
        }
        assert_eq!(count, chans, "Missing control point channel");
    }

    Ok((num_ctrl_pts, ctrl_pts))
}

fn load_weights_coefs<P: AsRef<Path>>(
    cam_model: &P,
    chans: usize,
    num_ctrl_pts: usize,
) -> Result<(Vec<f32>, Vec<f32>), Error> {
    let f = File::open(Path::join(cam_model.as_ref(), "raw2jpg_coefs.txt"))?;
    let mut f = BufReader::new(f);

    let mut first_line = String::new();
    let _ = f.read_line(&mut first_line)?;
    let lines = usize::from_str(first_line.trim())?;

    assert_eq!(lines, num_ctrl_pts + 4);

    let mut weights = vec![0.0; num_ctrl_pts * chans];
    let mut weights_chunked = weights.chunks_mut(chans).collect::<Vec<_>>();
    let weights_dst: &mut [&mut [f32]] = weights_chunked.as_mut_slice();

    let mut coefs = vec![0.0; 4 * chans];
    let mut coefs_chunked = coefs.chunks_mut(chans).collect::<Vec<_>>();
    let coefs_dst: &mut [&mut [f32]] = coefs_chunked.as_mut_slice();

    for (i, line) in f.lines().enumerate() {
        let line = line?;
        let words = line.split_whitespace();

        if i < num_ctrl_pts {
            let mut count = 0;
            for (j, word) in words.enumerate() {
                weights_dst[i][j] = f32::from_str(word)?;
                count += 1;
            }
            assert_eq!(count, chans, "Missing weight channel");
        } else if i < num_ctrl_pts + 4 {
            let mut count = 0;
            for (j, word) in words.enumerate() {
                coefs_dst[i - num_ctrl_pts][j] = f32::from_str(word)?;
                count += 1;
            }
            assert_eq!(count, chans, "Missing coefficient channel");
        } else {
            break;
        }
    }

    Ok((weights, coefs))
}

fn load_tonemap<P: AsRef<Path>>(cam_model: &P, chans: usize) -> Result<Vec<f32>, Error> {
    let f = File::open(Path::join(cam_model.as_ref(), "raw2jpg_respFcns.txt"))?;
    let mut f = BufReader::new(f);

    let mut first_line = String::new();
    let _ = f.read_line(&mut first_line)?;
    // TODO: What does the first line mean? Should we use it for something?

    let mut tonemap = vec![0.0; 256 * chans];
    let mut tonemap_chunked = tonemap.chunks_mut(chans).collect::<Vec<_>>();
    let result: &mut [&mut [f32]] = tonemap_chunked.as_mut_slice();

    for (i, line) in f.lines().enumerate() {
        if i >= 256 {
            break;
        }

        let line = line?;
        let words = line.split_whitespace();
        let mut count = 0;
        for (j, word) in words.enumerate() {
            result[i][j] = f32::from_str(word)?;
            count += 1;
        }
        assert_eq!(count, chans, "Missing tonemap channels");
    }

    Ok(tonemap)
}