ultralytics_inference/

preprocessing.rs

// Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

//! Image preprocessing for YOLO inference.
//!
//! This module handles all image preprocessing operations needed before
//! running YOLO model inference, including resizing, padding, and normalization.

#![allow(
    unsafe_code,
    clippy::similar_names,
    clippy::cast_precision_loss,
    clippy::cast_possible_wrap,
    clippy::cast_sign_loss,
    clippy::cast_possible_truncation,
    clippy::wildcard_imports,
    clippy::ptr_as_ptr,
    clippy::cast_lossless,
    clippy::single_match_else,
    clippy::suboptimal_flops,
    clippy::manual_div_ceil
)]

use std::cell::RefCell;
use std::num::NonZeroUsize;
use std::sync::Arc;

use half::f16;
use image::{DynamicImage, GenericImageView, RgbImage};
use lru::LruCache;
use ndarray::{Array3, Array4};

use crate::inference::Quantization;

/// Converts current and legacy preprocessing precision arguments to `quantize`.
#[doc(hidden)]
pub trait IntoQuantization {
    /// Convert to the unified precision representation.
    fn into_quantization(self) -> Option<Quantization>;
}

impl IntoQuantization for Option<Quantization> {
    fn into_quantization(self) -> Option<Quantization> {
        self
    }
}

impl IntoQuantization for bool {
    fn into_quantization(self) -> Option<Quantization> {
        crate::inference::handle_deprecated_precision(None, Some(self))
    }
}

/// Default letterbox padding color (gray).
pub const LETTERBOX_COLOR: [u8; 3] = [114, 114, 114];

/// Fixed-point scale factor for bilinear interpolation (2^11 = 2048).
/// Matches `OpenCV`'s `INTER_RESIZE_COEF_BITS = 11` for `INTER_LINEAR`.
const SCALE_BITS: i32 = 11;
const SCALE_INT: i32 = 1 << SCALE_BITS;

/// Double scale bits for single-pass bilinear interpolation.
const SCALE_BITS_2X: i32 = 2 * SCALE_BITS;

/// Rounding bias for fixed-point bilinear, added before the final right-shift
/// to achieve round-to-nearest behavior matching `OpenCV`'s `saturate_cast`.
const ROUND_BIAS: i32 = 1 << (SCALE_BITS_2X - 1);

/// Normalized letterbox padding color (114/255 ≈ 0.447).
const LETTERBOX_NORM: f32 = 114.0 / 255.0;

/// Reciprocal of 255 for normalization.
const INV_255: f32 = 1.0 / 255.0;

/// Maximum LRU cache size for X coordinate LUTs.
const LUT_CACHE_SIZE: usize = 8;

/// X LUT entry: (`x0_byte_offset`, `x1_byte_offset`, 1-fx, fx) using 11-bit fixed-point
/// weights matching `OpenCV`'s `INTER_LINEAR` coordinate mapping.
type XLutEntry = (usize, usize, i32, i32);
type XLutKey = (u32, u32);

thread_local! {
    static X_LUT_CACHE: RefCell<LruCache<XLutKey, Arc<Vec<XLutEntry>>>> =
        RefCell::new(LruCache::new(NonZeroUsize::new(LUT_CACHE_SIZE).unwrap()));
}

/// Result of preprocessing an image, containing the tensor and transform info.
#[derive(Debug, Clone)]
pub struct PreprocessResult {
    /// Preprocessed image tensor in NCHW format, normalized to [0, 1].
    pub tensor: Array4<f32>,
    /// Preprocessed FP16 tensor (if requested).
    pub tensor_f16: Option<Array4<f16>>,
    /// Original image dimensions (height, width).
    pub orig_shape: (u32, u32),
    /// Scale factors applied (`scale_y`, `scale_x`).
    pub scale: (f32, f32),
    /// Padding applied (`pad_top`, `pad_left`).
    pub padding: (f32, f32),
}

/// Geometry parameters resolved by letterbox layout.
#[derive(Clone, Copy)]
pub(crate) struct LetterboxGeometry {
    pub(crate) new_w: u32,
    pub(crate) new_h: u32,
    pub(crate) pad_left: u32,
    pub(crate) pad_top: u32,
}

impl LetterboxGeometry {
    /// Letterbox geometry for fitting `(orig_w, orig_h)` into `target` = `(height, width)`:
    /// a uniform scale (`min` over both axes), rounded resized extents, and centered
    /// padding. The returned `scale` is that uniform gain, used for coordinate
    /// back-projection.
    ///
    /// The single source of this math for both the CPU letterbox and the `cuda-preprocess`
    /// kernel, so the two can never disagree on where a pixel lands.
    #[allow(
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss
    )]
    pub(crate) fn compute(
        orig_w: u32,
        orig_h: u32,
        target: (usize, usize),
        stretch: bool,
    ) -> (Self, (f32, f32)) {
        // Scale-fill: the source covers the target exactly, so each axis keeps its own gain
        // and there is no padding. A zero-extent axis has no gain to report and yields 1.0,
        // keeping back-projection finite instead of dividing coordinates by infinity.
        if stretch {
            let gain = |extent: usize, orig: u32| {
                if orig == 0 {
                    1.0
                } else {
                    extent as f32 / orig as f32
                }
            };
            let geom = Self {
                new_w: target.1 as u32,
                new_h: target.0 as u32,
                pad_left: 0,
                pad_top: 0,
            };
            return (geom, (gain(target.0, orig_h), gain(target.1, orig_w)));
        }
        let (target_h, target_w) = (target.0 as f32, target.1 as f32);
        let (orig_hf, orig_wf) = (orig_h as f32, orig_w as f32);
        let scale = (target_h / orig_hf).min(target_w / orig_wf);
        // Rounding can collapse an extreme aspect ratio to zero: a 10000x1 source at imgsz 640
        // scales its height to 0.064, which would letterbox the entire tensor and throw the
        // image away. Keep at least one pixel per axis, leaving a genuinely empty source at 0.
        let extent = |orig_extent: f32, orig: u32| {
            if orig == 0 {
                0
            } else {
                ((orig_extent * scale).round() as u32).max(1)
            }
        };
        let new_w = extent(orig_wf, orig_w);
        let new_h = extent(orig_hf, orig_h);
        let pad_left = (target.1 as u32).saturating_sub(new_w) / 2;
        let pad_top = (target.0 as u32).saturating_sub(new_h) / 2;
        (
            Self {
                new_w,
                new_h,
                pad_left,
                pad_top,
            },
            (scale, scale),
        )
    }
}

