mosure / bevy_light_field Goto Github PK

        if person_timeout.elapsed_secs() > 3.0 {
            person_timeout.reset();

            info!("no person detected for 3 seconds, stop recording");

            let _session_entity = live_session.0.take().unwrap();
            let _raw_streams = stream_manager.stop_recording();

            // TODO: TODO: add a recording finished event when all streams are closed, then execute the following command if pipeline auto-processing is enabled (not ideal for fast recording)

            // commands.entity(session_entity)
            //     .insert(RawStreams {
            //         streams: raw_streams,
            //     })
            //     .insert(PipelineConfig::default());
        }

        person_timeout.tick(time.delta());

}


// TODO: move to MattingPlugin
#[cfg(feature = "person_matting")]
fn create_mask_streams(
    mut commands: Commands,

}


#[derive(Component, Default)]
pub struct YoloFrames {
    pub frames: HashMap<StreamId, Vec<Vec<BoundingBox>>>,
    pub directory: String,
}
impl YoloFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/yolo_frames", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut yolo_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        yolo_frames.reload();

        yolo_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("json"))
                    .map(|entry| std::fs::File::open(entry.path()).unwrap())
                    .map(|yolo_json_file| {
                        let bounding_boxes: Vec<BoundingBox> = serde_json::from_reader(&yolo_json_file).unwrap();

                        bounding_boxes
                    })
                    .collect::<Vec<_>>();

                // TODO: parse the json at each frame path to get the bounding boxes

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn write(&self) {
        self.frames.iter()
            .for_each(|(stream_id, frames)| {
                let output_directory = format!("{}/{}", self.directory, stream_id.0);
                std::fs::create_dir_all(&output_directory).unwrap();

                frames.iter()
                    .enumerate()
                    .for_each(|(frame_idx, bounding_boxes)| {
                        let path = format!("{}/{}.json", output_directory, frame_idx);
                        let _ = serde_json::to_writer(std::fs::File::create(path).unwrap(), bounding_boxes);
                    });
            });
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/yolo_frames", session.directory);
        std::fs::metadata(output_directory).is_ok()
    }

    pub fn image(&self, _camera: usize, _frame: usize) -> Option<Image> {
        todo!()
    }
}



#[derive(Component, Default)]
pub struct RotatedFrames {
    pub frames: HashMap<StreamId, Vec<String>>,
    pub directory: String,
}
impl RotatedFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/rotated_frames", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut raw_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        raw_frames.reload();

        raw_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("png"))
                    .map(|entry| entry.path().to_str().unwrap().to_string())
                    .collect::<Vec<_>>();

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/rotated_frames", session.directory);
        std::fs::metadata(output_directory).is_ok()
    }

    pub fn image(&self, _camera: usize, _frame: usize) -> Option<Image> {
        todo!()
    }
}



#[derive(Component, Default, Reflect)]
pub struct MaskFrames {
    pub frames: HashMap<StreamId, Vec<String>>,
    pub directory: String
}
impl MaskFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/masks", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut mask_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        mask_frames.reload();

        mask_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("png"))
                    .map(|entry| entry.path().to_str().unwrap().to_string())
                    .collect::<Vec<_>>();

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/masks", session.directory);
        std::fs::metadata(output_directory).is_ok()
    }

    pub fn image(&self, _camera: usize, _frame: usize) -> Option<Image> {

}


// TODO: use the async task pool for all generate systems https://crates.io/crates/bevy-async-task
fn generate_raw_frames(
    mut commands: Commands,
    raw_streams: Query<

}


fn generate_rotated_frames(
    mut commands: Commands,
    descriptors: Res<StreamDescriptors>,
    raw_frames: Query<
        (
            Entity,
            &PipelineConfig,
            &RawFrames,
            &Session,
        ),
        Without<RotatedFrames>,
    >,
) {
    // TODO: create a caching/loading system wrapper over run_node interior
    for (
        entity,
        config,
        raw_frames,
        session,
    ) in raw_frames.iter() {
        // TODO: get stream descriptor rotation

        if config.rotate_raw_frames {
            let run_node = !RotatedFrames::exists(session);
            let mut rotated_frames = RotatedFrames::load_from_session(session);

            if run_node {
                let rotations: HashMap<StreamId, f32> = descriptors.0.iter()
                    .enumerate()
                    .map(|(id, descriptor)| (StreamId(id), descriptor.rotation.unwrap_or_default()))
                    .collect();

                info!("generating rotated frames for session {}", session.id);

                raw_frames.frames.iter()
                    .for_each(|(stream_id, frames)| {
                        let output_directory = format!("{}/{}", rotated_frames.directory, stream_id.0);
                        std::fs::create_dir_all(&output_directory).unwrap();

                        let frames = frames.par_iter()
                            .map(|frame| {
                                let frame_idx = std::path::Path::new(frame).file_stem().unwrap().to_str().unwrap();
                                let output_path = format!("{}/{}.png", output_directory, frame_idx);

                                rotate_image(
                                    std::path::Path::new(frame),
                                    std::path::Path::new(&output_path),
                                    rotations[stream_id],
                                ).unwrap();

                                output_path
                            })
                            .collect::<Vec<_>>();

                            rotated_frames.frames.insert(*stream_id, frames);
                    });
            } else {
                info!("rotated frames already exist for session {}", session.id);
            }

            commands.entity(entity).insert(rotated_frames);
        }
    }
}


