use std::{
cmp::{max, min},
path::{Path, PathBuf},
};
use opencv::{
core::{self, Mat, MatTraitConst, Point, Rect, Size},
imgproc,
};
use serde::{Deserialize, Serialize};
use super::{
core::{default_image_threshold, validate_threshold, MatchResult, Roi},
error::{McvError, Result},
frame::{RgbaFrame, VisionTemplate},
io::read_image_gray,
};
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ImageTemplateOptions {
pub threshold: f64,
pub max_pyramid_level: i32,
pub min_pyramid_size: i32,
pub refine_margin: i32,
}
impl Default for ImageTemplateOptions {
fn default() -> Self {
Self {
threshold: default_image_threshold(),
max_pyramid_level: 4,
min_pyramid_size: 16,
refine_margin: 8,
}
}
}
pub struct ImageTemplate {
template: Mat,
pyramid: Vec<Mat>,
width: i32,
height: i32,
opts: ImageTemplateOptions,
roi: Option<Roi>,
}
#[derive(Debug, Clone)]
pub struct ImageTemplateBuilder {
path: PathBuf,
opts: ImageTemplateOptions,
roi: Option<Roi>,
}
impl ImageTemplate {
pub fn new(path: impl Into<PathBuf>) -> Result<Self> {
Self::builder(path).open()
}
pub fn builder(path: impl Into<PathBuf>) -> ImageTemplateBuilder {
ImageTemplateBuilder {
path: path.into(),
opts: ImageTemplateOptions::default(),
roi: None,
}
}
pub fn with_options(path: impl AsRef<Path>, opts: ImageTemplateOptions) -> Result<Self> {
Self::from_path(path, opts)
}
pub fn from_path(path: impl AsRef<Path>, opts: ImageTemplateOptions) -> Result<Self> {
validate_threshold(opts.threshold)?;
let template = read_image_gray(path)?;
if template.empty() {
return Err(McvError::InvalidImage("template is empty".to_string()));
}
let width = template.cols();
let height = template.rows();
if width <= 0 || height <= 0 {
return Err(McvError::InvalidImage(format!(
"invalid template size {width}x{height}"
)));
}
validate_template_variance(&template)?;
Ok(Self {
template: template.clone(),
pyramid: vec![template],
width,
height,
opts,
roi: None,
})
}
pub fn with_roi(mut self, roi: Option<Roi>) -> Self {
self.roi = roi;
self
}
pub fn with_threshold(mut self, threshold: f64) -> Result<Self> {
validate_threshold(threshold)?;
self.opts.threshold = threshold;
Ok(self)
}
pub fn width(&self) -> i32 {
self.width
}
pub fn height(&self) -> i32 {
self.height
}
pub fn find_path(
&mut self,
image_path: impl AsRef<Path>,
roi: Option<Roi>,
threshold: Option<f64>,
) -> Result<Option<MatchResult>> {
let image = read_image_gray(image_path)?;
self.find_in_gray(&image, roi, threshold)
}
pub fn find_in_gray(
&mut self,
image: &Mat,
roi: Option<Roi>,
threshold: Option<f64>,
) -> Result<Option<MatchResult>> {
let threshold = threshold.unwrap_or(self.opts.threshold);
validate_threshold(threshold)?;
if image.empty() {
return Err(McvError::InvalidImage("search image is empty".to_string()));
}
let image_width = image.cols();
let image_height = image.rows();
let effective_roi = roi
.unwrap_or_else(|| Roi::full(image_width, image_height))
.clamp(image_width, image_height)
.ok_or_else(|| {
McvError::InvalidRoi(
"ROI is outside the image or has non-positive size".to_string(),
)
})?;
if effective_roi.width < self.width || effective_roi.height < self.height {
return Ok(None);
}
let search = Mat::roi(image, effective_roi.as_rect())?;
let level = self.select_pyramid_level(search.cols(), search.rows());
if level == 0 {
return self.match_single(&search, effective_roi.x, effective_roi.y, threshold);
}
let mut search_pyramid: Vec<Mat> = Vec::with_capacity((level + 1) as usize);
search_pyramid.push(search.try_clone()?);
for lvl in 1..=level {
let mut down = Mat::default();
imgproc::pyr_down(
&search_pyramid[(lvl - 1) as usize],
&mut down,
Size::new(0, 0),
core::BORDER_DEFAULT,
)?;
search_pyramid.push(down);
}
self.ensure_template_pyramid(level)?;
let coarse_threshold = threshold.max(0.50) - 0.15;
let coarse_threshold = coarse_threshold.max(0.50);
for lvl in (1..=level).rev() {
let search_lvl = &search_pyramid[lvl as usize];
let template_lvl = &self.pyramid[lvl as usize];
if search_lvl.cols() < template_lvl.cols() || search_lvl.rows() < template_lvl.rows() {
continue;
}
let (loc, score) = best_match(search_lvl, template_lvl)?;
if score < coarse_threshold {
continue;
}
if let Some(found) = self.refine_at_base(
&search_pyramid[0],
loc,
lvl,
effective_roi.x,
effective_roi.y,
threshold,
)? {
return Ok(Some(found));
}
}
self.match_single(
&search_pyramid[0],
effective_roi.x,
effective_roi.y,
threshold,
)
}
pub fn find_all_path(
&self,
image_path: impl AsRef<Path>,
roi: Option<Roi>,
threshold: Option<f64>,
max_count: usize,
) -> Result<Vec<MatchResult>> {
let image = read_image_gray(image_path)?;
self.find_all_in_gray(&image, roi, threshold, max_count)
}
pub fn find_all_in_gray(
&self,
image: &Mat,
roi: Option<Roi>,
threshold: Option<f64>,
max_count: usize,
) -> Result<Vec<MatchResult>> {
let threshold = threshold.unwrap_or(self.opts.threshold);
validate_threshold(threshold)?;
if max_count == 0 {
return Ok(Vec::new());
}
let effective_roi = roi
.unwrap_or_else(|| Roi::full(image.cols(), image.rows()))
.clamp(image.cols(), image.rows())
.ok_or_else(|| {
McvError::InvalidRoi(
"ROI is outside the image or has non-positive size".to_string(),
)
})?;
if effective_roi.width < self.width || effective_roi.height < self.height {
return Ok(Vec::new());
}
let search = Mat::roi(image, effective_roi.as_rect())?;
let mut result = Mat::default();
imgproc::match_template(
&search,
&self.template,
&mut result,
imgproc::TM_CCOEFF_NORMED,
&Mat::default(),
)?;
let mut matches = Vec::new();
for y in 0..result.rows() {
for x in 0..result.cols() {
let score = *result.at_2d::<f32>(y, x)? as f64;
if score >= threshold {
matches.push(ScoredMatch {
result: MatchResult::from_box(
x + effective_roi.x,
y + effective_roi.y,
self.width,
self.height,
),
score: score.clamp(0.0, 1.0),
});
}
}
}
matches.sort_by(|a, b| b.score.total_cmp(&a.score));
Ok(nms(matches, 0.5, max_count))
}
fn match_single<T: opencv::core::ToInputArray>(
&self,
search: &T,
offset_x: i32,
offset_y: i32,
threshold: f64,
) -> Result<Option<MatchResult>> {
let (loc, score) = best_match(search, &self.template)?;
if score < threshold {
return Ok(None);
}
Ok(Some(MatchResult::from_box(
offset_x + loc.x,
offset_y + loc.y,
self.width,
self.height,
)))
}
fn refine_at_base(
&self,
base_image: &Mat,
coarse_loc: Point,
coarse_level: i32,
offset_x: i32,
offset_y: i32,
threshold: f64,
) -> Result<Option<MatchResult>> {
let scale = 1_i32 << coarse_level;
let margin = max(scale * 2, self.opts.refine_margin);
let cx = coarse_loc.x * scale;
let cy = coarse_loc.y * scale;
let x1 = max(0, cx - margin);
let y1 = max(0, cy - margin);
let x2 = min(base_image.cols(), cx + self.width + margin);
let y2 = min(base_image.rows(), cy + self.height + margin);
if x2 - x1 < self.width || y2 - y1 < self.height {
return Ok(None);
}
let window = Mat::roi(base_image, Rect::new(x1, y1, x2 - x1, y2 - y1))?;
let (loc, score) = best_match(&window, &self.template)?;
if score < threshold {
return Ok(None);
}
Ok(Some(MatchResult::from_box(
offset_x + x1 + loc.x,
offset_y + y1 + loc.y,
self.width,
self.height,
)))
}
fn select_pyramid_level(&self, search_width: i32, search_height: i32) -> i32 {
let min_dim = [search_width, search_height, self.width, self.height]
.into_iter()
.min()
.unwrap_or(0);
if min_dim < 2 * self.opts.min_pyramid_size {
return 0;
}
let ratio = min_dim as f64 / self.opts.min_pyramid_size as f64;
let max_level = ratio.log2().floor() as i32;
self.opts.max_pyramid_level.clamp(0, max_level.max(0))
}
fn ensure_template_pyramid(&mut self, level: i32) -> Result<()> {
while self.pyramid.len() <= level as usize {
let mut down = Mat::default();
let last = self.pyramid.last().expect("template pyramid has base");
imgproc::pyr_down(last, &mut down, Size::new(0, 0), core::BORDER_DEFAULT)?;
self.pyramid.push(down);
}
Ok(())
}
}
fn validate_template_variance(template: &Mat) -> Result<()> {
let mut min_value = 0.0;
let mut max_value = 0.0;
core::min_max_loc(
template,
Some(&mut min_value),
Some(&mut max_value),
None,
None,
&Mat::default(),
)?;
if min_value == max_value {
return Err(McvError::InvalidImage(
"template must contain at least two distinct grayscale values".to_string(),
));
}
Ok(())
}
impl ImageTemplateBuilder {
pub fn options(mut self, opts: ImageTemplateOptions) -> Self {
self.opts = opts;
self
}
pub fn threshold(mut self, threshold: f64) -> Result<Self> {
validate_threshold(threshold)?;
self.opts.threshold = threshold;
Ok(self)
}
pub fn max_pyramid_level(mut self, max_pyramid_level: i32) -> Self {
self.opts.max_pyramid_level = max_pyramid_level;
self
}
pub fn min_pyramid_size(mut self, min_pyramid_size: i32) -> Self {
self.opts.min_pyramid_size = min_pyramid_size;
self
}
pub fn refine_margin(mut self, refine_margin: i32) -> Self {
self.opts.refine_margin = refine_margin;
self
}
pub fn roi(mut self, roi: Roi) -> Self {
self.roi = Some(roi);
self
}
pub fn optional_roi(mut self, roi: Option<Roi>) -> Self {
self.roi = roi;
self
}
pub fn open(self) -> Result<ImageTemplate> {
Ok(ImageTemplate::from_path(self.path, self.opts)?.with_roi(self.roi))
}
}
impl VisionTemplate for ImageTemplate {
fn find_with_roi(
&mut self,
image: &mut RgbaFrame<'_>,
roi: Option<Roi>,
) -> Result<Option<MatchResult>> {
let roi = roi.or(self.roi);
let gray = image.gray()?;
self.find_in_gray(gray, roi, None)
}
}
pub(crate) fn best_match<T: opencv::core::ToInputArray, U: opencv::core::ToInputArray>(
search: &T,
template: &U,
) -> Result<(Point, f64)> {
let mut result = Mat::default();
imgproc::match_template(
search,
template,
&mut result,
imgproc::TM_CCOEFF_NORMED,
&Mat::default(),
)?;
let mut min_val = 0.0;
let mut max_val = 0.0;
let mut min_loc = Point::default();
let mut max_loc = Point::default();
core::min_max_loc(
&result,
Some(&mut min_val),
Some(&mut max_val),
Some(&mut min_loc),
Some(&mut max_loc),
&Mat::default(),
)?;
Ok((max_loc, max_val.clamp(0.0, 1.0)))
}
struct ScoredMatch {
result: MatchResult,
score: f64,
}
fn nms(mut matches: Vec<ScoredMatch>, iou_threshold: f64, max_count: usize) -> Vec<MatchResult> {
let mut keep = Vec::new();
while !matches.is_empty() && keep.len() < max_count {
let current = matches.remove(0);
matches.retain(|other| iou(¤t.result, &other.result) <= iou_threshold);
keep.push(current.result);
}
keep
}
fn iou(a: &MatchResult, b: &MatchResult) -> f64 {
let (a_x1, a_y1) = a.top_left();
let (a_x2, a_y2) = a.bottom_right_exclusive();
let (b_x1, b_y1) = b.top_left();
let (b_x2, b_y2) = b.bottom_right_exclusive();
let x1 = max(a_x1, b_x1);
let y1 = max(a_y1, b_y1);
let x2 = min(a_x2, b_x2);
let y2 = min(a_y2, b_y2);
let inter_w = max(0, x2 - x1) as f64;
let inter_h = max(0, y2 - y1) as f64;
let intersection = inter_w * inter_h;
let area_a = (a.width() * a.height()) as f64;
let area_b = (b.width() * b.height()) as f64;
intersection / (area_a + area_b - intersection + 1e-8)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_constant_template_before_normalized_correlation() -> Result<()> {
let template = Mat::from_slice_2d(&[&[128_u8, 128], &[128, 128]])?;
let error = validate_template_variance(&template).expect_err("constant template");
assert!(matches!(error, McvError::InvalidImage(_)));
Ok(())
}
#[test]
fn accepts_template_with_nonzero_variance() -> Result<()> {
let template = Mat::from_slice_2d(&[&[0_u8, 128], &[128, 255]])?;
validate_template_variance(&template)?;
Ok(())
}
}