/// Build a `PreprocessResult` from a resolved letterbox geometry.
///
/// Runs the fused zero-copy resize/pad/normalize, optionally produces an FP16 tensor,
/// and packages the transform metadata.
#[allow(clippy::cast_precision_loss)]
fn build_preprocess_result(
    image: &DynamicImage,
    target_size: (usize, usize),
    geom: LetterboxGeometry,
    scale: (f32, f32),
    orig_shape: (u32, u32),
    fp16: bool,
) -> PreprocessResult {
    let (orig_width, orig_height) = image.dimensions();

    let tensor = match image {
        DynamicImage::ImageRgb8(rgb) => {
            fused_zerocopy_preprocess(rgb.as_raw(), orig_width, orig_height, target_size, &geom)
        }
        _ => {
            let src_rgb = image.to_rgb8();
            fused_zerocopy_preprocess(
                src_rgb.as_raw(),
                orig_width,
                orig_height,
                target_size,
                &geom,
            )
        }
    };

    let tensor_f16 = if fp16 {
        Some(tensor_f32_to_f16(&tensor))
    } else {
        None
    };

    PreprocessResult {
        tensor,
        tensor_f16,
        orig_shape,
        scale,
        padding: (geom.pad_top as f32, geom.pad_left as f32),
    }
}

/// Preprocess an image for YOLO inference.
///
/// Performs letterbox resizing, normalization to [0, 1], and conversion to
/// NCHW tensor format.
///
/// # Arguments
///
/// * `image` - Input image.
/// * `target_size` - Target size as (height, width).
/// * `stride` - Model stride for padding alignment (typically 32).
///
/// # Returns
///
/// Preprocessed tensor and transform information for post-processing.
#[must_use]
pub fn preprocess_image(
    image: &DynamicImage,
    target_size: (usize, usize),
    stride: u32,
) -> PreprocessResult {
    preprocess_image_with_precision(image, target_size, stride, None)
}

/// Preprocess an image for YOLO inference with optional FP16 output.
///
/// # Arguments
///
/// * `image` - Input image.
/// * `target_size` - Target size as (height, width).
/// * `stride` - Model stride for padding alignment (typically 32).
/// * `quantize` - Requested precision. FP16 also generates an FP16 tensor.
///
/// # Returns
///
/// Preprocessed tensor and transform information for post-processing.
#[must_use]
pub fn preprocess_image_with_precision(
    image: &DynamicImage,
    target_size: (usize, usize),
    stride: u32,
    quantize: impl IntoQuantization,
) -> PreprocessResult {
    let quantize = quantize.into_quantization();
    let (orig_width, orig_height) = image.dimensions();
    let orig_shape = (orig_height, orig_width);

    let (geom, scale) = calculate_letterbox_params(orig_width, orig_height, target_size, stride);
    build_preprocess_result(
        image,
        target_size,
        geom,
        scale,
        orig_shape,
        quantize == Some(Quantization::Fp16),
    )
}

/// Preprocess an image for RT-DETR inference.
///
/// Stretches the image onto `target_size` so it fills the tensor exactly, with no
/// aspect-ratio preservation and no padding. RT-DETR expects this layout; feeding it a
/// padded letterbox shifts every prediction.
///
/// # Arguments
///
/// * `image` - Input image.
/// * `target_size` - Target size as (height, width).
/// * `quantize` - Requested precision. FP16 also generates an FP16 tensor.
///
/// # Returns
///
/// Preprocessed tensor and transform information for post-processing.
#[must_use]
pub fn preprocess_image_stretch(
    image: &DynamicImage,
    target_size: (usize, usize),
    quantize: impl IntoQuantization,
) -> PreprocessResult {
    let quantize = quantize.into_quantization();
    let (orig_width, orig_height) = image.dimensions();
    let (geom, scale) = LetterboxGeometry::compute(orig_width, orig_height, target_size, true);
    build_preprocess_result(
        image,
        target_size,
        geom,
        scale,
        (orig_height, orig_width),
        quantize == Some(Quantization::Fp16),
    )
}