fn generate_mask_frames(
    mut commands: Commands,
    frames: Query<
        (
            Entity,
            &PipelineConfig,
            &RotatedFrames,
            &Session,
        ),
        Without<MaskFrames>,
    >,
    modnet: Res<Modnet>,
    onnx_assets: Res<Assets<Onnx>>,
) {
    for (
        entity,
        config,
        frames,
        session,
    ) in frames.iter() {
        if config.mask_frames {
            if onnx_assets.get(&modnet.onnx).is_none() {
                return;
            }

            let onnx = onnx_assets.get(&modnet.onnx).unwrap();
            let onnx_session_arc = onnx.session.clone();
            let onnx_session_lock = onnx_session_arc.lock().map_err(|e| e.to_string()).unwrap();
            let onnx_session = onnx_session_lock.as_ref().ok_or("failed to get session from ONNX asset").unwrap();

            let run_node = !MaskFrames::exists(session);
            let mut mask_frames = MaskFrames::load_from_session(session);

            if run_node {
                info!("generating mask frames for session {}", session.id);

                frames.frames.keys()
                    .for_each(|stream_id| {
                        let output_directory = format!("{}/{}", mask_frames.directory, stream_id.0);
                        std::fs::create_dir_all(output_directory).unwrap();
                    });

                let mask_images = frames.frames.iter()
                    .map(|(stream_id, frames)| {
                        let frames = frames.iter()
                            .map(|frame| {
                                let mut decoder = png::Decoder::new(std::fs::File::open(frame).unwrap());
                                decoder.set_transformations(Transformations::EXPAND | Transformations::ALPHA);
                                let mut reader = decoder.read_info().unwrap();
                                let mut img_data = vec![0; reader.output_buffer_size()];
                                let _ = reader.next_frame(&mut img_data).unwrap();

                                assert_eq!(reader.info().bytes_per_pixel(), 3);

                                let width = reader.info().width;
                                let height = reader.info().height;

                                // TODO: separate image loading and onnx inference (so the image loading result can be viewed in the pipeline grid view)
                                let image = Image::new(
                                    Extent3d {
                                        width,
                                        height,
                                        depth_or_array_layers: 1,
                                    },
                                    bevy::render::render_resource::TextureDimension::D2,
                                    img_data,
                                    bevy::render::render_resource::TextureFormat::Rgba8UnormSrgb,
                                    RenderAssetUsages::all(),
                                );

                                let frame_idx = std::path::Path::new(frame).file_stem().unwrap().to_str().unwrap();

                                (
                                    frame_idx,
                                    modnet_inference(
                                        onnx_session,
                                        &[&image],
                                        Some((512, 512)),
                                    ).pop().unwrap(),
                                )
                            })
                            .collect::<Vec<_>>();

                        (stream_id, frames)
                    })
                    .collect::<Vec<_>>();

                mask_images.iter()
                    .for_each(|(stream_id, frames)| {
                        let output_directory = format!("{}/{}", mask_frames.directory, stream_id.0);
                        let mask_paths = frames.iter()
                            .map(|(frame_idx, frame)| {
                                let path = format!("{}/{}.png", output_directory, frame_idx);

                                let buffer = ImageBuffer::<Luma<u8>, Vec<u8>>::from_raw(
                                    frame.width(),
                                    frame.height(),
                                    frame.data.clone(),
                                ).unwrap();

                                let _ = buffer.save(&path);

                                path
                            })
                            .collect::<Vec<_>>();

                        mask_frames.frames.insert(**stream_id, mask_paths);
                    });
            } else {
                info!("mask frames already exist for session {}", session.id);
            }

            commands.entity(entity).insert(mask_frames);
        }
    }
}


