Modernizes codebase with dependency updates and API improvements

Updates burn and related dependencies to latest versions for better performance and features. Refactors image utilities from struct-based to free functions for cleaner API. Adds tracing-based logging, MiMalloc allocator for memory efficiency, and workspace-wide linting rules. Enhances documentation with inline re-exports and improves error handling in inference and training modules.

Author: nusu-github

Cargo.lock

@@ -36,6 +36,8 @@ cfg-if = "1.0.4"
 libm = "0.2.15"
 rand = "0.9.2"
 stacker = "0.1.22"
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 walkdir = "2.5.0"
 
 # Testing and benchmarking
@@ -44,5 +46,43 @@ rstest = "0.26.1"
 # Model hub integration
 hf-hub = { version = "0.4.3", default-features = false, features = ["ureq", "rustls-tls"] }
 
+[workspace.lints.rust]
+ambiguous_negative_literals = "warn"
+missing_debug_implementations = "warn"
+redundant_imports = "warn"
+redundant_lifetimes = "warn"
+trivial_numeric_casts = "warn"
+unsafe_op_in_unsafe_fn = "warn"
+unused_lifetimes = "warn"
+
+[workspace.lints.clippy]
+cargo = { level = "warn", priority = -1 }
+complexity = { level = "warn", priority = -1 }
+correctness = { level = "warn", priority = -1 }
+pedantic = { level = "warn", priority = -1 }
+perf = { level = "warn", priority = -1 }
+style = { level = "warn", priority = -1 }
+suspicious = { level = "warn", priority = -1 }
+allow_attributes_without_reason = "warn"
+assertions_on_result_states = "warn"
+clone_on_ref_ptr = "warn"
+deref_by_slicing = "warn"
+empty_drop = "warn"
+empty_enum_variants_with_brackets = "warn"
+empty_structs_with_brackets = "warn"
+fn_to_numeric_cast_any = "warn"
+if_then_some_else_none = "warn"
+map_err_ignore = "warn"
+redundant_type_annotations = "warn"
+renamed_function_params = "warn"
+semicolon_outside_block = "warn"
+string_to_string = "warn"
+undocumented_unsafe_blocks = "warn"
+unnecessary_safety_comment = "warn"
+unnecessary_safety_doc = "warn"
+unneeded_field_pattern = "warn"
+unused_result_ok = "warn"
+literal_string_with_formatting_args = "allow"
+
 [profile.dev.package.image]
 opt-level = 3

Cargo.toml

@@ -5,6 +5,7 @@
 
 mod backbones;
 
+#[doc(inline)]
 pub use backbones::{
     pvt_v2::{PvtV2Config, PyramidVisionTransformerImpr},
     resnet::{ResNetBackbone, ResNetConfig},
@@ -276,4 +277,4 @@ mod tests {
     use burn::backend::Cpu;
 
     pub type TestBackend = Cpu;
-}
+}

crates/birefnet-backbones/src/lib.rs

@@ -9,14 +9,18 @@ mod identity;
 mod trunc_normal;
 
 // Convenient re-exports
+#[doc(inline)]
 pub use drop_path::{DropPath, DropPathConfig};
+#[doc(inline)]
 pub use erfinv::{Erfinv, erfinv};
+#[doc(inline)]
 pub use identity::Identity;
+#[doc(inline)]
 pub use trunc_normal::{trunc_normal, trunc_normal_};
 
 #[cfg(test)]
 mod tests {
     use burn::backend::Cpu;
 
     pub type TestBackend = Cpu;
-}
+}

crates/birefnet-extra-ops/src/lib.rs

@@ -1,10 +1,14 @@
 pub mod postprocessing;
 
-pub use postprocessing::*;
+#[doc(inline)]
+pub use postprocessing::{
+    apply_threshold, fill_holes, gaussian_blur, morphological_closing, morphological_opening,
+    postprocess_mask, remove_small_components, resize_tensor, tensor_to_image_data,
+};
 
 #[cfg(test)]
 mod tests {
     use burn::backend::Cpu;
 
     pub type TestBackend = Cpu;
-}
+}

crates/birefnet-inference/src/lib.rs

@@ -4,7 +4,7 @@
 //! segmentation masks and images.
 
 use anyhow::{Context, Result};
-use birefnet_util::ImageUtils;
+use birefnet_util::{dynamic_image_to_tensor, tensor_to_dynamic_image};
 use burn::tensor::{Tensor, backend::Backend};
 use image::{self, imageops::FilterType};
 