/// Get or compute the X coordinate LUT for bilinear interpolation.
///
/// Uses 11-bit fixed-point weights matching `OpenCV`'s `INTER_LINEAR` coordinate mapping:
/// `src_x = (dst_x + 0.5) * (src_w / dst_w) - 0.5`
///
/// Weight computation matches `OpenCV`'s `resize.cpp`:
/// `cbuf[0] = saturate_cast<short>((1-fx) * 2048); cbuf[1] = 2048 - cbuf[0];`
fn get_or_compute_x_lut(src_w: u32, dst_w: u32) -> Arc<Vec<XLutEntry>> {
    let key = (src_w, dst_w);

    X_LUT_CACHE.with(|cache| {
        let mut cache = cache.borrow_mut();

        if let Some(lut) = cache.get(&key) {
            return Arc::clone(lut);
        }

        let scale_x = src_w as f32 / dst_w as f32;
        let src_w_max = (src_w - 1) as i32;

        let lut: Arc<Vec<XLutEntry>> = Arc::new(
            (0..dst_w)
                .map(|dx| {
                    let sx = ((dx as f32 + 0.5) * scale_x - 0.5).max(0.0);
                    let x0 = sx.floor() as i32;
                    // Match OpenCV: cbuf[0] = saturate_cast<short>((1-fx)*SCALE),
                    //               cbuf[1] = SCALE - cbuf[0]
                    let fx_f = sx - x0 as f32;
                    let fx_inv = ((1.0 - fx_f) * SCALE_INT as f32 + 0.5) as i32;
                    let fx = SCALE_INT - fx_inv;
                    let x0c = x0.clamp(0, src_w_max) as usize * 3;
                    let x1c = (x0 + 1).clamp(0, src_w_max) as usize * 3;
                    (x0c, x1c, fx_inv, fx)
                })
                .collect(),
        );

        cache.put(key, Arc::clone(&lut));
        lut
    })
}

/// Zero-copy fused preprocessing for maximum performance.
///
/// Combines bilinear resize, letterbox padding, and NCHW normalization
/// in a single memory pass with parallel row processing.
fn fused_zerocopy_preprocess(
    src_raw: &[u8],
    src_w: u32,
    src_h: u32,
    target_size: (usize, usize),
    geom: &LetterboxGeometry,
) -> Array4<f32> {
    #[allow(clippy::wildcard_imports)] // native: rayon prelude; wasm: sequential shims
    use crate::parallel::*;
    use std::mem::MaybeUninit;
    use std::sync::atomic::{AtomicPtr, Ordering};

    let LetterboxGeometry {
        new_w: new_width,
        new_h: new_height,
        pad_left,
        pad_top,
    } = *geom;

    let (dst_h, dst_w) = target_size;

    // A zero-extent source or image region has no pixels to sample. Every row would take the
    // padding branch below anyway, but the LUT and scale setup would first underflow on
    // `src_w - 1` / `src_h - 1` and divide by zero, so return the letterbox fill directly.
    if src_w == 0 || src_h == 0 || new_width == 0 || new_height == 0 {
        return Array4::from_elem((1, 3, dst_h, dst_w), LETTERBOX_NORM);
    }

    let channel_size = dst_h * dst_w;
    let src_stride = (src_w * 3) as usize;

    // ALLOCATE UNINITIALIZED: Saves ~0.2ms by not zeroing memory
    let mut tensor: Array4<MaybeUninit<f32>> = Array4::uninit((1, 3, dst_h, dst_w));
    let out_ptr = tensor.as_mut_ptr() as *mut f32;

    // Use AtomicPtr for thread-safe pointer sharing (each thread writes to disjoint rows)
    let atomic_ptr = AtomicPtr::new(out_ptr);

    let x_lut = get_or_compute_x_lut(src_w, new_width);
    let scale_y = src_h as f32 / new_height as f32;
    let src_h_max = (src_h - 1) as i32;

    let pad_top_usize = pad_top as usize;
    let pad_left_usize = pad_left as usize;
    let new_height_usize = new_height as usize;
    let new_width_usize = new_width as usize;

    // Parallel row processing with raw pointers (no bounds checks)
    (0..dst_h).into_par_iter().for_each(|dy| {
        let data_ptr = atomic_ptr.load(Ordering::Relaxed);
        unsafe {
            // Calculate row pointers for R, G, B channels

            let r_row = data_ptr.add(dy * dst_w);
            let g_row = data_ptr.add(channel_size + dy * dst_w);
            let b_row = data_ptr.add(2 * channel_size + dy * dst_w);

            // Vertical padding (top/bottom rows)
            if dy < pad_top_usize || dy >= pad_top_usize + new_height_usize {
                for dx in 0..dst_w {
                    *r_row.add(dx) = LETTERBOX_NORM;
                    *g_row.add(dx) = LETTERBOX_NORM;
                    *b_row.add(dx) = LETTERBOX_NORM;
                }
                return;
            }

            // Image row calculations - 11-bit fixed-point bilinear matching
            // OpenCV's INTER_LINEAR (INTER_RESIZE_COEF_BITS = 11).
            let img_dy = dy - pad_top_usize;
            let sy = ((img_dy as f32 + 0.5) * scale_y - 0.5).max(0.0);
            let y0 = sy.floor() as i32;
            let fy_f = sy - y0 as f32;
            let fy_inv = ((1.0 - fy_f) * SCALE_INT as f32 + 0.5) as i32;
            let fy = SCALE_INT - fy_inv;

            let y0c = y0.clamp(0, src_h_max) as usize;
            let y1c = (y0 + 1).clamp(0, src_h_max) as usize;
            let row0_off = y0c * src_stride;
            let row1_off = y1c * src_stride;

            // Left padding
            for dx in 0..pad_left_usize {
                *r_row.add(dx) = LETTERBOX_NORM;
                *g_row.add(dx) = LETTERBOX_NORM;
                *b_row.add(dx) = LETTERBOX_NORM;
            }

            // Inner image pixels - fixed-point bilinear with rounding.
            // Uses untruncated weights (w = fx * fy, range [0, 2048^2]) and a
            // single shift with rounding bias, matching OpenCV's saturate_cast:
            //   result = (sum + ROUND_BIAS) >> 22
            // Max intermediate: 255 * 2048^2 + 2^21 ≈ 1.07B < i32::MAX.
            let mut img_dx = 0usize;
            let src_ptr = src_raw.as_ptr();

            while img_dx < new_width_usize {
                let (x0_off, x1_off, fx_inv, fx) = *x_lut.get_unchecked(img_dx);
                let w00 = fx_inv * fy_inv;
                let w10 = fx * fy_inv;
                let w01 = fx_inv * fy;
                let w11 = fx * fy;

                let p00 = src_ptr.add(row0_off + x0_off);
                let p10 = src_ptr.add(row0_off + x1_off);
                let p01 = src_ptr.add(row1_off + x0_off);
                let p11 = src_ptr.add(row1_off + x1_off);

                let out_x = pad_left_usize + img_dx;
                *r_row.add(out_x) = ((*p00 as i32 * w00
                    + *p10 as i32 * w10
                    + *p01 as i32 * w01
                    + *p11 as i32 * w11
                    + ROUND_BIAS)
                    >> SCALE_BITS_2X) as f32
                    * INV_255;
                *g_row.add(out_x) = ((*p00.add(1) as i32 * w00
                    + *p10.add(1) as i32 * w10
                    + *p01.add(1) as i32 * w01
                    + *p11.add(1) as i32 * w11
                    + ROUND_BIAS)
                    >> SCALE_BITS_2X) as f32
                    * INV_255;
                *b_row.add(out_x) = ((*p00.add(2) as i32 * w00
                    + *p10.add(2) as i32 * w10
                    + *p01.add(2) as i32 * w01
                    + *p11.add(2) as i32 * w11
                    + ROUND_BIAS)
                    >> SCALE_BITS_2X) as f32
                    * INV_255;

                img_dx += 1;
            }

            // Right padding
            for dx in (pad_left_usize + new_width_usize)..dst_w {
                *r_row.add(dx) = LETTERBOX_NORM;
                *g_row.add(dx) = LETTERBOX_NORM;
                *b_row.add(dx) = LETTERBOX_NORM;
            }
        }
    });

    // SAFETY: All elements have been initialized
    unsafe { tensor.assume_init() }
}