fn generate_yolo_frames(
    mut commands: Commands,
    raw_frames: Query<
        (
            Entity,
            &PipelineConfig,
            &RawFrames,
            &Session,
        ),
        Without<YoloFrames>,
    >,
    yolo: Res<Yolo>,
    onnx_assets: Res<Assets<Onnx>>,
) {
    for (
        entity,
        config,
        raw_frames,
        session,
    ) in raw_frames.iter() {
        if config.yolo {
            if onnx_assets.get(&yolo.onnx).is_none() {
                return;
            }

            let onnx = onnx_assets.get(&yolo.onnx).unwrap();
            let onnx_session_arc = onnx.session.clone();
            let onnx_session_lock = onnx_session_arc.lock().map_err(|e| e.to_string()).unwrap();
            let onnx_session = onnx_session_lock.as_ref().ok_or("failed to get session from ONNX asset").unwrap();

            let run_node = !YoloFrames::exists(session);
            let mut yolo_frames = YoloFrames::load_from_session(session);

            if run_node {
                info!("generating yolo frames for session {}", session.id);

                raw_frames.frames.keys()
                    .for_each(|stream_id| {
                        let output_directory = format!("{}/{}", yolo_frames.directory, stream_id.0);
                        std::fs::create_dir_all(output_directory).unwrap();
                    });

                // TODO: support async ort inference (re. progress bars)
                let bounding_box_streams = raw_frames.frames.iter()
                    .map(|(stream_id, frames)| {
                        let frames = frames.iter()
                            .map(|frame| {
                                let mut decoder = png::Decoder::new(std::fs::File::open(frame).unwrap());
                                decoder.set_transformations(Transformations::EXPAND | Transformations::ALPHA);
                                let mut reader = decoder.read_info().unwrap();
                                let mut img_data = vec![0; reader.output_buffer_size()];
                                let _ = reader.next_frame(&mut img_data).unwrap();

                                assert_eq!(reader.info().bytes_per_pixel(), 3);

                                let width = reader.info().width;
                                let height = reader.info().height;

                                // TODO: separate image loading and onnx inference (so the image loading result can be viewed in the pipeline grid view)
                                let image = Image::new(
                                    Extent3d {
                                        width,
                                        height,
                                        depth_or_array_layers: 1,
                                    },
                                    bevy::render::render_resource::TextureDimension::D2,
                                    img_data,
                                    bevy::render::render_resource::TextureFormat::Rgba8UnormSrgb,
                                    RenderAssetUsages::all(),
                                );

                                let frame_idx = std::path::Path::new(frame).file_stem().unwrap().to_str().unwrap();

                                (
                                    frame_idx,
                                    yolo_inference(
                                        onnx_session,
                                        &image,
                                        0.5,
                                    ),
                                )
                            })
                            .collect::<Vec<_>>();

                        (stream_id, frames)
                    })
                    .collect::<Vec<_>>();

                bounding_box_streams.iter()
                    .for_each(|(stream_id, frames)| {
                        let output_directory = format!("{}/{}", yolo_frames.directory, stream_id.0);
                        let bounding_boxes = frames.iter()
                            .map(|(frame_idx, bounding_boxes)| {
                                let path = format!("{}/{}.json", output_directory, frame_idx);

                                let _ = serde_json::to_writer(std::fs::File::create(path).unwrap(), bounding_boxes);

                                bounding_boxes.clone()
                            })
                            .collect::<Vec<_>>();

                        yolo_frames.frames.insert(**stream_id, bounding_boxes);
                    });
            } else {
                info!("yolo frames already exist for session {}", session.id);
            }

            commands.entity(entity).insert(yolo_frames);
        }
    }
}


// TODO: alphablend frames
#[derive(Component, Default)]
pub struct AlphablendFrames {
    pub frames: HashMap<StreamId, Vec<String>>,
    pub directory: String,
}
impl AlphablendFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/alphablend", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut alphablend_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        alphablend_frames.reload();

        alphablend_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("png"))
                    .map(|entry| entry.path().to_str().unwrap().to_string())
                    .collect::<Vec<_>>();

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/alphablend", session.directory);
        std::fs::metadata(output_directory).is_ok()
    }

    pub fn image(&self, _camera: usize, _frame: usize) -> Option<Image> {
        todo!()
    }
}



// TODO: support loading maskframes -> images into a pipeline mask viewer


#[derive(Component, Default)]
pub struct RawFrames {
    pub frames: HashMap<StreamId, Vec<String>>,
    pub directory: String,
}
impl RawFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/frames", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut raw_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        raw_frames.reload();

        raw_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("png"))
                    .map(|entry| entry.path().to_str().unwrap().to_string())
                    .collect::<Vec<_>>();

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn exists(

}


// TODO: add pipeline viewer /w left/right arrow keys and UI controls to switch between frames



fn calculate_grid_dimensions(window_width: f32, window_height: f32, num_streams: usize) -> (usize, usize, f32, f32) {
    let window_aspect_ratio = window_width / window_height;
    let stream_aspect_ratio: f32 = 16.0 / 9.0;

}


fn generate_rotated_frames(
    mut commands: Commands,
    descriptors: Res<StreamDescriptors>,
    raw_frames: Query<
        (
            Entity,
            &PipelineConfig,
            &RawFrames,
            &Session,
        ),
        Without<RotatedFrames>,
    >,
) {
    // TODO: create a caching/loading system wrapper over run_node interior
    for (
        entity,
        config,
        raw_frames,
        session,
    ) in raw_frames.iter() {
        // TODO: get stream descriptor rotation

        if config.rotate_raw_frames {
            let run_node = !RotatedFrames::exists(session);
            let mut rotated_frames = RotatedFrames::load_from_session(session);

            if run_node {
                let rotations: HashMap<StreamId, f32> = descriptors.0.iter()
                    .enumerate()
                    .map(|(id, descriptor)| (StreamId(id), descriptor.rotation.unwrap_or_default()))
                    .collect();

                info!("generating rotated frames for session {}", session.id);

                raw_frames.frames.iter()
                    .for_each(|(stream_id, frames)| {
                        let output_directory = format!("{}/{}", rotated_frames.directory, stream_id.0);
                        std::fs::create_dir_all(&output_directory).unwrap();

                        let frames = frames.par_iter()
                            .map(|frame| {
                                let frame_idx = std::path::Path::new(frame).file_stem().unwrap().to_str().unwrap();
                                let output_path = format!("{}/{}.png", output_directory, frame_idx);

                                rotate_image(
                                    std::path::Path::new(frame),
                                    std::path::Path::new(&output_path),
                                    rotations[stream_id],
                                ).unwrap();

                                output_path
                            })
                            .collect::<Vec<_>>();

                            rotated_frames.frames.insert(*stream_id, frames);
                    });
            } else {
                info!("rotated frames already exist for session {}", session.id);
            }

            commands.entity(entity).insert(rotated_frames);
        }
    }
}