@@ -209,7 +209,7 @@ pub fn resize_tensor<B: Backend>(
 
     // This approach converts tensor to image, resizes, then converts back to tensor.
     // It's inefficient but works without a direct tensor interpolation implementation.
-    let dynamic_image = ImageUtils::tensor_to_dynamic_image(tensor, false)
+    let dynamic_image = tensor_to_dynamic_image(tensor, false)
         .context("Failed to convert tensor to image for resizing")?;
 
     let resized_image = dynamic_image.resize_exact(
@@ -219,6 +219,6 @@ pub fn resize_tensor<B: Backend>(
     );
 
     // Convert back to tensor
-    ImageUtils::dynamic_image_to_tensor(resized_image, device)
+    dynamic_image_to_tensor(resized_image, device)
         .context("Failed to convert resized image back to tensor")
 }

crates/birefnet-inference/src/postprocessing.rs

@@ -78,15 +78,25 @@ mod threshold_regularization;
 mod birefnet_loss;
 
 // Re-export core loss functions and configurations
+#[doc(inline)]
 pub use birefnet_loss::{BiRefNetLoss, BiRefNetLossConfig};
+#[doc(inline)]
 pub use classification::{ClassificationLoss, ClassificationLossConfig};
+#[doc(inline)]
 pub use contour::{ContourLoss, ContourLossConfig};
+#[doc(inline)]
 pub use iou::{IoULoss, IoULossConfig};
+#[doc(inline)]
 pub use mae::{MaeLoss, MaeLossConfig};
+#[doc(inline)]
 pub use patch_iou::{PatchIoULoss, PatchIoULossConfig};
+#[doc(inline)]
 pub use pixel::{LossWeightsConfig, PixLoss, PixLossConfig};
+#[doc(inline)]
 pub use ssim::{SSIMLoss, SSIMLossConfig};
+#[doc(inline)]
 pub use structure::{StructureLoss, StructureLossConfig};
+#[doc(inline)]
 pub use threshold_regularization::{
     ThresholdRegularizationLoss, ThresholdRegularizationLossConfig,
 };
@@ -96,4 +106,4 @@ mod tests {
     use burn::backend::Cpu;
 
     pub type TestBackend = Cpu;
-}
+}

crates/birefnet-loss/src/lib.rs

@@ -78,21 +78,33 @@ pub mod utils;
 pub mod weighted_f_measure;
 
 // Re-export main types and traits
+#[doc(inline)]
 pub use biou::{BIoUMetric, BIoUMetricConfig};
+#[doc(inline)]
 pub use f_measure::{FMeasureMetric, FMeasureMetricConfig};
-pub use input::*;
+#[doc(inline)]
+pub use input::{
+    BIoUInput, BiRefNetLossInput, EMeasureInput, FMeasureInput, MAEInput, MSEInput, SMeasureInput,
+    WeightedFMeasureInput,
+};
+#[doc(inline)]
 pub use loss::LossMetric;
+#[doc(inline)]
 pub use mae::{MAEMetric, MAEMetricConfig};
+#[doc(inline)]
 pub use mse::{MSEMetric, MSEMetricConfig};
 // Re-export lesser-used items with warning
 #[deprecated(note = "S-measure is not yet fully implemented and tested")]
-pub use s_measure::*;
-pub use utils::*;
-pub use weighted_f_measure::{WeightedFMeasureMetric, WeightedFMeasureMetricConfig};
+#[doc(inline)]
+pub use s_measure::{SMeasureInput as SMeasureInputDeprecated, SMeasureMetric, SMeasureMetricConfig, calculate_s_measure};
+#[doc(inline)]
+pub use utils::{AllMetricsResult, calculate_all_metrics};
+#[doc(inline)]
+pub use weighted_f_measure::{WeightedFMeasureMetric, WeightedFMeasureMetricConfig, calculate_weighted_f_measure};
 
 #[cfg(test)]
 mod tests {
     use burn::backend::Cpu;
 
     pub type TestBackend = Cpu;
-}
+}

crates/birefnet-metric/src/lib.rs

@@ -31,7 +31,7 @@ pub struct AllMetricsResult {
 pub fn calculate_all_metrics<B: Backend>(
     predictions: Tensor<B, 4>,
     targets: Tensor<B, 4>,
-    threshold: f64,
+    _threshold: f64,
 ) -> AllMetricsResult {
     // Ensure predictions and targets have correct shape
     let [batch_size, _, _, _] = predictions.dims();

crates/birefnet-metric/src/utils.rs

@@ -28,14 +28,18 @@ mod models;
 pub mod training;
 
 // Re-export specific configuration types instead of wildcard
+#[doc(inline)]
 pub use config::{
     Backbone, BackboneConfig, DecBlk, DecChannelsInter, DecoderAttention, DecoderConfig,
     InterpolationStrategy, LateralBlock, ModelConfig, MultiScaleInput, Optimizer, PathConfig,
     PreprocMethods, PromptForLocation, Refine, RefineConfig, SqueezeBlock, Task, TaskConfig,
 };
+#[doc(inline)]
 pub use error::{BiRefNetError, BiRefNetResult};
+#[doc(inline)]
 pub use models::birefnet::{BiRefNet, BiRefNetConfig, BiRefNetRecord};
 #[cfg(feature = "train")]
+#[doc(inline)]
 pub use training::{BiRefNetBatch, BiRefNetOutput};
 
 #[cfg(test)]
@@ -45,4 +49,4 @@ mod tests {
     pub type TestBackend = Cpu;
 
     pub type TestAutodiffBackend = Autodiff<TestBackend>;
-}
+}