/// Convert f32 tensor to f16 tensor.
fn tensor_f32_to_f16(tensor: &Array4<f32>) -> Array4<half::f16> {
    use half::slice::HalfFloatSliceExt;
    // `mapv` converts one element at a time; `half`'s slice conversion is vectorized where
    // the target supports it. Falls back to `mapv` if the tensor is not contiguous.
    let Some(src) = tensor.as_slice() else {
        return tensor.mapv(half::f16::from_f32);
    };
    let mut out = vec![half::f16::ZERO; src.len()];
    out.convert_from_f32_slice(src);
    Array4::from_shape_vec(tensor.raw_dim(), out).expect("shape matches the source tensor")
}

/// Calculate target size for rectangular inference mode.
///
/// Adjusts `target_size` such that the image's aspect ratio is preserved,
/// and both dimensions are multiples of `stride`.
///
/// # Arguments
///
/// * `orig_width` - Original image width.
/// * `orig_height` - Original image height.
/// * `target_size` - Base target size (e.g. 640x640).
/// * `stride` - Model stride for alignment.
///
/// # Returns
///
/// Adjusted target size as (height, width).
#[must_use]
pub fn calculate_rect_size(
    orig_width: u32,
    orig_height: u32,
    target_size: (usize, usize),
    stride: u32,
) -> (usize, usize) {
    let (target_h, target_w) = target_size;

    #[allow(clippy::cast_precision_loss)]
    let orig_h = orig_height as f32;
    #[allow(clippy::cast_precision_loss)]
    let orig_w = orig_width as f32;
    #[allow(clippy::cast_precision_loss)]
    let target_h_f = target_h as f32;
    #[allow(clippy::cast_precision_loss)]
    let target_w_f = target_w as f32;

    // Calculate scale to fit within target while maintaining aspect ratio
    let scale = (target_h_f / orig_h).min(target_w_f / orig_w);

    // New dimensions after scaling
    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    let new_h = (orig_h * scale).round() as usize;
    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    let new_w = (orig_w * scale).round() as usize;

    // Round up to nearest multiple of stride. A degenerate source rounds an extent to zero
    // (a zero-width image, or an aspect ratio past what the target can represent), which would
    // hand the model a zero-width input and the CUDA preprocess a zero-sized kernel grid, so
    // keep at least one stride step on each axis.
    let stride = stride as usize;
    let rect_h = (((new_h + stride - 1) / stride) * stride).max(stride);
    let rect_w = (((new_w + stride - 1) / stride) * stride).max(stride);

    (rect_h, rect_w)
}