fn generate_mask_frames(
    mut commands: Commands,
    frames: Query<
        (
            Entity,
            &PipelineConfig,
            &RotatedFrames,
            &Session,
        ),
        Without<MaskFrames>,
    >,
    modnet: Res<Modnet>,
    onnx_assets: Res<Assets<Onnx>>,
) {
    for (
        entity,
        config,
        frames,
        session,
    ) in frames.iter() {
        if config.mask_frames {
            if onnx_assets.get(&modnet.onnx).is_none() {
                return;
            }

            let onnx = onnx_assets.get(&modnet.onnx).unwrap();
            let onnx_session_arc = onnx.session.clone();
            let onnx_session_lock = onnx_session_arc.lock().map_err(|e| e.to_string()).unwrap();
            let onnx_session = onnx_session_lock.as_ref().ok_or("failed to get session from ONNX asset").unwrap();

            let run_node = !MaskFrames::exists(session);
            let mut mask_frames = MaskFrames::load_from_session(session);

            if run_node {
                info!("generating mask frames for session {}", session.id);

                frames.frames.keys()
                    .for_each(|stream_id| {
                        let output_directory = format!("{}/{}", mask_frames.directory, stream_id.0);
                        std::fs::create_dir_all(output_directory).unwrap();
                    });

                let mask_images = frames.frames.iter()
                    .map(|(stream_id, frames)| {
                        let frames = frames.iter()
                            .map(|frame| {
                                let mut decoder = png::Decoder::new(std::fs::File::open(frame).unwrap());
                                decoder.set_transformations(Transformations::EXPAND | Transformations::ALPHA);
                                let mut reader = decoder.read_info().unwrap();
                                let mut img_data = vec![0; reader.output_buffer_size()];
                                let _ = reader.next_frame(&mut img_data).unwrap();

                                assert_eq!(reader.info().bytes_per_pixel(), 3);

                                let width = reader.info().width;
                                let height = reader.info().height;

                                // TODO: separate image loading and onnx inference (so the image loading result can be viewed in the pipeline grid view)
                                let image = Image::new(
                                    Extent3d {
                                        width,
                                        height,
                                        depth_or_array_layers: 1,
                                    },
                                    bevy::render::render_resource::TextureDimension::D2,
                                    img_data,
                                    bevy::render::render_resource::TextureFormat::Rgba8UnormSrgb,
                                    RenderAssetUsages::all(),
                                );

                                let frame_idx = std::path::Path::new(frame).file_stem().unwrap().to_str().unwrap();

                                (
                                    frame_idx,
                                    modnet_inference(
                                        onnx_session,
                                        &[&image],
                                        Some((512, 512)),
                                    ).pop().unwrap(),
                                )
                            })
                            .collect::<Vec<_>>();

                        (stream_id, frames)
                    })
                    .collect::<Vec<_>>();

                mask_images.iter()
                    .for_each(|(stream_id, frames)| {
                        let output_directory = format!("{}/{}", mask_frames.directory, stream_id.0);
                        let mask_paths = frames.iter()
                            .map(|(frame_idx, frame)| {
                                let path = format!("{}/{}.png", output_directory, frame_idx);

                                let buffer = ImageBuffer::<Luma<u8>, Vec<u8>>::from_raw(
                                    frame.width(),
                                    frame.height(),
                                    frame.data.clone(),
                                ).unwrap();

                                let _ = buffer.save(&path);

                                path
                            })
                            .collect::<Vec<_>>();

                        mask_frames.frames.insert(**stream_id, mask_paths);
                    });
            } else {
                info!("mask frames already exist for session {}", session.id);
            }

            commands.entity(entity).insert(mask_frames);
        }
    }
}


