From 48e7e6487cc838cb14c16b19a01256932e1ee6c0 Mon Sep 17 00:00:00 2001 From: 4o3F <4o3f@proton.me> Date: Sat, 5 Oct 2024 14:54:18 +0800 Subject: [PATCH] feat: change split image with filter behaviour --- .github/workflows/build.yml | 109 +++++----- Cargo.toml | 2 +- src/common/augment.rs | 418 +++++++++++++++++------------------- src/common/operation.rs | 69 +++++- src/main.rs | 30 ++- 5 files changed, 353 insertions(+), 275 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d977ddd..1117669 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -22,60 +22,71 @@ on: # A workflow run is made up of one or more jobs that can run sequentially or in parallel jobs: # Builds for linux - # ubuntu-build: - # # The type of runner that the job will run on - # runs-on: ubuntu-latest + ubuntu-build: + # The type of runner that the job will run on + runs-on: ubuntu-latest - # # Steps represent a sequence of tasks that will be executed as part of the job - # steps: - # # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it - # - uses: actions/checkout@v3 + # Steps represent a sequence of tasks that will be executed as part of the job + steps: + # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it + - uses: actions/checkout@v3 - # # Install rust - # - name: Install rust - # run: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y - - # # Install opencv - # #- name: Install opencv - # # run: sudo apt install libopencv-dev clang libclang-dev + # Install rust + - name: Install rust + run: curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y - # # Pull opencv source, unzip, configure cmake, build, and install - # # Disable image formats like jpeg, png, tiff, as we use rust image crate instead. See https://docs.opencv.org/4.x/db/d05/tutorial_config_reference.html - # # Inspired from https://github.com/twistedfall/opencv-rust/issues/364 - # - name: Install OpenCV static lib - # run: | - # wget -O opencv.zip https://github.com/opencv/opencv/archive/refs/tags/4.8.1.zip - # wget -O opencv_contrib.zip https://github.com/opencv/opencv_contrib/archive/refs/tags/4.8.1.zip - # unzip opencv.zip && rm opencv.zip - # unzip opencv_contrib.zip && rm opencv_contrib.zip - # mkdir -p build && cd build - # cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=NO -DCMAKE_INSTALL_PREFIX=/opt/opencv -DBUILD_DOCS=OFF -DBUILD_EXAMPLES=OFF -DBUILD_TESTS=OFF -DBUILD_PERF_TESTS=OFF -DWITH_PNG=OFF -DWITH_JPEG=OFF -DWITH_TIFF=OFF -DWITH_WEBP=OFF -DWITH_OPENJPEG=OFF -DWITH_JASPER=OFF -DWITH_OPENEXR=OFF -DWITH_V4L=OFF -DBUILD_opencv_java=OFF -DBUILD_opencv_python=OFF -DOPENCV_EXTRA_MODULES_PATH=../opencv_contrib-4.8.1/modules ../opencv-4.8.1 - # cmake --build . --target install --config Release --parallel 8 - # cmake --install . --prefix /opt/opencv - # cd .. + # Install opencv + #- name: Install opencv + # run: sudo apt install libopencv-dev clang libclang-dev + - name: Cache OpenCV build + uses: actions/cache@v4 + with: + path: | + /opt/opencv + key: ${{ runner.os }}-opencv-${{ hashFiles('**/CMakeLists.txt') }} # 使用 CMakeLists.txt 的哈希作为缓存键 + restore-keys: | + ${{ runner.os }}-opencv- + # Pull opencv source, unzip, configure cmake, build, and install + # Disable image formats like jpeg, png, tiff, as we use rust image crate instead. See https://docs.opencv.org/4.x/db/d05/tutorial_config_reference.html + # Inspired from https://github.com/twistedfall/opencv-rust/issues/364 + - name: Install OpenCV static lib + run: | + wget -O opencv.zip https://github.com/opencv/opencv/archive/refs/tags/4.8.1.zip + wget -O opencv_contrib.zip https://github.com/opencv/opencv_contrib/archive/refs/tags/4.8.1.zip + unzip opencv.zip && rm opencv.zip + unzip opencv_contrib.zip && rm opencv_contrib.zip + mkdir -p build && cd build + cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=NO -DCMAKE_INSTALL_PREFIX=/opt/opencv -DBUILD_DOCS=OFF -DBUILD_EXAMPLES=OFF -DBUILD_TESTS=OFF -DBUILD_PERF_TESTS=OFF -DWITH_WEBP=OFF -DWITH_JASPER=OFF -DWITH_OPENEXR=OFF -DWITH_V4L=OFF -DBUILD_opencv_java=OFF -DBUILD_opencv_python=OFF -DOPENCV_EXTRA_MODULES_PATH=../opencv_contrib-4.8.1/modules ../opencv-4.8.1 + cmake --build . --target install --config Release --parallel 8 + cmake --install . --prefix /opt/opencv + cd .. - # # Note: OPENCV_LINK_LIBS ordering matters for linux. Put lower level deps after higher level. See https://answers.opencv.org/question/186124/undefined-reference-to-cvsoftdoubleoperator/ - # # libclang files are in /usr/lib/llvm-##/lib. We symlink it to one of the opencv_link_paths - # # OpenCV-rust looks for "opencv2/core/version.hpp" for the OpenCV version: https://github.com/twistedfall/opencv-rust/issues/368 - # # which is under //include/opencv4 for linux - # # Build - # - name: Build - # run: | - # export OPENCV_LINK_LIBS="opencv_videoio,opencv_imgcodecs,opencv_imgproc,opencv_core,libippiw,libittnotify,libippicv,z" - # export OPENCV_LINK_PATHS=/opt/opencv/lib,/opt/opencv/lib/opencv4/3rdparty,/usr/lib/x86_64-linux-gnu - # export OPENCV_INCLUDE_PATHS=/opt/opencv/include,/opt/opencv/include/opencv4 - # sudo ln -s /usr/lib/llvm-15/lib/libclang.so.1 /usr/lib/x86_64-linux-gnu/libclang.so - # ls -R /opt/opencv - # ls -R /usr/lib - # cargo build --release + - name: Setup tmate session + uses: mxschmitt/action-tmate@v3 + if: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_enabled }} + + # Note: OPENCV_LINK_LIBS ordering matters for linux. Put lower level deps after higher level. See https://answers.opencv.org/question/186124/undefined-reference-to-cvsoftdoubleoperator/ + # libclang files are in /usr/lib/llvm-##/lib. We symlink it to one of the opencv_link_paths + # OpenCV-rust looks for "opencv2/core/version.hpp" for the OpenCV version: https://github.com/twistedfall/opencv-rust/issues/368 + # which is under //include/opencv4 for linux + # Build + - name: Build + run: | + export OPENCV_LINK_LIBS="opencv_videoio,opencv_imgcodecs,opencv_imgproc,opencv_core,libippiw,libittnotify,libippicv,libippicvmt,zlib,IlmImf,libjpeg-turbo,libopenjp2,libpng,libtiff,libwebp" + export OPENCV_LINK_PATHS=/opt/opencv/lib,/opt/opencv/lib/opencv4/3rdparty,/usr/lib/x86_64-linux-gnu + export OPENCV_INCLUDE_PATHS=/opt/opencv/include,/opt/opencv/include/opencv4 + sudo ln -s /usr/lib/llvm-15/lib/libclang.so.1 /usr/lib/x86_64-linux-gnu/libclang.so + ls -R /opt/opencv + ls -R /usr/lib + cargo build --release --target x86_64-unknown-linux-musl - # # Upload artifact: https://github.com/actions/upload-artifact - # - name: Upload Artifacts - # uses: actions/upload-artifact@v3 - # with: - # name: imagetools - # path: target/release/imagetools - # retention-days: 14 + # Upload artifact: https://github.com/actions/upload-artifact + - name: Upload Artifacts + uses: actions/upload-artifact@v3 + with: + name: imagetools + path: target/release/imagetools + retention-days: 14 windows-build: # The type of runner that the job will run on diff --git a/Cargo.toml b/Cargo.toml index e5f6260..cf09aa5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,7 +18,7 @@ opencv = { version = "0.89.0", default-features = false, features = [ "imgproc", "imgcodecs", ] } -# gdal = "0.17.0" +# gdal = "0.17" itertools = "0.12.1" clap = { version = "4.5.9", features = ["derive"] } tracing-subscriber = "0.3.18" diff --git a/src/common/augment.rs b/src/common/augment.rs index 1715d6c..51cf774 100644 --- a/src/common/augment.rs +++ b/src/common/augment.rs @@ -1,5 +1,6 @@ use std::{ - fs, + fs::{self}, + path::PathBuf, sync::{Arc, RwLock}, }; @@ -17,6 +18,9 @@ pub async fn split_images(dataset_path: &String, target_height: &u32, target_wid for entry in entries { let entry = entry.expect_or_log("Failed to iterate entries"); + if entry.path().is_dir() { + continue; + } let dataset_path = dataset_path.clone(); let target_height = *target_height; let target_width = *target_width; @@ -62,27 +66,27 @@ pub async fn split_images(dataset_path: &String, target_height: &u32, target_wid let (width, height) = (size.width as u32, size.height as u32); - let vertical_count = height / target_height; - let horizontal_count = width / target_width; + let y_count = height / target_height; + let x_count = width / target_width; tracing::trace!( "Image {} vertical_count {} horizontal_count {}", file_name, - vertical_count, - horizontal_count + y_count, + x_count ); // LTR - 'outer: for horizontal_index in 0..horizontal_count { - 'inner: for vertical_index in 0..vertical_count { - let start_y = horizontal_index * target_height; - let start_x = vertical_index * target_width; + 'outer: for x_index in 0..x_count { + 'inner: for y_index in 0..y_count { + let start_x = x_index * target_width; + let start_y = y_index * target_height; if start_x + target_width > width { - continue 'inner; + continue 'outer; } if start_y + target_height > height { - continue 'outer; + continue 'inner; } let cropped_img = core::Mat::roi( @@ -96,8 +100,8 @@ pub async fn split_images(dataset_path: &String, target_height: &u32, target_wid ) .expect_or_log("Failed to crop image"); let path = format!( - "{}\\..\\output\\{}_LTR_h{}_v{}.{}", - dataset_path, file_stem, horizontal_index, vertical_index, file_extension + "{}\\..\\output\\{}_LTR_x{}_y{}.{}", + dataset_path, file_stem, x_index, y_index, file_extension ); let result = imgcodecs::imwrite(path.as_str(), &cropped_img, &core::Vector::new()) @@ -116,22 +120,22 @@ pub async fn split_images(dataset_path: &String, target_height: &u32, target_wid tracing::info!("Image LTR {} done", entry.file_name().to_str().unwrap()); // RTL - 'outer: for horizontal_index in 0..horizontal_count { - 'inner: for vertical_index in 0..vertical_count { - let start_y = height as i32 - (horizontal_index * target_height) as i32; - let start_x = width as i32 - (vertical_index * target_width) as i32; + 'outer: for x_index in 0..x_count { + 'inner: for y_index in 0..y_count { + let start_y = height as i32 - (y_index * target_height) as i32; + let start_x = width as i32 - (x_index * target_width) as i32; if start_x < 0 || start_x as u32 >= width || start_x as u32 + target_width > width { - continue 'inner; + continue 'outer; } if start_y < 0 || start_y as u32 >= height || start_y as u32 + target_height > height { - continue 'outer; + continue 'inner; } let cropped_img = core::Mat::roi( @@ -145,8 +149,8 @@ pub async fn split_images(dataset_path: &String, target_height: &u32, target_wid ) .expect_or_log("Failed to crop image"); let path = format!( - "{}\\..\\output\\{}_RTL_h{}_v{}.{}", - dataset_path, file_stem, horizontal_index, vertical_index, file_extension + "{}\\..\\output\\{}_RTL_x{}_y{}.{}", + dataset_path, file_stem, x_index, y_index, file_extension ); let result = imgcodecs::imwrite(path.as_str(), &cropped_img, &core::Vector::new()) @@ -209,22 +213,22 @@ pub async fn split_images_with_bias( let img = imgcodecs::imread(&entry_path, imgcodecs::IMREAD_UNCHANGED).unwrap(); let size = img.size().unwrap(); let (width, height) = (size.width as u32, size.height as u32); - let vertical_count = height / target_height; - let horizontal_count = width / target_width; - let bias_max = horizontal_count.max(vertical_count); + let y_count = height / target_height; + let x_count = width / target_width; + let bias_max = x_count.max(y_count); // LTR for bias in 0..bias_max { - 'outer: for horizontal_index in 0..horizontal_count { - 'inner: for vertical_index in 0..vertical_count { - let start_y = horizontal_index * target_height + bias * bias_step; - let start_x = vertical_index * target_width + bias * bias_step; + 'outer: for x_index in 0..x_count { + 'inner: for y_index in 0..y_count { + let start_y = y_index * target_height + bias * bias_step; + let start_x = x_index * target_width + bias * bias_step; if start_x + target_width > width { - continue 'inner; + continue 'outer; } if start_y + target_height > height { - continue 'outer; + continue 'inner; } let cropped_img = core::Mat::roi( @@ -239,12 +243,12 @@ pub async fn split_images_with_bias( .unwrap(); imgcodecs::imwrite( format!( - "{}\\..\\output\\{}_LTR_bias{}_h{}_v{}.{}", + "{}\\..\\output\\{}_LTR_bias{}_x{}_y{}.{}", dataset_path, entry.path().file_stem().unwrap().to_str().unwrap(), bias, - horizontal_index, - vertical_index, + x_index, + y_index, entry.path().extension().unwrap().to_str().unwrap() ) .as_str(), @@ -263,24 +267,24 @@ pub async fn split_images_with_bias( // RTL for bias in 0..bias_max { - 'outer: for horizontal_index in 0..horizontal_count { - 'inner: for vertical_index in 0..vertical_count { - let start_y = height as i32 - - (horizontal_index * target_height + bias * bias_step) as i32; - let start_x = width as i32 - - (vertical_index * target_width + bias * bias_step) as i32; + 'outer: for x_index in 0..x_count { + 'inner: for y_index in 0..y_count { + let start_y = + height as i32 - (y_index * target_height + bias * bias_step) as i32; + let start_x = + width as i32 - (x_index * target_width + bias * bias_step) as i32; if start_x < 0 || start_x as u32 >= width || start_x as u32 + target_width > width { - continue 'inner; + continue 'outer; } if start_y < 0 || start_y as u32 >= height || start_y as u32 + target_height > height { - continue 'outer; + continue 'inner; } let cropped_img = core::Mat::roi( @@ -295,12 +299,12 @@ pub async fn split_images_with_bias( .unwrap(); imgcodecs::imwrite( format!( - "{}\\..\\output\\{}_RTL_bias{}_h{}_v{}.png", + "{}\\..\\output\\{}_RTL_bias{}_x{}_y{}.png", dataset_path, entry.file_name().to_str().unwrap().replace(".png", ""), bias, - horizontal_index, - vertical_index + x_index, + y_index ) .as_str(), &cropped_img, @@ -322,7 +326,7 @@ pub async fn split_images_with_bias( while threads.join_next().await.is_some() {} } -pub fn check_valid_pixel_count(img: &Mat, valid_rgb_list: &[core::Vec3b]) -> bool { +pub fn check_valid_pixel_count(img: &Mat, rgb_list: &[core::Vec3b], mode: bool) -> bool { let mut count = 0; let size = img.size().unwrap(); let (width, height) = (size.width, size.height); @@ -330,8 +334,17 @@ pub fn check_valid_pixel_count(img: &Mat, valid_rgb_list: &[core::Vec3b]) -> boo for y in 0..height { for x in 0..width { let pixel = img.at_2d::(y, x).unwrap(); - if valid_rgb_list.contains(&pixel) { - count += 1; + match mode { + true => { + if rgb_list.contains(&pixel) { + count += 1; + } + } + false => { + if !rgb_list.contains(&pixel) { + count += 1; + } + } } } } @@ -347,45 +360,110 @@ pub async fn split_images_with_filter( label_path: &String, target_height: &u32, target_width: &u32, - valid_rgb_list_str: &str, + rgb_list_str: &str, + valid_rgb_mode: bool, ) { - let valid_rgb_list: Arc>> = Arc::new(RwLock::new(Vec::new())); - for rgb in valid_rgb_list_str.split(";") { - let valid_rgb_list = Arc::clone(&valid_rgb_list); - let mut valid_rgb_list = valid_rgb_list.write().unwrap(); + let rgb_list: Arc>> = Arc::new(RwLock::new(Vec::new())); + for rgb in rgb_list_str.split(";") { + let rgb_list = Arc::clone(&rgb_list); + let mut rgb_list = rgb_list.write().unwrap(); let rgb_vec: Vec = rgb.split(',').map(|s| s.parse::().unwrap()).collect(); - valid_rgb_list.push(core::Vec3b::from([rgb_vec[0], rgb_vec[1], rgb_vec[2]])); + rgb_list.push(core::Vec3b::from([rgb_vec[0], rgb_vec[1], rgb_vec[2]])); } tracing::info!( - "Valid rgb list length: {}", - valid_rgb_list.read().unwrap().len() + "RGB list length: {} Mode: {}", + rgb_list.read().unwrap().len(), + match valid_rgb_mode { + true => "Valid", + false => "Invalid", + } ); - let image_entries = fs::read_dir(image_path).unwrap(); - let label_entries = fs::read_dir(label_path).unwrap(); - let sem = Arc::new(Semaphore::new(3)); - // let cropped_images = Arc::new(Mutex::new(HashMap::::new())); - // let cropped_labels = Arc::new(Mutex::new(HashMap::::new())); - let invalid_id = Arc::new(RwLock::new(Vec::::new())); + let image_entries: Vec; + let label_entries: Vec; + + let image_output_path: String; + let label_output_path: String; + + let image_path = PathBuf::from(image_path.as_str()); + let label_path = PathBuf::from(label_path.as_str()); + if label_path.is_dir() ^ image_path.is_dir() { + tracing::error!("Image and label path should be both directories or both files"); + return (); + } + + if image_path.is_dir() { + image_entries = fs::read_dir(&image_path) + .unwrap() + .map(|e| e.unwrap().path()) + .collect(); + + image_output_path = format!("{}\\output\\", image_path.to_str().unwrap()); + } else { + image_entries = vec![image_path.clone()]; + image_output_path = format!( + "{}\\..\\output\\images\\", + image_path.parent().unwrap().to_str().unwrap() + ); + } + + if label_path.is_dir() { + label_entries = fs::read_dir(&label_path) + .unwrap() + .map(|e| e.unwrap().path()) + .collect(); + + label_output_path = format!("{}\\output\\", label_path.to_str().unwrap()); + } else { + label_entries = vec![label_path.clone()]; + label_output_path = format!( + "{}\\..\\output\\labels\\", + label_path.parent().unwrap().to_str().unwrap() + ); + } + + let sem = Arc::new(Semaphore::new(5)); + let valid_id = Arc::new(RwLock::new(Vec::::new())); let mut threads = tokio::task::JoinSet::new(); - fs::create_dir_all(format!("{}\\output\\", image_path)).unwrap(); - fs::create_dir_all(format!("{}\\output\\", label_path)).unwrap(); + match fs::create_dir_all(&image_output_path) { + Ok(_) => { + tracing::info!("Image output directory created"); + } + Err(e) => { + if e.kind() == std::io::ErrorKind::AlreadyExists { + tracing::info!("Image output directory already exists"); + } else { + tracing::error!("Failed to create directory: {}", e); + return (); + } + } + } + match fs::create_dir_all(&label_output_path) { + Ok(_) => { + tracing::info!("Label output directory created"); + } + Err(e) => { + if e.kind() == std::io::ErrorKind::AlreadyExists { + tracing::info!("Label output directory already exists"); + } else { + tracing::error!("Failed to create directory: {}", e); + return (); + } + } + } let mut label_extension = None; for entry in label_entries { - let entry = entry.unwrap(); - - if !entry.path().is_file() { + if !entry.is_file() { continue; } if label_extension.is_none() { let extension = entry - .path() .extension() .unwrap() .to_os_string() @@ -395,42 +473,34 @@ pub async fn split_images_with_filter( } let permit = Arc::clone(&sem); - // let cropped_labels = Arc::clone(&cropped_labels); - // let cropped_images = Arc::clone(&cropped_images); - let invalid_id = Arc::clone(&invalid_id); + let valid_id = Arc::clone(&valid_id); let target_width = *target_width; let target_height = *target_height; - let valid_rgb_list = Arc::clone(&valid_rgb_list); - let label_path = label_path.clone(); + let rgb_list = Arc::clone(&rgb_list); let label_extension = label_extension.clone(); + let label_output_path = label_output_path.clone(); threads.spawn(async move { let _ = permit.acquire().await.unwrap(); - let label_id = entry - .path() - .file_stem() - .unwrap() - .to_str() - .unwrap() - .to_string(); + let label_id = entry.file_stem().unwrap().to_str().unwrap().to_string(); let img = - imgcodecs::imread(entry.path().to_str().unwrap(), imgcodecs::IMREAD_UNCHANGED) - .unwrap(); + imgcodecs::imread(entry.to_str().unwrap(), imgcodecs::IMREAD_UNCHANGED).unwrap(); - let vertical_count = img.rows() / target_height as i32; - let horizontal_count = img.cols() / target_width as i32; + let size = img.size().unwrap(); + let (width, height) = (size.width, size.height); + let y_count = height / target_height as i32; + let x_count = width / target_width as i32; // let mut labels_map = HashMap::::new(); // Crop horizontally from left - for horizontal_index in 0..horizontal_count { - for vertical_index in 0..vertical_count { - let label_id = - format!("{}_rb2lt_{}_{}", label_id, horizontal_index, vertical_index); + for x_index in 0..x_count { + for y_index in 0..y_count { + let label_id = format!("{}_lt2rb_{}_{}", label_id, x_index, y_index); let cropped = core::Mat::roi( &img, core::Rect::new( - horizontal_index * target_width as i32, - vertical_index * target_height as i32, + x_index * target_width as i32, + y_index * target_height as i32, target_width as i32, target_height as i32, ), @@ -438,12 +508,12 @@ pub async fn split_images_with_filter( .unwrap() .clone_pointee(); - let valid_rgb_list = valid_rgb_list.read().unwrap(); - if check_valid_pixel_count(&cropped, &valid_rgb_list) { + let rgb_list = rgb_list.read().unwrap(); + if check_valid_pixel_count(&cropped, &rgb_list, valid_rgb_mode) { imgcodecs::imwrite( &format!( - "{}\\output\\{}.{}", - label_path, + "{}\\{}.{}", + label_output_path, label_id, label_extension.as_ref().unwrap() ), @@ -451,41 +521,35 @@ pub async fn split_images_with_filter( &core::Vector::new(), ) .unwrap(); - } else { - invalid_id.write().unwrap().push(label_id.clone()); + valid_id.write().unwrap().push(label_id.clone()); } tracing::info!("Label {} processed", label_id); - // labels_map.insert( - // format!("{}_lt2rb_{}_{}", label_id, horizontal_index, vertical_index), - // cropped, - // ); } } tracing::info!("Label {} LTR iteration done", label_id); // Crop horizontally from right - for horizontal_index in 0..horizontal_count { - for vertical_index in 0..vertical_count { - let label_id = - format!("{}_rb2lt_{}_{}", label_id, horizontal_index, vertical_index); + for x_index in 0..x_count { + for y_index in 0..y_count { + let label_id = format!("{}_rb2lt_{}_{}", label_id, x_index, y_index); let cropped = core::Mat::roi( &img, core::Rect::new( - img.cols() - (horizontal_index + 1) * target_width as i32, - img.rows() - (vertical_index + 1) * target_height as i32, + width - (x_index + 1) * target_width as i32, + height - (y_index + 1) * target_height as i32, target_width as i32, target_height as i32, ), ) .unwrap() .clone_pointee(); - let valid_rgb_list = valid_rgb_list.read().unwrap(); - if check_valid_pixel_count(&cropped, &valid_rgb_list) { + let rgb_list = rgb_list.read().unwrap(); + if check_valid_pixel_count(&cropped, &rgb_list, valid_rgb_mode) { imgcodecs::imwrite( &format!( - "{}\\output\\{}.{}", - label_path, + "{}\\{}.{}", + label_output_path, label_id, label_extension.as_ref().unwrap() ), @@ -493,8 +557,7 @@ pub async fn split_images_with_filter( &core::Vector::new(), ) .unwrap(); - } else { - invalid_id.write().unwrap().push(label_id.clone()); + valid_id.write().unwrap().push(label_id.clone()); } tracing::info!("Label {} processed", label_id); } @@ -502,41 +565,6 @@ pub async fn split_images_with_filter( tracing::info!("Label {} RTL iteration done", label_id); - // let mut useless_img_id = Vec::::new(); - // let valid_rgb_list = valid_rgb_list.read().unwrap(); - // for (label_id, label) in labels_map.iter() { - // if check_valid_pixel_count(label, &valid_rgb_list) { - // continue; - // } - // useless_img_id.push(label_id.clone()); - // } - - // // let mut cropped_labels = cropped_labels.lock().unwrap(); - // // let mut cropped_images = cropped_images.lock().unwrap(); - // let mut invalid_id = invalid_id - // .write() - // .expect_or_log("Failed to lock invalid_id vec"); - // for img_id in useless_img_id { - // labels_map.remove(&img_id); - // // cropped_images.remove(&img_id); - // invalid_id.push(img_id); - // } - - // tracing::info!("Label {} save start", label_id); - // for (label_id, label) in labels_map.iter() { - // imgcodecs::imwrite( - // &format!( - // "{}\\output\\{}.{}", - // label_path, - // label_id, - // label_extension.as_ref().unwrap() - // ), - // label, - // &core::Vector::new(), - // ) - // .unwrap(); - // } - // cropped_labels.extend(labels_map); tracing::info!("Label {} process done", label_id); }); } @@ -547,15 +575,12 @@ pub async fn split_images_with_filter( let mut image_extension = None; for entry in image_entries { - let entry = entry.unwrap(); - - if !entry.path().is_file() { + if !entry.is_file() { continue; } if image_extension.is_none() { let extension = entry - .path() .extension() .unwrap() .to_os_string() @@ -566,56 +591,44 @@ pub async fn split_images_with_filter( let permit = Arc::clone(&sem); // let cropped_images = Arc::clone(&cropped_images); - let invalid_id = Arc::clone(&invalid_id); + let valid_id = Arc::clone(&valid_id); let target_height = *target_height; let target_width = *target_width; - let image_path = image_path.clone(); let image_extension = image_extension.clone(); - + let image_output_path = image_output_path.clone(); threads.spawn(async move { let _ = permit.acquire().await.unwrap(); - let img_id = entry - .path() - .file_stem() - .unwrap() - .to_str() - .unwrap() - .to_string(); + let img_id = entry.file_stem().unwrap().to_str().unwrap().to_string(); let img = - imgcodecs::imread(entry.path().to_str().unwrap(), imgcodecs::IMREAD_UNCHANGED) - .unwrap(); + imgcodecs::imread(entry.to_str().unwrap(), imgcodecs::IMREAD_UNCHANGED).unwrap(); - let vertical_count = img.rows() / target_height as i32; - let horizontal_count = img.cols() / target_width as i32; - // let mut imgs_map = HashMap::::new(); + let size = img.size().unwrap(); + let (width, height) = (size.width, size.height); + let y_count = height / target_height as i32; + let x_count = width / target_width as i32; // Crop horizontally from left - for horizontal_index in 0..horizontal_count { - for vertical_index in 0..vertical_count { - let img_id = - format!("{}_lt2rb_{}_{}", img_id, horizontal_index, vertical_index); - if invalid_id.read().unwrap().contains(&img_id) { + for x_index in 0..x_count { + for y_index in 0..y_count { + let img_id = format!("{}_lt2rb_{}_{}", img_id, x_index, y_index); + if !valid_id.read().unwrap().contains(&img_id) { continue; } let cropped = core::Mat::roi( &img, core::Rect::new( - horizontal_index * target_width as i32, - vertical_index * target_height as i32, + x_index * target_width as i32, + y_index * target_height as i32, target_width as i32, target_height as i32, ), ) .unwrap() .clone_pointee(); - // imgs_map.insert( - // img_id, - // cropped, - // ); imgcodecs::imwrite( &format!( - "{}\\output\\{}.{}", - image_path, + "{}\\{}.{}", + image_output_path, img_id, image_extension.as_ref().unwrap() ), @@ -630,18 +643,17 @@ pub async fn split_images_with_filter( tracing::info!("Image {} RTL iteration done", img_id); // Crop horizontally from right - for horizontal_index in 0..horizontal_count { - for vertical_index in 0..vertical_count { - let img_id = - format!("{}_rb2lt_{}_{}", img_id, horizontal_index, vertical_index); - if invalid_id.read().unwrap().contains(&img_id) { + for x_index in 0..x_count { + for y_index in 0..y_count { + let img_id = format!("{}_rb2lt_{}_{}", img_id, x_index, y_index); + if !valid_id.read().unwrap().contains(&img_id) { continue; } let cropped = core::Mat::roi( &img, core::Rect::new( - img.cols() - (horizontal_index + 1) * target_width as i32, - img.rows() - (vertical_index + 1) * target_height as i32, + width - (x_index + 1) * target_width as i32, + height - (y_index + 1) * target_height as i32, target_width as i32, target_height as i32, ), @@ -650,8 +662,8 @@ pub async fn split_images_with_filter( .clone_pointee(); imgcodecs::imwrite( &format!( - "{}\\output\\{}.{}", - image_path, + "{}\\{}.{}", + image_output_path, img_id, image_extension.as_ref().unwrap() ), @@ -671,36 +683,4 @@ pub async fn split_images_with_filter( } while threads.join_next().await.is_some() {} - - // let cropped_labels = cropped_labels.lock().unwrap(); - // for (label_id, label) in cropped_labels.iter() { - // imgcodecs::imwrite( - // &format!( - // "{}\\output\\{}.{}", - // label_path, - // label_id, - // label_extension.as_ref().unwrap() - // ), - // label, - // &core::Vector::new(), - // ) - // .unwrap(); - // tracing::info!("Label {} saved", label_id); - // } - - // let cropped_images = cropped_images.lock().unwrap(); - // for (img_id, img) in cropped_images.iter() { - // imgcodecs::imwrite( - // &format!( - // "{}\\output\\{}.{}", - // image_path, - // img_id, - // image_extension.as_ref().unwrap() - // ), - // img, - // &core::Vector::new(), - // ) - // .unwrap(); - // tracing::info!("Image {} saved", img_id); - // } } diff --git a/src/common/operation.rs b/src/common/operation.rs index 6e89b89..d57605f 100644 --- a/src/common/operation.rs +++ b/src/common/operation.rs @@ -1,7 +1,11 @@ use image::imageops::FilterType; -use opencv::{core, core::MatTraitConst, imgcodecs}; +use opencv::{ + core::{self, MatTraitConst}, + imgcodecs, +}; use std::{fs, io::Cursor, sync::Arc}; use tokio::{fs::File, io::AsyncWriteExt, sync::Semaphore, task::JoinSet}; +use tracing_unwrap::{OptionExt, ResultExt}; pub async fn resize_images( dataset_path: &String, @@ -130,9 +134,70 @@ pub fn crop_rectangle_region(source_path: &String, target_path: &String, corners let cropped_img = core::Mat::roi( &img, - core::Rect::new(cords[0].0, cords[0].1, cords[1].0 - cords[0].0, cords[1].1 - cords[0].1), + core::Rect::new( + cords[0].0, + cords[0].1, + cords[1].0 - cords[0].0, + cords[1].1 - cords[0].1, + ), ) .unwrap(); imgcodecs::imwrite(&target_path, &cropped_img, &core::Vector::new()).unwrap(); tracing::info!("Image {} done", source_path); } + +pub fn normalize(dataset_path: &String, target_max: &f64, target_min: &f64) { + let entries = fs::read_dir(dataset_path).expect_or_log("Failed to read directory"); + + fs::create_dir_all(format!("{}\\output\\", dataset_path)) + .expect_or_log("Failed to create directory"); + + for entry in entries { + let entry = entry.unwrap(); + if entry.path().is_dir() { + continue; + } + let img = imgcodecs::imread( + entry + .path() + .to_str() + .expect_or_log("Failed to get image path"), + imgcodecs::IMREAD_UNCHANGED, + ) + .expect_or_log("Failed to read image"); + + let size = img.size().expect_or_log("Failed to get image size"); + let mut dst = core::Mat::new_rows_cols_with_default( + size.height, + size.width, + img.typ(), + opencv::core::Scalar::all(0.), + ) + .expect_or_log("Failed to create dst mat"); + core::normalize( + &img, + &mut dst, + target_min.clone(), + target_max.clone(), + core::NORM_MINMAX, + -1, + &core::no_array(), + ) + .expect_or_log("Failed to normalize"); + + let dst_path = format!( + "{}\\output\\{}", + dataset_path, + entry + .path() + .file_name() + .expect_or_log("Failed to get file name") + .to_str() + .expect_or_log("Failed to convert file name to string") + .to_owned() + ); + imgcodecs::imwrite(&dst_path, &dst, &core::Vector::new()) + .expect_or_log("Failed to write image"); + tracing::info!("Image {} done", dst_path); + } +} diff --git a/src/main.rs b/src/main.rs index 4ecfcbc..3d11a53 100644 --- a/src/main.rs +++ b/src/main.rs @@ -44,6 +44,16 @@ enum CommonCommands { rectangle: String, }, + /// Normalize image to given range + Normalize { + #[arg(short, long, help = "The path for the original image")] + dataset_path: String, + #[arg(long, help = "The max value for normalization")] + max: f64, + #[arg(long, help = "The min value for normalization")] + min: f64, + }, + // TODO: Add arg for image extension selection /// Map one RGB color to another in a given PNG image file MapColor { @@ -119,14 +129,17 @@ enum CommonCommands { #[arg(short, long, help = "The path for the folder containing labels")] label_path: String, - #[arg(short, long, help = "Valid RGB list, in R0,G0,B0;R1,G1,B1 format")] - valid_rgb_list: String, + #[arg(short, long, help = "RGB list, in R0,G0,B0;R1,G1,B1 format")] + rgb_list: String, #[arg(long = "height", help = "Height for each split")] target_height: u32, #[arg(long = "width", help = "Width for each split")] target_width: u32, + + #[arg(short, help = "Use valid RGB filter mode")] + valid_rgb_mode: bool, }, /// Map 8 bit grayscale PNG class image to RGB image @@ -295,6 +308,13 @@ async fn main() { } => { common::operation::crop_rectangle_region(image_path, save_path, rectangle); } + CommonCommands::Normalize { + dataset_path, + max, + min, + } => { + common::operation::normalize(dataset_path, max, min); + } CommonCommands::MapColor { original_color, new_color, @@ -338,14 +358,16 @@ async fn main() { label_path, target_height, target_width, - valid_rgb_list, + rgb_list, + valid_rgb_mode } => { common::augment::split_images_with_filter( image_path, label_path, target_height, target_width, - valid_rgb_list, + rgb_list, + *valid_rgb_mode ) .await; }