oar-ocr-core 0.6.3

//! Prediction result types for the OCR pipeline.
//!
//! This module defines various types and traits for representing and working with
//! prediction results in the OCR pipeline. It includes enums for different types
//! of predictions (detection, recognition, classification, rectification) and
//! traits for converting between different representations.

use image::RgbImage;
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::sync::Arc;

/// Enum representing different types of prediction results.
///
/// This enum is used to represent the results of different types of predictions
/// in the OCR pipeline, such as text detection, text recognition, image classification,
/// and image rectification.
///
/// # Type Parameters
///
/// * `'a` - The lifetime of the borrowed data.
/// * `I` - The type of the input images.
#[derive(Debug, Clone)]
pub enum PredictionResult<'a, I = Arc<RgbImage>> {
    /// Results from text detection.
    Detection {
        /// The input paths of the images.
        input_path: Vec<Cow<'a, str>>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        input_img: Vec<I>,
        /// The detected polygons.
        dt_polys: Vec<Vec<crate::processors::BoundingBox>>,
        /// The scores for the detected polygons.
        dt_scores: Vec<Vec<f32>>,
    },
    /// Results from text recognition.
    Recognition {
        /// The input paths of the images.
        input_path: Vec<Cow<'a, str>>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        input_img: Vec<I>,
        /// The recognized text.
        rec_text: Vec<Cow<'a, str>>,
        /// The scores for the recognized text.
        rec_score: Vec<f32>,
    },
    /// Results from image classification.
    Classification {
        /// The input paths of the images.
        input_path: Vec<Cow<'a, str>>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        input_img: Vec<I>,
        /// The class IDs for the classifications.
        class_ids: Vec<Vec<usize>>,
        /// The scores for the classifications.
        scores: Vec<Vec<f32>>,
        /// The label names for the classifications.
        label_names: Vec<Vec<Cow<'a, str>>>,
    },
    /// Results from image rectification.
    Rectification {
        /// The input paths of the images.
        input_path: Vec<Cow<'a, str>>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        input_img: Vec<I>,
        /// The rectified images.
        rectified_img: Vec<I>,
    },
}

/// Enum representing owned prediction results.
///
/// This enum is similar to PredictionResult, but uses owned String values instead
/// of borrowed Cow values. It also implements Serialize and Deserialize traits
/// for easy serialization and deserialization.
///
/// # Type Parameters
///
/// * `I` - The type of the input images.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OwnedPredictionResult<I = Arc<RgbImage>> {
    /// Results from text detection.
    Detection {
        /// The input paths of the images.
        input_path: Vec<String>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        #[serde(skip)]
        input_img: Vec<I>,
        /// The detected polygons.
        dt_polys: Vec<Vec<crate::processors::BoundingBox>>,
        /// The scores for the detected polygons.
        dt_scores: Vec<Vec<f32>>,
    },
    /// Results from text recognition.
    Recognition {
        /// The input paths of the images.
        input_path: Vec<String>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        #[serde(skip)]
        input_img: Vec<I>,
        /// The recognized text.
        rec_text: Vec<String>,
        /// The scores for the recognized text.
        rec_score: Vec<f32>,
    },
    /// Results from image classification.
    Classification {
        /// The input paths of the images.
        input_path: Vec<String>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        #[serde(skip)]
        input_img: Vec<I>,
        /// The class IDs for the classifications.
        class_ids: Vec<Vec<usize>>,
        /// The scores for the classifications.
        scores: Vec<Vec<f32>>,
        /// The label names for the classifications.
        label_names: Vec<Vec<String>>,
    },
    /// Results from image rectification.
    Rectification {
        /// The input paths of the images.
        input_path: Vec<String>,
        /// The indices of the images in the batch.
        index: Vec<usize>,
        /// The input images.
        #[serde(skip)]
        input_img: Vec<I>,
        /// The rectified images.
        #[serde(skip)]
        rectified_img: Vec<I>,
    },
}

/// Implementation of methods for PredictionResult.
impl<'a, I> PredictionResult<'a, I> {
    /// Gets the input paths of the images.
    ///
    /// # Returns
    ///
    /// A slice of the input paths.
    pub fn input_paths(&self) -> &[Cow<'a, str>] {
        match self {
            PredictionResult::Detection { input_path, .. } => input_path,
            PredictionResult::Recognition { input_path, .. } => input_path,
            PredictionResult::Classification { input_path, .. } => input_path,
            PredictionResult::Rectification { input_path, .. } => input_path,
        }
    }