fn generate_yolo_frames(
    mut commands: Commands,
    raw_frames: Query<
        (
            Entity,
            &PipelineConfig,
            &RawFrames,
            &Session,
        ),
        Without<YoloFrames>,
    >,
    yolo: Res<Yolo>,
    onnx_assets: Res<Assets<Onnx>>,
) {
    for (
        entity,
        config,
        raw_frames,
        session,
    ) in raw_frames.iter() {
        if config.yolo {
            if onnx_assets.get(&yolo.onnx).is_none() {
                return;
            }

            let onnx = onnx_assets.get(&yolo.onnx).unwrap();
            let onnx_session_arc = onnx.session.clone();
            let onnx_session_lock = onnx_session_arc.lock().map_err(|e| e.to_string()).unwrap();
            let onnx_session = onnx_session_lock.as_ref().ok_or("failed to get session from ONNX asset").unwrap();

            let run_node = !YoloFrames::exists(session);
            let mut yolo_frames = YoloFrames::load_from_session(session);

            if run_node {
                info!("generating yolo frames for session {}", session.id);

                raw_frames.frames.keys()
                    .for_each(|stream_id| {
                        let output_directory = format!("{}/{}", yolo_frames.directory, stream_id.0);
                        std::fs::create_dir_all(output_directory).unwrap();
                    });

                // TODO: support async ort inference (re. progress bars)
                let bounding_box_streams = raw_frames.frames.iter()
                    .map(|(stream_id, frames)| {
                        let frames = frames.iter()
                            .map(|frame| {
                                let mut decoder = png::Decoder::new(std::fs::File::open(frame).unwrap());
                                decoder.set_transformations(Transformations::EXPAND | Transformations::ALPHA);
                                let mut reader = decoder.read_info().unwrap();
                                let mut img_data = vec![0; reader.output_buffer_size()];
                                let _ = reader.next_frame(&mut img_data).unwrap();

                                assert_eq!(reader.info().bytes_per_pixel(), 3);

                                let width = reader.info().width;
                                let height = reader.info().height;

                                // TODO: separate image loading and onnx inference (so the image loading result can be viewed in the pipeline grid view)
                                let image = Image::new(
                                    Extent3d {
                                        width,
                                        height,
                                        depth_or_array_layers: 1,
                                    },
                                    bevy::render::render_resource::TextureDimension::D2,
                                    img_data,
                                    bevy::render::render_resource::TextureFormat::Rgba8UnormSrgb,
                                    RenderAssetUsages::all(),
                                );

                                let frame_idx = std::path::Path::new(frame).file_stem().unwrap().to_str().unwrap();

                                (
                                    frame_idx,
                                    yolo_inference(
                                        onnx_session,
                                        &image,
                                        0.5,
                                    ),
                                )
                            })
                            .collect::<Vec<_>>();

                        (stream_id, frames)
                    })
                    .collect::<Vec<_>>();

                bounding_box_streams.iter()
                    .for_each(|(stream_id, frames)| {
                        let output_directory = format!("{}/{}", yolo_frames.directory, stream_id.0);
                        let bounding_boxes = frames.iter()
                            .map(|(frame_idx, bounding_boxes)| {
                                let path = format!("{}/{}.json", output_directory, frame_idx);

                                let _ = serde_json::to_writer(std::fs::File::create(path).unwrap(), bounding_boxes);

                                bounding_boxes.clone()
                            })
                            .collect::<Vec<_>>();

                        yolo_frames.frames.insert(**stream_id, bounding_boxes);
                    });
            } else {
                info!("yolo frames already exist for session {}", session.id);
            }

            commands.entity(entity).insert(yolo_frames);
        }
    }
}


// TODO: alphablend frames
#[derive(Component, Default)]
pub struct AlphablendFrames {
    pub frames: HashMap<StreamId, Vec<String>>,
    pub directory: String,
}
impl AlphablendFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/alphablend", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut alphablend_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        alphablend_frames.reload();

        alphablend_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("png"))
                    .map(|entry| entry.path().to_str().unwrap().to_string())
                    .collect::<Vec<_>>();

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/alphablend", session.directory);
        std::fs::metadata(output_directory).is_ok()
    }

    pub fn image(&self, _camera: usize, _frame: usize) -> Option<Image> {
        todo!()
    }
}



// TODO: support loading maskframes -> images into a pipeline mask viewer


#[derive(Component, Default)]
pub struct RawFrames {
    pub frames: HashMap<StreamId, Vec<String>>,
    pub directory: String,
}
impl RawFrames {
    pub fn load_from_session(
        session: &Session,
    ) -> Self {
        let directory = format!("{}/frames", session.directory);
        std::fs::create_dir_all(&directory).unwrap();

        let mut raw_frames = Self {
            frames: HashMap::new(),
            directory,
        };
        raw_frames.reload();

        raw_frames
    }

    pub fn reload(&mut self) {
        std::fs::read_dir(&self.directory)
            .unwrap()
            .filter_map(|entry| entry.ok())
            .filter(|entry| entry.path().is_dir())
            .map(|stream_dir| {
                let stream_id = StreamId(stream_dir.path().file_name().unwrap().to_str().unwrap().parse::<usize>().unwrap());

                let frames = std::fs::read_dir(stream_dir.path()).unwrap()
                    .filter_map(|entry| entry.ok())
                    .filter(|entry| entry.path().is_file() && entry.path().extension().and_then(|s| s.to_str()) == Some("png"))
                    .map(|entry| entry.path().to_str().unwrap().to_string())
                    .collect::<Vec<_>>();

                (stream_id, frames)
            })
            .for_each(|(stream_id, frames)| {
                self.frames.insert(stream_id, frames);
            });
    }

    pub fn exists(

        ..default()
    };

    // TODO: support enabling/disabling decoding/matting per stream (e.g. during 'record mode')

