Refactors image loading and processing

Improves image loading error handling in distributed processing.

Streamlines image processing by removing unnecessary conversions
and using a more direct method for creating the tensor from the image data.

Improves code readability by introducing formatting and consistent coding styles.

Author: nusu-github

Cargo.lock

@@ -38,8 +38,7 @@ image = { version = "0.25", default-features = false }
 imageops-kit = { git = "https://github.com/nusu-github/imageops-kit", version = "0.1.0" }
 imageproc = { version = "0.25", default-features = false }
 indicatif = "0.18"
-ndarray = "0.16"
-nshare = { version = "0.10", default-features = false, features = ["image", "ndarray"] }
+ndarray = "0.16.1"
 ort = "2.0.0-rc.10"
 parking_lot = { version = "0.12", features = ["hardware-lock-elision"] }
 rayon = "1.11"

Cargo.toml

@@ -284,11 +284,12 @@ impl<M: BatchImageSegmentationModel + 'static> BatchProcessor for GpuInferenceBa
                 batch_for_load
                     .par_iter()
                     .map(|job| {
-                        image::open(&job.payload.input_path)
-                            .map_err(|e| (job.id.clone(), e))
+                        image::open(&job.payload.input_path).map_err(|e| (job.id.clone(), e))
                     })
                     .collect()
-            }).await.unwrap();
+            })
+            .await
+            .unwrap();
 
             let mut images = Vec::with_capacity(batch_size);
             let mut error_indices = Vec::new();
@@ -305,10 +306,10 @@ impl<M: BatchImageSegmentationModel + 'static> BatchProcessor for GpuInferenceBa
 
             for idx in error_indices.iter().rev() {
                 let mut failed_job = batch.remove(*idx);
-                failed_job.payload.metadata.insert(
-                    "error".to_string(),
-                    "Image loading error".to_string(),
-                );
+                failed_job
+                    .payload
+                    .metadata
+                    .insert("error".to_string(), "Image loading error".to_string());
                 failed_job.job_type = JobType::Postprocessing;
             }
 
@@ -345,10 +346,9 @@ impl<M: BatchImageSegmentationModel + 'static> BatchProcessor for GpuInferenceBa
                             return;
                         }
 
-                        job.payload.metadata.insert(
-                            "segmentation_complete".to_string(),
-                            "true".to_string(),
-                        );
+                        job.payload
+                            .metadata
+                            .insert("segmentation_complete".to_string(), "true".to_string());
                         job.payload
                             .metadata
                             .insert("batch_size".to_string(), batch_size.to_string());
@@ -359,7 +359,9 @@ impl<M: BatchImageSegmentationModel + 'static> BatchProcessor for GpuInferenceBa
                         job.job_type = JobType::Postprocessing;
                     });
                 batch
-            }).await.unwrap();
+            })
+            .await
+            .unwrap();
 
             Ok(batch)
         })

src/distributed.rs

@@ -11,7 +11,6 @@ use image::{
 use imageops_kit::{AlphaMaskError, Image, ModifyAlphaExt, NormalizedFrom, PaddingExt};
 use imageproc::map::map_colors;
 use ndarray::prelude::*;
-use nshare::AsNdarray3;
 use ort::value::TensorRef;
 use ort::{
     execution_providers::{
@@ -102,10 +101,7 @@ impl Model {
             })?
             .with_optimized_model_path(cache_dir.join(format!(
                 "{}.optimized.onnx",
-                model_path
-                    .file_stem()
-                    .unwrap_or_default()
-                    .to_string_lossy()
+                model_path.file_stem().unwrap_or_default().to_string_lossy()
             )))
             .map_err(|e| AnimeSegError::Model {
                 operation: "optimized model path configuration".to_string(),
@@ -215,7 +211,12 @@ impl crate::traits::BatchImageSegmentationModel for Model {
             dimensions.push(img.dimensions());
         }
 
-        let batch_shape = (batch_size, 3, self.image_size as usize, self.image_size as usize);
+        let batch_shape = (
+            batch_size,
+            3,
+            self.image_size as usize,
+            self.image_size as usize,
+        );
         let mut batch_tensor = Array4::<f32>::zeros(batch_shape);
 
         for (i, tensor) in batch_tensors.iter().enumerate() {
@@ -260,15 +261,27 @@ where
     let image = imageops::resize(image, image_size, image_size, FilterType::Lanczos3);
     let (w, h) = image.dimensions();
     let zero = S::zero();
-    let (image, (x, y)) = image
-        .to_square(Rgb([zero, zero, zero]))
+    let (image, (x, y)) =
+        image
+            .to_square(Rgb([zero, zero, zero]))
+            .map_err(|e| AnimeSegError::ImageProcessing {
+                path: "unknown".to_string(),
+                operation: "パディング追加".to_string(),
+                source: Box::new(e),
+            })?;
+
+    let (width, height) = image.dimensions();
+    let channels = 3;
+    let raw = image.into_raw();
+
+    let array_hwc = Array3::from_shape_vec((height as usize, width as usize, channels), raw)
         .map_err(|e| AnimeSegError::ImageProcessing {
             path: "unknown".to_string(),
-            operation: "パディング追加".to_string(),
+            operation: "Converting an image to an ndarray".to_string(),
             source: Box::new(e),
         })?;
+    let tensor = array_hwc.permuted_axes([2, 0, 1]).insert_axis(Axis(0));
 
-    let tensor = image.as_ndarray3().slice_move(s![NewAxis, ..;-1, .., ..]);
     let max = S::DEFAULT_MAX_VALUE.into();
     let tensor = if max == (<f32 as Primitive>::DEFAULT_MAX_VALUE) {
         tensor.map(|v| (*v).into())

src/model.rs

@@ -119,9 +119,7 @@ impl QueueProvider for InMemoryQueueProvider {
         let notifier = self.get_or_create_notifier(queue_name);
         {
             let mut queues = self.queues.lock();
-            let queue = queues
-                .entry(queue_name.to_string())
-                .or_default();
+            let queue = queues.entry(queue_name.to_string()).or_default();
             queue.push_back(job);
         }
         notifier.notify_one();
@@ -133,9 +131,7 @@ impl QueueProvider for InMemoryQueueProvider {
         loop {
             {
                 let mut queues = self.queues.lock();
-                let queue = queues
-                    .entry(queue_name.to_string())
-                    .or_default();
+                let queue = queues.entry(queue_name.to_string()).or_default();
 
                 if queue.len() < max_size {
                     queue.push_back(job);
@@ -151,9 +147,7 @@ impl QueueProvider for InMemoryQueueProvider {
 
     async fn dequeue(&self, queue_name: &str) -> Result<Option<Job>> {
         let mut queues = self.queues.lock();
-        let queue = queues
-            .entry(queue_name.to_string())
-            .or_default();
+        let queue = queues.entry(queue_name.to_string()).or_default();
         Ok(queue.pop_front())
     }