    /// Gets the indices of the images in the batch.
    ///
    /// # Returns
    ///
    /// A slice of the indices.
    pub fn indices(&self) -> &[usize] {
        match self {
            PredictionResult::Detection { index, .. } => index,
            PredictionResult::Recognition { index, .. } => index,
            PredictionResult::Classification { index, .. } => index,
            PredictionResult::Rectification { index, .. } => index,
        }
    }

    /// Gets the input images.
    ///
    /// # Returns
    ///
    /// A slice of the input images.
    pub fn input_images(&self) -> &[I] {
        match self {
            PredictionResult::Detection { input_img, .. } => input_img,
            PredictionResult::Recognition { input_img, .. } => input_img,
            PredictionResult::Classification { input_img, .. } => input_img,
            PredictionResult::Rectification { input_img, .. } => input_img,
        }
    }

    /// Checks if the prediction result is a detection result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a detection result, false otherwise.
    pub fn is_detection(&self) -> bool {
        matches!(self, PredictionResult::Detection { .. })
    }

    /// Checks if the prediction result is a recognition result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a recognition result, false otherwise.
    pub fn is_recognition(&self) -> bool {
        matches!(self, PredictionResult::Recognition { .. })
    }

    /// Checks if the prediction result is a classification result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a classification result, false otherwise.
    pub fn is_classification(&self) -> bool {
        matches!(self, PredictionResult::Classification { .. })
    }

    /// Checks if the prediction result is a rectification result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a rectification result, false otherwise.
    pub fn is_rectification(&self) -> bool {
        matches!(self, PredictionResult::Rectification { .. })
    }

    /// Converts the prediction result to an owned prediction result.
    ///
    /// # Returns
    ///
    /// An OwnedPredictionResult with the same data.
    pub fn into_owned(self) -> OwnedPredictionResult<I> {
        match self {
            PredictionResult::Detection {
                input_path,
                index,
                input_img,
                dt_polys,
                dt_scores,
            } => OwnedPredictionResult::Detection {
                input_path: input_path.into_iter().map(|cow| cow.into_owned()).collect(),
                index,
                input_img,
                dt_polys,
                dt_scores,
            },
            PredictionResult::Recognition {
                input_path,
                index,
                input_img,
                rec_text,
                rec_score,
            } => OwnedPredictionResult::Recognition {
                input_path: input_path.into_iter().map(|cow| cow.into_owned()).collect(),
                index,
                input_img,
                rec_text: rec_text.into_iter().map(|cow| cow.into_owned()).collect(),
                rec_score,
            },
            PredictionResult::Classification {
                input_path,
                index,
                input_img,
                class_ids,
                scores,
                label_names,
            } => OwnedPredictionResult::Classification {
                input_path: input_path.into_iter().map(|cow| cow.into_owned()).collect(),
                index,
                input_img,
                class_ids,
                scores,
                label_names: label_names
                    .into_iter()
                    .map(|vec| vec.into_iter().map(|cow| cow.into_owned()).collect())
                    .collect(),
            },
            PredictionResult::Rectification {
                input_path,
                index,
                input_img,
                rectified_img,
            } => OwnedPredictionResult::Rectification {
                input_path: input_path.into_iter().map(|cow| cow.into_owned()).collect(),
                index,
                input_img,
                rectified_img,
            },
        }
    }
}

/// Implementation of methods for OwnedPredictionResult.
impl<I> OwnedPredictionResult<I> {
    /// Gets the input paths of the images.
    ///
    /// # Returns
    ///
    /// A slice of the input paths.
    pub fn input_paths(&self) -> &[String] {
        match self {
            OwnedPredictionResult::Detection { input_path, .. } => input_path,
            OwnedPredictionResult::Recognition { input_path, .. } => input_path,
            OwnedPredictionResult::Classification { input_path, .. } => input_path,
            OwnedPredictionResult::Rectification { input_path, .. } => input_path,
        }
    }

    /// Gets the indices of the images in the batch.
    ///
    /// # Returns
    ///
    /// A slice of the indices.
    pub fn indices(&self) -> &[usize] {
        match self {
            OwnedPredictionResult::Detection { index, .. } => index,
            OwnedPredictionResult::Recognition { index, .. } => index,
            OwnedPredictionResult::Classification { index, .. } => index,
            OwnedPredictionResult::Rectification { index, .. } => index,
        }
    }

    /// Gets the input images.
    ///
    /// # Returns
    ///
    /// A slice of the input images.
    pub fn input_images(&self) -> &[I] {
        match self {
            OwnedPredictionResult::Detection { input_img, .. } => input_img,
            OwnedPredictionResult::Recognition { input_img, .. } => input_img,
            OwnedPredictionResult::Classification { input_img, .. } => input_img,
            OwnedPredictionResult::Rectification { input_img, .. } => input_img,
        }
    }