    let input_images: Vec<Handle<Image>> = stream_uris.0.iter()
        .enumerate()
        .map(|(index, descriptor)| {
            let entity = commands.spawn_empty().id();

            let mut image = Image {

];


// TODO: add bevy_args
fn main() {
    let primary_window = Some(Window {
        mode: bevy::window::WindowMode::Windowed,

use bevy::prelude::*;
use image::DynamicImage;
use rayon::prelude::*;


pub struct PipelinePlugin;
impl Plugin for PipelinePlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Update, generate_annotations);
    }
}


#[derive(Component, Reflect)]
pub struct PipelineConfig {
    pub raw_frames: bool,
    pub subject_refinement: bool,           // https://github.com/onnx/models/tree/main/validated/vision/body_analysis/ultraface
    pub repair_frames: bool,                // https://huggingface.co/docs/diffusers/en/optimization/onnx & https://github.com/bnm6900030/swintormer
    pub upsample_frames: bool,              // https://huggingface.co/ssube/stable-diffusion-x4-upscaler-onnx
    pub mask_frames: bool,                  // https://github.com/ZHKKKe/MODNet
    pub light_field_cameras: bool,          // https://github.com/jasonyzhang/RayDiffusion
    pub depth_maps: bool,                   // https://github.com/fabio-sim/Depth-Anything-ONNX
    pub gaussian_cloud: bool,
}

impl Default for PipelineConfig {
    fn default() -> Self {
        Self {
            raw_frames: true,
            subject_refinement: false,
            repair_frames: false,
            upsample_frames: false,
            mask_frames: false,
            light_field_cameras: false,
            depth_maps: false,
            gaussian_cloud: false,
        }
    }
}


#[derive(Bundle, Default, Reflect)]
pub struct StreamSessionBundle {
    pub config: PipelineConfig,
    pub raw_streams: RawStreams,
    pub session: Session,
}

// TODO: use an entity saver to write Session and it's components (e.g. `0/session.ron`)


#[derive(Component, Default, Reflect)]
pub struct Session {
    pub id: usize,
    pub directory: String,
}

impl Session {
    pub fn new(directory: String) -> Self {
        let id = get_next_session_id(&directory);
        let directory = format!("{}/{}", directory, id);
        std::fs::create_dir_all(&directory).unwrap();

        Self { id, directory }
    }
}


pub trait PipelineNode {
    fn new(session: &Session) -> Self;
    fn exists(session: &Session) -> bool;
}


#[derive(Component, Default, Reflect)]
pub struct RawStreams {
    pub streams: Vec<String>,
}
impl RawStreams {
    pub fn decoders(&self) -> () {
        // TODO: create decoders for each h264 file stream
        todo!()
    }
}


// TODO: add a pipeline config for the session describing which components to run


fn generate_annotations(
    mut commands: Commands,
    raw_streams: Query<
        (
            Entity,
            &PipelineConfig,
            &RawStreams,
            &Session,
        ),
        Without<RawFrames>,
    >,
) {
    for (
        entity,
        config,
        raw_streams,
        session,
    ) in raw_streams.iter() {
        let raw_frames = RawFrames::new(session);

        if config.raw_frames {
            let run_node = !RawFrames::exists(session);
            let mut raw_frames = RawFrames::new(session);

            if run_node {
                // TODO: add (identifier == "{}_{}", stream, frame) output to this first stage (forward to mask stage /w image in RAM)
                let frames = raw_streams.streams.par_iter()
                    .map(|stream| {
                        let decoder = todo!();

                        (stream, decoder)
                    })
                    .map(|(stream, decoder)| {
                        let frames: Vec<DynamicImage> = vec![];

                        // TODO: read all frames

                        (stream, frames)
                    })
                    .collect::<Vec<_>>();

                frames.par_iter()
                    .for_each(|(stream, frames)| {
                        frames.iter().enumerate()
                            .for_each(|(i, frame)| {
                                let path = format!(
                                    "{}/frames/{}_{}.png",
                                    session.directory,
                                    stream,
                                    i,
                                );
                                frame.save(path).unwrap();
                            });
                    });
            }
        }

        commands.entity(entity).insert(raw_frames);
    }
}


#[derive(Component, Default, Reflect)]
pub struct RawFrames {
    pub frames: Vec<String>,
}
impl RawFrames {
    // TODO: move new and exists to a trait

    pub fn new(
        session: &Session,
    ) -> Self {
        let output_directory = format!("{}/frames", session.directory);
        std::fs::create_dir_all(&output_directory).unwrap();

        // TODO: load all files that are already in the directory

        Self {
            frames: vec![],
        }
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/frames", session.directory);
        std::fs::metadata(&output_directory).is_ok()
    }

    pub fn image(&self, camera: usize, frame: usize) -> Option<Image> {
        todo!()
    }
}


#[derive(Component, Default, Reflect)]
pub struct MaskFrames {
    pub frames: Vec<String>,
}
impl MaskFrames {
    pub fn new(
        session: &Session,
    ) -> Self {
        let output_directory = format!("{}/masks", session.directory);
        std::fs::create_dir_all(&output_directory).unwrap();

        Self {
            frames: vec![],
        }
    }

    pub fn image(&self, camera: usize, frame: usize) -> Option<Image> {
        todo!()
    }
}


#[derive(Default, Clone, Reflect)]
pub struct LightFieldCamera {
    // TODO: intrinsics/extrinsics
}

#[derive(Component, Default, Reflect)]
pub struct LightFieldCameras {
    pub cameras: Vec<LightFieldCamera>,
}



fn get_next_session_id(output_directory: &str) -> usize {
    match std::fs::read_dir(output_directory) {
        Ok(entries) => entries.filter_map(|entry| {
            let entry = entry.ok()?;
                if entry.path().is_dir() {
                    entry.file_name().to_string_lossy().parse::<usize>().ok()
                } else {
                    None
                }
            })
            .max()
            .map_or(0, |max_id| max_id + 1),
        Err(_) => 0,
    }
}

#[derive(Component)]
pub struct DetectPersons;  // TODO: add option for mask vs yolo detection


#[derive(Debug, Clone, Reflect, PartialEq)]

use bevy::prelude::*;
use image::DynamicImage;
use rayon::prelude::*;


pub struct PipelinePlugin;
impl Plugin for PipelinePlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(Update, generate_annotations);
    }
}


#[derive(Component, Reflect)]
pub struct PipelineConfig {
    pub raw_frames: bool,
    pub subject_refinement: bool,           // https://github.com/onnx/models/tree/main/validated/vision/body_analysis/ultraface
    pub repair_frames: bool,                // https://huggingface.co/docs/diffusers/en/optimization/onnx & https://github.com/bnm6900030/swintormer
    pub upsample_frames: bool,              // https://huggingface.co/ssube/stable-diffusion-x4-upscaler-onnx
    pub mask_frames: bool,                  // https://github.com/ZHKKKe/MODNet
    pub light_field_cameras: bool,          // https://github.com/jasonyzhang/RayDiffusion
    pub depth_maps: bool,                   // https://github.com/fabio-sim/Depth-Anything-ONNX
    pub gaussian_cloud: bool,
}

impl Default for PipelineConfig {
    fn default() -> Self {
        Self {
            raw_frames: true,
            subject_refinement: false,
            repair_frames: false,
            upsample_frames: false,
            mask_frames: false,
            light_field_cameras: false,
            depth_maps: false,
            gaussian_cloud: false,
        }
    }
}


#[derive(Bundle, Default, Reflect)]
pub struct StreamSessionBundle {
    pub config: PipelineConfig,
    pub raw_streams: RawStreams,
    pub session: Session,
}

// TODO: use an entity saver to write Session and it's components (e.g. `0/session.ron`)


#[derive(Component, Default, Reflect)]
pub struct Session {
    pub id: usize,
    pub directory: String,
}

impl Session {
    pub fn new(directory: String) -> Self {
        let id = get_next_session_id(&directory);
        let directory = format!("{}/{}", directory, id);
        std::fs::create_dir_all(&directory).unwrap();

        Self { id, directory }
    }
}


pub trait PipelineNode {
    fn new(session: &Session) -> Self;
    fn exists(session: &Session) -> bool;
}


#[derive(Component, Default, Reflect)]
pub struct RawStreams {
    pub streams: Vec<String>,
}
impl RawStreams {
    pub fn decoders(&self) -> () {
        // TODO: create decoders for each h264 file stream
        todo!()
    }
}


// TODO: add a pipeline config for the session describing which components to run


fn generate_annotations(
    mut commands: Commands,
    raw_streams: Query<
        (
            Entity,
            &PipelineConfig,
            &RawStreams,
            &Session,
        ),
        Without<RawFrames>,
    >,
) {
    for (
        entity,
        config,
        raw_streams,
        session,
    ) in raw_streams.iter() {
        let raw_frames = RawFrames::new(session);

        if config.raw_frames {
            let run_node = !RawFrames::exists(session);
            let mut raw_frames = RawFrames::new(session);

            if run_node {
                // TODO: add (identifier == "{}_{}", stream, frame) output to this first stage (forward to mask stage /w image in RAM)
                let frames = raw_streams.streams.par_iter()
                    .map(|stream| {
                        let decoder = todo!();

                        (stream, decoder)
                    })
                    .map(|(stream, decoder)| {
                        let frames: Vec<DynamicImage> = vec![];

                        // TODO: read all frames

                        (stream, frames)
                    })
                    .collect::<Vec<_>>();

                frames.par_iter()
                    .for_each(|(stream, frames)| {
                        frames.iter().enumerate()
                            .for_each(|(i, frame)| {
                                let path = format!(
                                    "{}/frames/{}_{}.png",
                                    session.directory,
                                    stream,
                                    i,
                                );
                                frame.save(path).unwrap();
                            });
                    });
            }
        }

        commands.entity(entity).insert(raw_frames);
    }
}


#[derive(Component, Default, Reflect)]
pub struct RawFrames {
    pub frames: Vec<String>,
}
impl RawFrames {
    // TODO: move new and exists to a trait