/// Calculate letterbox parameters for resizing.
///
/// Computes new dimensions and padding to fit the image within the target size while maintaining aspect ratio.
///
/// # Arguments
///
/// * `orig_width` - Original image width.
/// * `orig_height` - Original image height.
/// * `target_size` - Target size as (height, width).
/// * `stride` - Model stride for alignment (unused in calculation but kept for API compatibility).
///
/// # Returns
///
/// Tuple containing:
/// 1. `new_width`: Scaled width.
/// 2. `new_height`: Scaled height.
/// 3. `pad_left`: Left padding.
/// 4. `pad_top`: Top padding.
/// 5. `(scale_y, scale_x)`: Scale factors.
fn calculate_letterbox_params(
    orig_width: u32,
    orig_height: u32,
    target_size: (usize, usize),
    _stride: u32,
) -> (LetterboxGeometry, (f32, f32)) {
    // A single uniform `gain` on both axes for coordinate back-projection; per-axis gains
    // from the rounded extents can diverge slightly, shifting boxes and changing NMS.
    LetterboxGeometry::compute(orig_width, orig_height, target_size, false)
}

/// Convert an RGB image to a normalized NCHW tensor, planar (CHW) layout.
///
/// Allocates the `(1, 3, H, W)` tensor, splits it into per-channel slices, and fills each
/// pixel with `convert` applied to the source byte. The caller supplies the element type's
/// zero value and the per-byte conversion, so the f32 and f16 callers each keep their exact
/// arithmetic (the f16 one hoists its scale constant into the closure).
fn image_to_tensor<T: Clone>(
    image: &RgbImage,
    zero: T,
    mut convert: impl FnMut(u8) -> T,
) -> Array4<T> {
    let (width, height) = image.dimensions();
    let (w, h) = (width as usize, height as usize);
    let pixels = image.as_raw();

    let mut tensor = Array4::from_elem((1, 3, h, w), zero);

    // Get mutable slices for each channel for faster access
    let (r_slice, rest) = tensor.as_slice_mut().unwrap().split_at_mut(h * w);
    let (g_slice, b_slice) = rest.split_at_mut(h * w);

    for (i, chunk) in pixels.as_chunks::<3>().0.iter().enumerate() {
        r_slice[i] = convert(chunk[0]);
        g_slice[i] = convert(chunk[1]);
        b_slice[i] = convert(chunk[2]);
    }

    tensor
}

/// Convert a `DynamicImage` to an HWC ndarray.
///
/// # Panics
///
/// Panics if the array cannot be created from the image pixels (e.g. dimension mismatch).
#[must_use]
pub fn image_to_array(image: &DynamicImage) -> Array3<u8> {
    let rgb = image.to_rgb8();
    let (width, height) = rgb.dimensions();
    let pixels = rgb.into_raw();

    Array3::from_shape_vec((height as usize, width as usize, 3), pixels)
        .expect("Failed to create array from image pixels")
}

/// Scale coordinates from model output space back to original image space.
///
/// # Arguments
///
/// * `coords` - Coordinates in model space (after letterbox).
/// * `scale` - Scale factors (`scale_y`, `scale_x`) from preprocessing.
/// * `padding` - Padding (`pad_top`, `pad_left`) from preprocessing.
///
/// # Returns
///
/// Coordinates in original image space.
#[must_use]
pub fn scale_coords(coords: &[f32; 4], scale: (f32, f32), padding: (f32, f32)) -> [f32; 4] {
    let (scale_y, scale_x) = scale;
    let (pad_top, pad_left) = padding;

    [
        (coords[0] - pad_left) / scale_x, // x1
        (coords[1] - pad_top) / scale_y,  // y1
        (coords[2] - pad_left) / scale_x, // x2
        (coords[3] - pad_top) / scale_y,  // y2
    ]
}

/// Clip coordinates to image bounds.
///
/// # Arguments
///
/// * `coords` - Box coordinates [x1, y1, x2, y2].
/// * `shape` - Image shape (height, width).
///
/// # Returns
///
/// Clipped coordinates.
#[must_use]
pub const fn clip_coords(coords: &[f32; 4], shape: (u32, u32)) -> [f32; 4] {
    #[allow(clippy::cast_precision_loss)]
    let (h, w) = (shape.0 as f32, shape.1 as f32);
    [
        coords[0].clamp(0.0, w),
        coords[1].clamp(0.0, h),
        coords[2].clamp(0.0, w),
        coords[3].clamp(0.0, h),
    ]
}

