-
Notifications
You must be signed in to change notification settings - Fork 0
Open
Labels
Description
bevy_light_field/src/pipeline.rs
Line 662 in a28416a
| // TODO: parse the json at each frame path to get the bounding boxes |
}
#[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> {