    pub fn new(
        session: &Session,
    ) -> Self {
        let output_directory = format!("{}/frames", session.directory);
        std::fs::create_dir_all(&output_directory).unwrap();

        // TODO: load all files that are already in the directory

        Self {
            frames: vec![],
        }
    }

    pub fn exists(
        session: &Session,
    ) -> bool {
        let output_directory = format!("{}/frames", session.directory);
        std::fs::metadata(&output_directory).is_ok()
    }

    pub fn image(&self, camera: usize, frame: usize) -> Option<Image> {
        todo!()
    }
}


#[derive(Component, Default, Reflect)]
pub struct MaskFrames {
    pub frames: Vec<String>,
}
impl MaskFrames {
    pub fn new(
        session: &Session,
    ) -> Self {
        let output_directory = format!("{}/masks", session.directory);
        std::fs::create_dir_all(&output_directory).unwrap();

        Self {
            frames: vec![],
        }
    }

    pub fn image(&self, camera: usize, frame: usize) -> Option<Image> {
        todo!()
    }
}


#[derive(Default, Clone, Reflect)]
pub struct LightFieldCamera {
    // TODO: intrinsics/extrinsics
}

#[derive(Component, Default, Reflect)]
pub struct LightFieldCameras {
    pub cameras: Vec<LightFieldCamera>,
}



fn get_next_session_id(output_directory: &str) -> usize {
    match std::fs::read_dir(output_directory) {
        Ok(entries) => entries.filter_map(|entry| {
            let entry = entry.ok()?;
                if entry.path().is_dir() {
                    entry.file_name().to_string_lossy().parse::<usize>().ok()
                } else {
                    None
                }
            })
            .max()
            .map_or(0, |max_id| max_id + 1),
        Err(_) => 0,
    }
}

    },
    codec::VideoFrame,
};
use url::Url;