/// Preprocess an image for YOLO classification (Center Crop).
///
/// Resizes the image so the shortest side matches `target_size`, then center crops.
///
/// A zero-width or zero-height source yields a `target_size` frame of the letterbox
/// padding color, matching [`preprocess_image`]. `target_size` is returned as requested,
/// so a zero target yields a zero-extent tensor there too.
///
/// # Arguments
///
/// * `image` - Input image.
/// * `target_size` - Target size as (height, width).
/// * `quantize` - Requested precision. FP16 also generates an FP16 tensor.
///
/// # Returns
///
/// Preprocessed tensor and transform information.
#[must_use]
pub fn preprocess_image_center_crop(
    image: &DynamicImage,
    target_size: (usize, usize),
    quantize: impl IntoQuantization,
) -> PreprocessResult {
    let quantize = quantize.into_quantization();
    let (orig_width, orig_height) = image.dimensions();
    let orig_shape = (orig_height, orig_width);

    // Perform center crop resize
    let (cropped, scale) = center_crop_image(image, target_size);

    // Convert to normalized NCHW tensor
    let tensor = image_to_tensor(&cropped, 0.0, |v| f32::from(v) / 255.0);

    // Optionally compute FP16 tensor, converting u8 straight to f16 with a hoisted 1/255.
    let tensor_f16 = (quantize == Some(Quantization::Fp16)).then(|| {
        let scale = f16::from_f32(1.0 / 255.0);
        image_to_tensor(&cropped, f16::ZERO, move |v| {
            f16::from_f32(f32::from(v)) * scale
        })
    });

    // For classification, we don't need complex coordinate mapping back to original
    // But we provide approximate scale/padding to satisfy strict types if needed.
    // In classification, we rarely map bounding boxes back, so these are less critical.
    let padding = (0.0, 0.0);

    PreprocessResult {
        tensor,
        tensor_f16,
        orig_shape,
        scale,
        padding,
    }
}

/// Resize and center crop image.
///
/// Resizes the image such that the shortest side equals the target dimension,
/// maintaining aspect ratio, then crops the center `target_size`.
///
/// # Arguments
///
/// * `image` - Source dynamic image.
/// * `target_size` - Desired output dimensions (height, width).
///
/// # Returns
///
/// Tuple containing:
/// 1. `cropped`: The processed `RgbImage`.
/// 2. `scale`: Scale factors applied (same for x and y).
#[allow(clippy::similar_names)]
fn center_crop_image(image: &DynamicImage, target_size: (usize, usize)) -> (RgbImage, (f32, f32)) {
    use fast_image_resize::{
        PixelType, ResizeAlg, ResizeOptions, Resizer,
        images::{Image, ImageRef},
    };

    let (src_w, src_h) = image.dimensions();
    #[allow(clippy::cast_possible_truncation)]
    let (target_h, target_w) = (target_size.0 as u32, target_size.1 as u32);

    // Any failure below yields no usable pixels; return a blank frame rather than panicking.
    let blank = || {
        (
            RgbImage::from_pixel(target_w, target_h, image::Rgb(LETTERBOX_COLOR)),
            (1.0, 1.0),
        )
    };

    // The cover scale below divides by each source extent, so `target / 0` is infinite: the
    // resized extents come out garbage and the allocator is asked for terabytes.
    if src_w == 0 || src_h == 0 {
        return blank();
    }

    // Calculate scale to "cover" the target area
    // scale = max(target_w / src_w, target_h / src_h)
    #[allow(clippy::cast_precision_loss)]
    let scale_x = target_w as f32 / src_w as f32;
    #[allow(clippy::cast_precision_loss)]
    let scale_y = target_h as f32 / src_h as f32;
    let scale = scale_x.max(scale_y);

    let (new_w, new_h) = if scale_x >= scale_y {
        #[allow(
            clippy::cast_possible_truncation,
            clippy::cast_sign_loss,
            clippy::cast_precision_loss
        )]
        (target_w, (src_h as f32 * scale_x) as u32)
    } else {
        #[allow(
            clippy::cast_possible_truncation,
            clippy::cast_sign_loss,
            clippy::cast_precision_loss
        )]
        ((src_w as f32 * scale_y) as u32, target_h)
    };

    // Resize first. Borrow the source samples when the image is already RGB8 (the common
    // case): `to_rgb8` would clone the full frame just to hand it straight to the resizer.
    let owned_rgb;
    let src_bytes: &[u8] = match image {
        DynamicImage::ImageRgb8(rgb) => rgb.as_raw(),
        other => {
            owned_rgb = other.to_rgb8();
            owned_rgb.as_raw()
        }
    };
    let Ok(src_image) = ImageRef::new(src_w, src_h, src_bytes, PixelType::U8x3) else {
        return blank();
    };

    // Valid dimensions check
    let safe_new_w = new_w.max(1);
    let safe_new_h = new_h.max(1);

    let mut dst_image = Image::new(safe_new_w, safe_new_h, PixelType::U8x3);

    let mut resizer = Resizer::new();
    let options = ResizeOptions::new().resize_alg(ResizeAlg::Convolution(
        fast_image_resize::FilterType::Bilinear,
    ));
    if resizer
        .resize(&src_image, &mut dst_image, Some(&options))
        .is_err()
    {
        return blank();
    }

    // Convert back to RgbImage to crop
    let resized_buffer = dst_image.into_vec();
    let Some(resized_rgb) = RgbImage::from_raw(safe_new_w, safe_new_h, resized_buffer) else {
        return blank();
    };

    // Calculate crop offsets using Banker's Rounding (round half to even).
    #[allow(clippy::cast_precision_loss)]
    let crop_x_float = (new_w.saturating_sub(target_w)) as f32 / 2.0;
    #[allow(clippy::cast_precision_loss)]
    let crop_y_float = (new_h.saturating_sub(target_h)) as f32 / 2.0;

    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    let crop_x = bankers_round(crop_x_float) as u32;
    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
    let crop_y = bankers_round(crop_y_float) as u32;

    let cropped =
        image::imageops::crop_imm(&resized_rgb, crop_x, crop_y, target_w, target_h).to_image();

    (cropped, (scale, scale))
}

/// Round float to nearest integer, rounding half to even (Banker's Rounding).
fn bankers_round(v: f32) -> f32 {
    let n = v.floor();
    let d = v - n;
    if (d - 0.5).abs() < 1e-6 {
        if n % 2.0 == 0.0 { n } else { n + 1.0 }
    } else {
        v.round()
    }
}

