Skip to content
Snippets Groups Projects
lib.rs 8.4 KiB
Newer Older
  • Learn to ignore specific revisions
  • rarbore2's avatar
    rarbore2 committed
    #![cfg(feature = "opencv")]
    
    rarbore2's avatar
    rarbore2 committed
    #![feature(concat_idents)]
    
    mod edge_detection_rust;
    
    use hercules_rt::{runner, HerculesImmBox, HerculesImmBoxTo, HerculesMutBox};
    
    use std::slice::from_raw_parts;
    
    use clap::Parser;
    
    use opencv::core::{Mat, Size, CV_32F, CV_8U};
    use opencv::highgui::{imshow, wait_key};
    use opencv::imgproc::{cvt_color_def, ColorConversionCodes};
    use opencv::prelude::{MatTraitConst, VideoCaptureTrait, VideoCaptureTraitConst};
    use opencv::videoio::{VideoCapture, VideoCaptureProperties, VideoWriter, VideoWriterTrait};
    
    juno_build::juno!("edge_detection");
    
    #[derive(Parser)]
    #[clap(author, version, about, long_about = None)]
    pub struct EdgeDetectionInputs {
        pub input: String,
        #[clap(short, long)]
        pub display: bool,
        #[clap(short, long, value_name = "PATH")]
        pub output: Option<String>,
        #[clap(short, long)]
        pub verify: bool,
        #[clap(long = "display-verify")]
        pub display_verify: bool,
        #[clap(long = "output-verify", value_name = "PATH")]
        pub output_verify: Option<String>,
        #[clap(short, long, value_name = "COUNT")]
        pub frames: Option<usize>,
    }
    
    pub fn load_frame(video: &mut VideoCapture) -> Mat {
        let mut frame = Mat::default();
    
        let Ok(true) = video.read(&mut frame) else {
            panic!("Failed to load frame");
        };
        let result = if frame.channels() == 3 {
            let mut converted = Mat::default();
            let () = cvt_color_def(
                &frame,
                &mut converted,
                ColorConversionCodes::COLOR_BGR2GRAY.into(),
            )
            .expect("Failure in conversion to grayscale");
            let mut result = Mat::default();
            let () = converted
                .convert_to(&mut result, CV_32F, 1.0 / 255.0, 0.0)
                .expect("Failure in conversion to f32");
            result
        } else if frame.channels() == 1 {
            let mut result = Mat::default();
            let () = frame
                .convert_to(&mut result, CV_32F, 1.0 / 255.0, 0.0)
                .expect("Failure in conversion to f32");
            result
        } else {
            panic!("Expected either RGB or grayscale image");
        };
    
        assert!(result.is_continuous());
        result
    }
    
    pub fn frame_from_slice(frame: &[f32], height: usize, width: usize) -> Mat {
        let result = Mat::from_slice(frame)
            .expect("Failed to create matrix from result")
            .reshape(1, height as i32)
            .expect("Failed to reshape result matrix")
            .clone_pointee();
        assert!(result.cols() == width as i32);
    
        // Convert to u8 since the VideoWriter seems to require that
        let mut converted = Mat::default();
        let () = result
            .convert_to(&mut converted, CV_8U, 255.0, 0.0)
            .expect("Failure in conversion to u8");
    
        converted
    }
    
    pub fn edge_detection_harness(args: EdgeDetectionInputs) {
        let EdgeDetectionInputs {
            input,
            display,
            output,
            verify,
            display_verify,
            output_verify,
            frames,
        } = args;
    
        let gs: usize = 7;
        let gaussian_filter: Vec<f32> = vec![
            0.000036, 0.000363, 0.001446, 0.002291, 0.001446, 0.000363, 0.000036, 0.000363, 0.003676,
            0.014662, 0.023226, 0.014662, 0.003676, 0.000363, 0.001446, 0.014662, 0.058488, 0.092651,
            0.058488, 0.014662, 0.001446, 0.002291, 0.023226, 0.092651, 0.146768, 0.092651, 0.023226,
            0.002291, 0.001446, 0.014662, 0.058488, 0.092651, 0.058488, 0.014662, 0.001446, 0.000363,
            0.003676, 0.014662, 0.023226, 0.014662, 0.003676, 0.000363, 0.000036, 0.000363, 0.001446,
            0.002291, 0.001446, 0.000363, 0.000036,
        ];
        let gaussian_filter_h = HerculesImmBox::from(gaussian_filter.as_slice());
    
        let sz: usize = 3;
        let structure: Vec<f32> = vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
        let structure_h = HerculesImmBox::from(structure.as_slice());
    
        let sb: usize = 3;
        let sx: Vec<f32> = vec![-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0];
        let sx_h = HerculesImmBox::from(sx.as_slice());
    
        let sy: Vec<f32> = vec![-1.0, -2.0, -1.0, 0.0, 0.0, 0.0, 1.0, 2.0, 1.0];
        let sy_h = HerculesImmBox::from(sy.as_slice());
    
        let theta: f32 = 0.1;
    
        let mut video = VideoCapture::from_file_def(&input).expect("Error loading video");
        assert!(video.is_opened().unwrap());
    
        let fps = video
            .get(VideoCaptureProperties::CAP_PROP_FPS.into())
            .expect("Error getting fps");
    
        let num_frames = video
            .get(VideoCaptureProperties::CAP_PROP_FRAME_COUNT.into())
            .expect("Error getting number of frames") as usize;
        let width = video
            .get(VideoCaptureProperties::CAP_PROP_FRAME_WIDTH.into())
            .expect("Error getting width") as usize;
        let height = video
            .get(VideoCaptureProperties::CAP_PROP_FRAME_HEIGHT.into())
            .expect("Error getting height") as usize;
    
        let num_frames = if let Some(frames) = frames {
            usize::min(frames, num_frames)
        } else {
            num_frames
        };
    
    
    rarbore2's avatar
    rarbore2 committed
        println!(
            "Running edge with {} rows, {} columns, {} gs, {} sz, and {} sb.",
            height, width, gs, sz, sb,
        );
    
    
    rarbore2's avatar
    rarbore2 committed
        let mut r = runner!(edge_detection);
    
        let mut output = output.map(|filename| {
            VideoWriter::new(
                &filename,
                VideoWriter::fourcc('m', 'p', '4', 'v').unwrap(),
                fps,
                Size {
                    width: width as i32,
                    height: height as i32,
                },
                false,
            )
            .expect("Error opening output video")
        });
    
        let mut output_verify = output_verify.map(|filename| {
            VideoWriter::new(
                &filename,
                VideoWriter::fourcc('m', 'p', '4', 'v').unwrap(),
                fps,
                Size {
                    width: width as i32,
                    height: height as i32,
                },
                false,
            )
            .expect("Error opening output video")
        });
    
        for i in 0..num_frames {
            let frame = load_frame(&mut video);
            let ptr = frame.ptr_def().unwrap() as *const f32;
    
            assert!(frame.rows() as usize == height);
            assert!(frame.cols() as usize == width);
    
            let input = unsafe { from_raw_parts(ptr, height * width) };
    
            let input_h = HerculesImmBox::from(input);
    
            let result = async_std::task::block_on(async {
                HerculesMutBox::from(
                    r.run(
                        height as u64,
                        width as u64,
                        input_h.to(),
                        gaussian_filter_h.to(),
                        structure_h.to(),
                        sx_h.to(),
                        sy_h.to(),
                        theta,
                    )
                    .await,
                )
            })
            .as_slice()
            .to_vec();
    
            if display {
                let result = frame_from_slice(&result, height, width);
                let () = imshow("Juno", &result).expect("Failure in displaying image");
            }
            if let Some(ref mut output) = output {
                let result = frame_from_slice(&result, height, width);
                let () = output.write(&result).expect("Failure in writing frame");
            }
    
            if verify {
                let rust_result = edge_detection_rust::edge_detection(
                    height,
                    width,
                    gs,
                    sz,
                    sb,
                    input,
                    &gaussian_filter,
                    &structure,
                    &sx,
                    &sy,
                    theta,
                );
    
    
    rarbore2's avatar
    rarbore2 committed
                let mut all = true;
                for idx in 0..rust_result.len() {
                    if result[idx] != rust_result[idx] {
                        all = false;
                        println!("Found mismatch in images at {}.", idx);
                    }
                }
                assert!(all);
    
    rarbore2's avatar
    rarbore2 committed
                println!("Frames {} match", i);
    
                if display_verify {
                    let rust_result = frame_from_slice(&rust_result, height, width);
                    let () = imshow("Rust", &rust_result).expect("Failure in displaying image");
                }
                if let Some(ref mut output) = output_verify {
                    let result = frame_from_slice(&rust_result, height, width);
                    let () = output.write(&result).expect("Failure in writing frame");
                }
            }
    
            if display || (verify && display_verify) {
                let _ = wait_key(0);
            }
        }
    
        if let Some(mut output) = output {
            let () = output.release().expect("Failure releasing output video");
        }
        if let Some(mut output) = output_verify {
            let () = output.release().expect("Failure releasing output video");
        }
    }