1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101

use super::bbox_modify;
use rand::distributions::{IndependentSample, Range};
use std::fs::create_dir_all;
use std::path::Path;
use std::path::PathBuf;
use std::env::current_dir;
use imageproc::affine::{translate};
use image::ColorType;
use image;
use rand::ThreadRng;
use super::translation_lambda;

#[derive(Debug)]
pub struct FileMapping {
    pub jpg_pathbuf_in: PathBuf,
    pub xml_pathbuf_in: PathBuf,
    pub jpg_pathbuf_out: PathBuf,
    pub xml_pathbuf_out: PathBuf,
}

impl FileMapping {
    pub fn new(jpg_pathbuf_in: PathBuf, output_dirname: &str) -> FileMapping {
        let jpg_filename = jpg_pathbuf_in.file_name()
            .expect("Failed to get filename").to_str()
            .expect("Failed to convert filename to str");
        let jpg_parent_pathbuf = jpg_pathbuf_in
            .parent()
            .expect("Failed to get parent dir");
        let jpg_parent_dirname = jpg_parent_pathbuf
            .file_name()
            .expect("Failed to get dirname of parent")
            .to_str()
            .expect("Failed to convert parent dirname to str");
        assert_eq!(jpg_parent_dirname, "JPEGImages");
        let input_pathbuf = jpg_parent_pathbuf
            .parent()
            .expect("Failed to get data root.");
        let data_parent = match input_pathbuf.parent() {
            None => panic!("Failed to get the parent folder of data root."),
            Some(p) if p == Path::new("") => current_dir().expect("Failed to get cwd"),
            Some(p) => p.to_owned(),
        };
        let xml_pathbuf_in = input_pathbuf
            .join("Annotations")
            .join(jpg_filename)
            .with_extension("xml");
        let output_pathbuf = data_parent
            .join(output_dirname);

        let output_jpg_subpath = output_pathbuf.join("JPEGImages");
        let output_xml_subpath = output_pathbuf.join("Annotations");
        if !output_jpg_subpath.exists() {
            create_dir_all(output_jpg_subpath.as_path())
                .expect("Failed to create jpg subpath.");
        }
        if !output_xml_subpath.exists() {
            create_dir_all(output_xml_subpath.as_path())
                .expect("Failed to create xml subpath.");
        }
        let jpg_pathbuf_out = output_jpg_subpath
            .join(Path::new(jpg_filename));
        let xml_pathbuf_out = output_xml_subpath
            .join(Path::new(jpg_filename))
            .with_extension("xml");
        FileMapping {
            jpg_pathbuf_in: jpg_pathbuf_in.clone(),
            xml_pathbuf_in,
            jpg_pathbuf_out,
            xml_pathbuf_out,
        }
    }

    pub fn process(&self, delta_x_percent: f32, delta_y_percent: f32, rng: &mut ThreadRng) {

        let img = image::open(self.jpg_pathbuf_in.as_path())
            .expect("Failed to read image.")
            .to_rgb();
        let width = img.width();
        let height = img.height();
        let w = width as f32;
        let h = height as f32;
        let delta_x_max = (w * delta_x_percent / 100.0) as i32;
        let delta_y_max = (h * delta_y_percent / 100.0) as i32;
        let x_range = Range::new(-delta_x_max, delta_x_max);
        let y_range = Range::new(-delta_y_max, delta_y_max);
        let delta_x = x_range.ind_sample(rng);
        let delta_y = y_range.ind_sample(rng);
        bbox_modify(self.xml_pathbuf_in.as_path(),
                    self.xml_pathbuf_out.as_path(),
                    &*translation_lambda(delta_x, delta_y))
            .expect("Failed to modify xml");
        let translated_img = translate(&img, (delta_x, delta_y));
        image::save_buffer(self.jpg_pathbuf_out.as_path(),
                           &translated_img,
                           width,
                           height,
                           ColorType::RGB(8))
            .expect("Failed to save modified image.");
    }
}