const RTSP_URIS: [&str; 2] = [
    // "rtsp://localhost:554/lizard",
    // "rtsp://rtspstream:[email protected]/movie",
    // "rtsp://rtspstream:[email protected]/pattern",
    "rtsp://192.168.1.23/user=admin&password=admin123&channel=1&stream=0.sdp?",
    "rtsp://192.168.1.24/user=admin&password=admin123&channel=1&stream=0.sdp?",
];


fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_systems(Update, press_esc_close)
        .add_systems(Startup, spawn_tasks)
        .add_systems(Update, task_loop)
        .add_systems(Update, handle_tasks)
        .run();
}


#[derive(Component)]
struct ReadRtspFrame(Task<CommandQueue>);

#[derive(Component, Clone)]
struct RtspStream {
    uri: String,
    index: usize,
    demuxed: Arc<Mutex<Option<Demuxed>>>,
    decoder: Arc<Mutex<Option<Decoder>>>,
}


// TODO: decouple bevy async tasks and the multi-stream RTSP handling
//       a bevy system should query the readiness of a frame for each RTSP stream and update texture as needed
//       the RTSP handling should be a separate async tokio pool that updates the readiness of each RTSP stream /w buffer for transfer to texture


fn spawn_tasks(
    mut commands: Commands,
) {
    RTSP_URIS.iter()
        .enumerate()
        .for_each(|(index, &url)| {
            let entity = commands.spawn_empty().id();

            let api = openh264::OpenH264API::from_source();
            let decoder = Decoder::new(api).unwrap();

            let rtsp_stream = RtspStream {
                uri: url.to_string(),
                index,
                demuxed: Arc::new(Mutex::new(None)),
                decoder: Arc::new(Mutex::new(decoder.into())),
            };
            let demuxed_arc = rtsp_stream.demuxed.clone();

            let session_result = futures::executor::block_on(Compat::new(create_session(&url)));
            match session_result {
                Ok(playing) => {
                    let mut demuxed = demuxed_arc.lock().unwrap();
                    *demuxed = Some(playing.demuxed().unwrap());

                    println!("created demuxer for {}", url);
                },
                Err(e) => panic!("Failed to create session: {}", e),
            }

            queue_rtsp_frame(
                &rtsp_stream,
                &mut commands,
                entity,
            );

            commands.entity(entity).insert(rtsp_stream);
        });
}


fn task_loop(
    mut commands: Commands,
    streams: Query<
        (
            Entity,
            &RtspStream,
        ),
        Without<ReadRtspFrame>,
    >,
) {
    for (entity, rtsp_stream) in &mut streams.iter() {
        queue_rtsp_frame(
            rtsp_stream,
            &mut commands,
            entity,
        );
    }
}


fn handle_tasks(
    mut commands: Commands,
    mut tasks: Query<&mut ReadRtspFrame>,
) {
    for mut task in &mut tasks {
        if let Some(mut commands_queue) = block_on(future::poll_once(&mut task.0)) {
            commands.append(&mut commands_queue);
        }
    }
}


fn convert_h264(data: &mut [u8]) -> Result<(), Error> {
    let mut i = 0;
    while i < data.len() - 3 {
        // Replace each NAL's length with the Annex B start code b"\x00\x00\x00\x01".
        let len = u32::from_be_bytes([data[i], data[i + 1], data[i + 2], data[i + 3]]) as usize;
        data[i] = 0;
        data[i + 1] = 0;
        data[i + 2] = 0;
        data[i + 3] = 1;
        i += 4 + len;
        if i > data.len() {
            bail!("partial NAL body");
        }
    }
    if i < data.len() {
        bail!("partial NAL length");
    }
    Ok(())
}

fn queue_rtsp_frame(
    rtsp_stream: &RtspStream,
    commands: &mut Commands,
    entity: Entity,
) {
    let idx = rtsp_stream.index;
    let demuxed_arc = rtsp_stream.demuxed.clone();
    let decoder_arc = rtsp_stream.decoder.clone();

    let task = AsyncComputeTaskPool::get().spawn(async move {
        let result = futures::executor::block_on(Compat::new(capture_frame(demuxed_arc.clone())));

        // TODO: clean this mess
        match result {
            Ok(frame) => {
                let mut decoder_lock = decoder_arc.lock().unwrap();
                let decoder = decoder_lock.as_mut().unwrap();

                let mut data = frame.into_data();
                let annex_b_convert = convert_h264(&mut data);

                match annex_b_convert {
                    Ok(_) => for packet in nal_units(&data) {
                        let result = decoder.decode(packet);

                        match result {
                            Ok(decoded_frame) => {
                                println!("decoded frame from stream {}", idx);
                                // TODO: populate 'latest frame' with new data
                            },
                            Err(e) => println!("failed to decode frame for stream {}: {}", idx, e),
                        }
                    },
                    Err(e) => println!("failed to convert NAL unit to Annex B format: {}", e)
                }
            },
            Err(e) => println!("failed to capture frame for stream {}: {}", idx, e),
        }

        let mut command_queue = CommandQueue::default();
        command_queue.push(move |world: &mut World| {
            world.entity_mut(entity).remove::<ReadRtspFrame>();
        });

        command_queue
    });

    commands.entity(entity).insert(ReadRtspFrame(task));
}

async fn create_session(url: &str) -> Result<Session<Playing>, Box<dyn std::error::Error + Send + Sync>> {
    let parsed_url = Url::parse(url)?;

    let username = parsed_url.username();