#[allow(clippy::similar_names)]
#[cfg(test)]
mod tests {
    use super::*;

    /// An aspect ratio worse than the target size rounds one extent to zero, which used to
    /// discard the image and hand the model a uniformly gray tensor.
    #[test]
    fn test_extreme_aspect_ratio_keeps_image_content() {
        let (geom, _) = calculate_letterbox_params(10000, 1, (640, 640), 32);
        assert!(geom.new_h >= 1, "height collapsed to {}", geom.new_h);
        assert!(geom.new_w >= 1);

        let img = DynamicImage::ImageRgb8(image::RgbImage::from_pixel(
            10000,
            1,
            image::Rgb([255, 0, 0]),
        ));
        let res = preprocess_image(&img, (640, 640), 32);
        let tensor = res.tensor;
        assert!(
            tensor.iter().any(|&v| (v - LETTERBOX_NORM).abs() > 1e-6),
            "every pixel is letterbox fill, so the image was thrown away"
        );
    }

    /// A zero-dimension image underflowed `src_w - 1` in the LUT builder, which panics under
    /// the overflow checks the dev profile enables.
    #[test]
    fn test_center_crop_degenerate_input_does_not_panic() {
        for (w, h) in [(0, 0), (0, 64), (64, 0)] {
            let img = DynamicImage::ImageRgb8(image::RgbImage::new(w, h));
            let res = preprocess_image_center_crop(&img, (224, 224), false);
            assert_eq!(res.tensor.shape(), &[1, 3, 224, 224]);
        }
    }

    #[test]
    fn test_zero_dimension_image_does_not_panic() {
        for (w, h) in [(0, 0), (0, 64), (64, 0)] {
            let img = DynamicImage::ImageRgb8(image::RgbImage::new(w, h));
            let res = preprocess_image(&img, (640, 640), 32);
            assert_eq!(res.tensor.shape(), &[1, 3, 640, 640]);
            assert!(
                res.tensor
                    .iter()
                    .all(|&v| (v - LETTERBOX_NORM).abs() < 1e-6),
                "an empty source has no pixels, so the tensor is all letterbox fill"
            );

            // `rect` mode picks its own target before preprocessing, so that target has to
            // stay nonzero too: a zero extent would hand the model an empty input and the
            // CUDA preprocess a zero-sized kernel grid.
            let rect = calculate_rect_size(w, h, (640, 640), 32);
            assert!(
                rect.0 >= 32 && rect.1 >= 32,
                "rect target {rect:?} collapsed for a {w}x{h} source"
            );
            let res = preprocess_image(&img, rect, 32);
            assert_eq!(res.tensor.shape(), &[1, 3, rect.0, rect.1]);
        }
    }

    #[test]
    fn test_letterbox_params() {
        // Square input into a square target: exact fit, no padding.
        let (geom, _scale) = calculate_letterbox_params(640, 640, (640, 640), 32);
        assert_eq!((geom.new_w, geom.new_h), (640, 640));
        assert_eq!((geom.pad_left, geom.pad_top), (0, 0));

        // Wide input is scaled down to fit and padded top/bottom, not left/right.
        let (geom, _) = calculate_letterbox_params(1280, 720, (640, 640), 32);
        assert!(geom.new_w <= 640 && geom.new_h <= 640);
        assert_eq!(geom.pad_left, 0);

        // Tall input is the mirror case: padded left/right, not top/bottom.
        let (geom, _) = calculate_letterbox_params(480, 640, (640, 640), 32);
        assert!(geom.pad_left > 0);
        assert_eq!(geom.pad_top, 0);
    }

    #[test]
    fn test_scale_coords() {
        let coords = [100.0, 100.0, 200.0, 200.0];
        let scale = (1.0, 1.0);
        let padding = (10.0, 10.0);

        let scaled = scale_coords(&coords, scale, padding);

        assert!((scaled[0] - 90.0).abs() < 1e-6);
        assert!((scaled[1] - 90.0).abs() < 1e-6);
        assert!((scaled[2] - 190.0).abs() < 1e-6);
        assert!((scaled[3] - 190.0).abs() < 1e-6);
    }

    #[test]
    fn test_clip_coords() {
        let coords = [-10.0, -20.0, 700.0, 500.0];
        let clipped = clip_coords(&coords, (480, 640));

        assert!((clipped[0] - 0.0).abs() < 1e-6);
        assert!((clipped[1] - 0.0).abs() < 1e-6);
        assert!((clipped[2] - 640.0).abs() < 1e-6);
        assert!((clipped[3] - 480.0).abs() < 1e-6);
    }

    #[test]
    fn test_preprocess_image_center_crop() {
        // Classification path: center-crop to the model size, normalize to [0, 1], and
        // emit the FP16 tensor only when asked. A non-square source exercises the crop.
        let img = image::DynamicImage::new_rgb8(400, 300);
        for quantize in [None, Some(Quantization::Fp16)] {
            let res = preprocess_image_center_crop(&img, (224, 224), quantize);
            assert_eq!(res.tensor.dim(), (1, 3, 224, 224));
            assert_eq!(res.orig_shape, (300, 400));
            assert_eq!(res.padding, (0.0, 0.0));
            assert!(res.tensor.iter().all(|v| (0.0..=1.0).contains(v)));
            assert_eq!(res.tensor_f16.is_some(), quantize.is_some());
            if let Some(t16) = &res.tensor_f16 {
                assert_eq!(t16.dim(), res.tensor.dim());
            }
        }
    }