    /// Checks if the prediction result is a detection result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a detection result, false otherwise.
    pub fn is_detection(&self) -> bool {
        matches!(self, OwnedPredictionResult::Detection { .. })
    }

    /// Checks if the prediction result is a recognition result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a recognition result, false otherwise.
    pub fn is_recognition(&self) -> bool {
        matches!(self, OwnedPredictionResult::Recognition { .. })
    }

    /// Checks if the prediction result is a classification result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a classification result, false otherwise.
    pub fn is_classification(&self) -> bool {
        matches!(self, OwnedPredictionResult::Classification { .. })
    }

    /// Checks if the prediction result is a rectification result.
    ///
    /// # Returns
    ///
    /// True if the prediction result is a rectification result, false otherwise.
    pub fn is_rectification(&self) -> bool {
        matches!(self, OwnedPredictionResult::Rectification { .. })
    }

    /// Converts the owned prediction result to a borrowed prediction result.
    ///
    /// # Returns
    ///
    /// A PredictionResult with borrowed data.
    pub fn as_prediction_result(&self) -> PredictionResult<'_, &I> {
        match self {
            OwnedPredictionResult::Detection {
                input_path,
                index,
                input_img,
                dt_polys,
                dt_scores,
            } => PredictionResult::Detection {
                input_path: input_path
                    .iter()
                    .map(|s| Cow::Borrowed(s.as_str()))
                    .collect(),
                index: index.clone(),
                input_img: input_img.iter().collect(),
                dt_polys: dt_polys.clone(),
                dt_scores: dt_scores.clone(),
            },
            OwnedPredictionResult::Recognition {
                input_path,
                index,
                input_img,
                rec_text,
                rec_score,
            } => PredictionResult::Recognition {
                input_path: input_path
                    .iter()
                    .map(|s| Cow::Borrowed(s.as_str()))
                    .collect(),
                index: index.clone(),
                input_img: input_img.iter().collect(),
                rec_text: rec_text.iter().map(|s| Cow::Borrowed(s.as_str())).collect(),
                rec_score: rec_score.clone(),
            },
            OwnedPredictionResult::Classification {
                input_path,
                index,
                input_img,
                class_ids,
                scores,
                label_names,
            } => PredictionResult::Classification {
                input_path: input_path
                    .iter()
                    .map(|s| Cow::Borrowed(s.as_str()))
                    .collect(),
                index: index.clone(),
                input_img: input_img.iter().collect(),
                class_ids: class_ids.clone(),
                scores: scores.clone(),
                label_names: label_names
                    .iter()
                    .map(|vec| vec.iter().map(|s| Cow::Borrowed(s.as_str())).collect())
                    .collect(),
            },
            OwnedPredictionResult::Rectification {
                input_path,
                index,
                input_img,
                rectified_img,
            } => PredictionResult::Rectification {
                input_path: input_path
                    .iter()
                    .map(|s| Cow::Borrowed(s.as_str()))
                    .collect(),
                index: index.clone(),
                input_img: input_img.iter().collect(),
                rectified_img: rectified_img.iter().collect(),
            },
        }
    }
}

/// Trait for converting a type into a prediction result.
///
/// This trait is used to convert a type into a prediction result.
pub trait IntoPrediction {
    /// The output type.
    type Out;
    /// Converts the type into a prediction result.
    ///
    /// # Returns
    ///
    /// The prediction result.
    fn into_prediction(self) -> Self::Out;
}

/// Trait for converting a type into an owned prediction result.
///
/// This trait is used to convert a type into an owned prediction result.
pub trait IntoOwnedPrediction {
    /// The output type.
    type Out;
    /// Converts the type into an owned prediction result.
    ///
    /// # Returns
    ///
    /// The owned prediction result.
    fn into_owned_prediction(self) -> Self::Out;
}

/// Implementation of IntoOwnedPrediction for types that implement IntoPrediction.
///
/// This implementation allows types that implement IntoPrediction to be converted
/// into owned prediction results.
impl<T> IntoOwnedPrediction for T
where
    T: IntoPrediction,
    T::Out: Into<OwnedPredictionResult>,
{
    type Out = OwnedPredictionResult;

    fn into_owned_prediction(self) -> Self::Out {
        self.into_prediction().into()
    }
}

/// Implementation of From for converting PredictionResult to OwnedPredictionResult.
///
/// This implementation allows PredictionResult to be converted to OwnedPredictionResult.
impl<I> From<PredictionResult<'_, I>> for OwnedPredictionResult<I> {
    fn from(result: PredictionResult<'_, I>) -> Self {
        result.into_owned()
    }
}