    #[test]
    fn test_preprocess_image_static_centered_letterbox() {
        // 480×640 wide image, target=1024×1024, is_dynamic=false.
        // Scale = min(1024/480, 1024/640) = 1024/640 = 1.6; nh=768, nw=1024, pad_top≈128.
        let img = image::DynamicImage::new_rgb8(640, 480);
        let res = preprocess_image_with_precision(&img, (1024, 1024), 32, None);
        let (_, _, h, w) = res.tensor.dim();
        assert_eq!(h, 1024);
        assert_eq!(w, 1024);
        // Wide image: horizontal padding is 0, vertical padding is non-zero.
        assert!(res.padding.1.abs() < 1e-6, "wide image: no left padding");
        assert!(res.padding.0 > 0.0, "wide image: top padding expected");
    }

    #[test]
    fn test_preprocess_image_rect_uses_centered_letterbox() {
        // Mirrors Ultralytics LetterBox((1024, 1024), auto=True, stride=32) for a 333×640 image:
        // resized content is 533×1024 and centered in a 544×1024 rect with 5/6 vertical padding.
        let img = image::DynamicImage::new_rgb8(640, 333);
        let rect_size = calculate_rect_size(640, 333, (1024, 1024), 32);
        assert_eq!(rect_size, (544, 1024));
        let res = preprocess_image_with_precision(&img, rect_size, 32, None);
        let (_, _, h, w) = res.tensor.dim();
        assert_eq!((h, w), rect_size);
        assert_eq!(res.padding, (5.0, 0.0));
    }

    #[test]
    fn test_preprocess_image_stretch_fills_target() {
        // RT-DETR scale-fill: a 810x1080 source stretches onto 640x640 with no padding and
        // a separate gain per axis, so back-projection returns the original extents.
        let img = image::DynamicImage::new_rgb8(810, 1080);
        let res = preprocess_image_stretch(&img, (640, 640), None);
        let (_, c, h, w) = res.tensor.dim();
        assert_eq!((c, h, w), (3, 640, 640));
        assert_eq!(res.padding, (0.0, 0.0));
        assert_eq!(res.scale, (640.0 / 1080.0, 640.0 / 810.0));
        let full = scale_coords(&[0.0, 0.0, 640.0, 640.0], res.scale, res.padding);
        assert!((full[2] - 810.0).abs() < 1e-3 && (full[3] - 1080.0).abs() < 1e-3);

        // A zero-extent source has no gain to report: 1.0 keeps the back-projection finite
        // instead of dividing coordinates by infinity.
        let (geom, scale) = LetterboxGeometry::compute(0, 0, (640, 640), true);
        assert_eq!(scale, (1.0, 1.0));
        assert_eq!((geom.new_w, geom.new_h), (640, 640));
    }

    #[test]
    fn test_preprocess_image_public_wrapper() {
        let img = image::DynamicImage::new_rgb8(320, 240);
        let res = preprocess_image(&img, (640, 640), 32);
        let (_, c, h, w) = res.tensor.dim();
        assert_eq!((c, h, w), (3, 640, 640));
        assert!(res.tensor_f16.is_none());
    }

    #[test]
    fn test_preprocess_image_fp16_path() {
        let img = image::DynamicImage::new_rgb8(320, 240);
        let res = preprocess_image_with_precision(&img, (640, 640), 32, Some(Quantization::Fp16));
        // quantize=16 also produces the FP16 tensor mirroring the FP32 one.
        let f16 = res.tensor_f16.expect("fp16 tensor present");
        assert_eq!(f16.dim(), res.tensor.dim());
    }

    #[test]
    fn test_preprocess_various_aspect_ratios() {
        // Tall, wide, square, and tiny inputs all produce the requested tensor size.
        for (w, h) in [(100u32, 400u32), (400, 100), (1, 1), (640, 640)] {
            let img = image::DynamicImage::new_rgb8(w, h);
            let res = preprocess_image(&img, (320, 320), 32);
            let (_, c, th, tw) = res.tensor.dim();
            assert_eq!((c, th, tw), (3, 320, 320));
            assert_eq!(res.orig_shape, (h, w));
        }
    }

    #[test]
    fn test_x_lut_cache_reuse() {
        // Two preprocess calls at the same width reuse the cached x-LUT (cache-hit path).
        let img = image::DynamicImage::new_rgb8(200, 150);
        let a = preprocess_image(&img, (320, 320), 32);
        let b = preprocess_image(&img, (320, 320), 32);
        assert_eq!(a.tensor.dim(), b.tensor.dim());
    }

    #[test]
    fn test_calculate_rect_size() {
        // A square input at the target size is returned unchanged.
        assert_eq!(calculate_rect_size(640, 640, (640, 640), 32), (640, 640));

        // Portrait, landscape, and non-aligned inputs all round to stride multiples
        // and stay within the target.
        for (w, h) in [(400u32, 1000u32), (1000, 400), (800, 600)] {
            let (rh, rw) = calculate_rect_size(w, h, (640, 640), 32);
            assert_eq!((rh % 32, rw % 32), (0, 0));
            assert!(rh <= 640 && rw <= 640);
        }
    }
}