use std::path::Path;
use opencv::{
core::{Mat, MatTraitConst},
prelude::MatTraitConstManual,
};
use serde::{Deserialize, Serialize};
use super::{
core::{default_color_threshold, validate_threshold, MatchResult, Roi},
error::{McvError, Result},
frame::{RgbaFrame, VisionTemplate},
io::{read_image_bgr, validate_bgr_image},
};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ColorSpec {
Hex(String),
Rgb([u8; 3]),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RgbColor {
pub r: u8,
pub g: u8,
pub b: u8,
}
impl RgbColor {
pub fn from_spec(spec: &ColorSpec) -> Result<Self> {
match spec {
ColorSpec::Rgb([r, g, b]) => Ok(Self {
r: *r,
g: *g,
b: *b,
}),
ColorSpec::Hex(value) => Self::from_hex(value),
}
}
pub fn from_hex(value: &str) -> Result<Self> {
let code = value.trim().trim_start_matches('#');
if code.len() != 6 {
return Err(McvError::InvalidColor(format!(
"hex color must contain exactly 6 characters: {value}"
)));
}
let r = u8::from_str_radix(&code[0..2], 16)
.map_err(|err| McvError::InvalidColor(format!("invalid red channel: {err}")))?;
let g = u8::from_str_radix(&code[2..4], 16)
.map_err(|err| McvError::InvalidColor(format!("invalid green channel: {err}")))?;
let b = u8::from_str_radix(&code[4..6], 16)
.map_err(|err| McvError::InvalidColor(format!("invalid blue channel: {err}")))?;
Ok(Self { r, g, b })
}
fn max_bgr_diff(self, bgr: &[u8]) -> u8 {
let b_diff = self.b.abs_diff(bgr[0]);
let g_diff = self.g.abs_diff(bgr[1]);
let r_diff = self.r.abs_diff(bgr[2]);
b_diff.max(g_diff).max(r_diff)
}
}
impl From<&str> for ColorSpec {
fn from(value: &str) -> Self {
Self::Hex(value.to_string())
}
}
impl From<String> for ColorSpec {
fn from(value: String) -> Self {
Self::Hex(value)
}
}
impl From<[u8; 3]> for ColorSpec {
fn from(value: [u8; 3]) -> Self {
Self::Rgb(value)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ColorOffset {
pub dx: i32,
pub dy: i32,
pub color: ColorSpec,
}
impl ColorOffset {
pub fn new(dx: i32, dy: i32, color: impl Into<ColorSpec>) -> Self {
Self {
dx,
dy,
color: color.into(),
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MultiColorTemplateOptions {
pub threshold: f64,
}
impl Default for MultiColorTemplateOptions {
fn default() -> Self {
Self {
threshold: default_color_threshold(),
}
}
}
pub struct MultiColorTemplate {
first_color: RgbColor,
offsets: Vec<(i32, i32, RgbColor)>,
opts: MultiColorTemplateOptions,
roi: Option<Roi>,
}
#[derive(Debug, Clone)]
pub struct MultiColorTemplateBuilder {
first_color: ColorSpec,
offsets: Vec<ColorOffset>,
opts: MultiColorTemplateOptions,
roi: Option<Roi>,
}
impl MultiColorTemplate {
pub fn builder(
first_color: impl Into<ColorSpec>,
offsets: impl IntoIterator<Item = ColorOffset>,
) -> MultiColorTemplateBuilder {
MultiColorTemplateBuilder {
first_color: first_color.into(),
offsets: offsets.into_iter().collect(),
opts: MultiColorTemplateOptions::default(),
roi: None,
}
}
pub fn with_options(
first_color: impl Into<ColorSpec>,
offsets: impl IntoIterator<Item = ColorOffset>,
opts: MultiColorTemplateOptions,
) -> Result<Self> {
let offsets = offsets.into_iter().collect::<Vec<_>>();
validate_threshold(opts.threshold)?;
validate_offsets(&offsets)?;
let first_color = first_color.into();
let first_color = RgbColor::from_spec(&first_color)?;
let offsets = offsets
.iter()
.map(|offset| Ok((offset.dx, offset.dy, RgbColor::from_spec(&offset.color)?)))
.collect::<Result<Vec<_>>>()?;
Ok(Self {
first_color,
offsets,
opts,
roi: None,
})
}
pub fn new(
first_color: impl Into<ColorSpec>,
offsets: impl IntoIterator<Item = ColorOffset>,
) -> Result<Self> {
Self::with_options(first_color, offsets, MultiColorTemplateOptions::default())
}
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 find_path(
&self,
image_path: impl AsRef<Path>,
roi: Option<Roi>,
threshold: Option<f64>,
) -> Result<Option<MatchResult>> {
let image = read_image_bgr(image_path)?;
let mut results = self.find_all_in_bgr(&image, roi, threshold, 1)?;
Ok(results.pop())
}
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_bgr(image_path)?;
self.find_all_in_bgr(&image, roi, threshold, max_count)
}
pub fn find_in_bgr(
&self,
image: &Mat,
roi: Option<Roi>,
threshold: Option<f64>,
) -> Result<Option<MatchResult>> {
let mut results = self.find_all_in_bgr(image, roi, threshold, 1)?;
Ok(results.pop())
}
pub fn find_all_in_bgr(
&self,
image: &Mat,
roi: Option<Roi>,
threshold: Option<f64>,
max_count: usize,
) -> Result<Vec<MatchResult>> {
if max_count == 0 {
return Ok(Vec::new());
}
let threshold = threshold.unwrap_or(self.opts.threshold);
validate_threshold(threshold)?;
validate_bgr_image(image)?;
let width = image.cols();
let height = image.rows();
let effective_roi = roi
.unwrap_or_else(|| Roi::full(width, height))
.clamp(width, height)
.ok_or_else(|| {
McvError::InvalidRoi(
"ROI is outside the image or has non-positive size".to_string(),
)
})?;
let tolerance = ((1.0 - threshold) * 255.0).round().clamp(0.0, 255.0) as u8;
let bytes = image.data_bytes()?;
let stride = width as usize * 3;
if bytes.len() < stride * height as usize {
return Err(McvError::InvalidImage(
"BGR image buffer is smaller than rows * cols * 3".to_string(),
));
}
let mut results = Vec::new();
let y_end = effective_roi.y + effective_roi.height;
let x_end = effective_roi.x + effective_roi.width;
for y in effective_roi.y..y_end {
'candidate: for x in effective_roi.x..x_end {
let idx = pixel_index(x, y, width);
let mut max_diff = self.first_color.max_bgr_diff(&bytes[idx..idx + 3]);
if max_diff > tolerance {
continue;
}
for (dx, dy, color) in &self.offsets {
let tx = x + dx;
let ty = y + dy;
if tx < effective_roi.x || ty < effective_roi.y || tx >= x_end || ty >= y_end {
continue 'candidate;
}
let offset_idx = pixel_index(tx, ty, width);
let diff = color.max_bgr_diff(&bytes[offset_idx..offset_idx + 3]);
if diff > tolerance {
continue 'candidate;
}
max_diff = max_diff.max(diff);
}
let score = 1.0 - f64::from(max_diff) / 255.0;
results.push(ScoredColorMatch {
result: MatchResult::from_point(x, y),
score,
});
}
}
results.sort_by(|a, b| b.score.total_cmp(&a.score));
results.truncate(max_count);
Ok(results.into_iter().map(|item| item.result).collect())
}
}
fn validate_offsets(offsets: &[ColorOffset]) -> Result<()> {
if offsets.is_empty() {
return Err(McvError::InvalidColor(
"multi-color template requires at least one offset color".to_string(),
));
}
Ok(())
}
impl MultiColorTemplateBuilder {
pub fn options(mut self, opts: MultiColorTemplateOptions) -> Self {
self.opts = opts;
self
}
pub fn offsets(mut self, offsets: impl IntoIterator<Item = ColorOffset>) -> Self {
self.offsets = offsets.into_iter().collect();
self
}
pub fn threshold(mut self, threshold: f64) -> Result<Self> {
validate_threshold(threshold)?;
self.opts.threshold = threshold;
Ok(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<MultiColorTemplate> {
Ok(
MultiColorTemplate::with_options(self.first_color, self.offsets, self.opts)?
.with_roi(self.roi),
)
}
}
struct ScoredColorMatch {
result: MatchResult,
score: f64,
}
impl VisionTemplate for MultiColorTemplate {
fn find_with_roi(
&mut self,
image: &mut RgbaFrame<'_>,
roi: Option<Roi>,
) -> Result<Option<MatchResult>> {
let bgr = image.bgr()?;
self.find_in_bgr(bgr, roi.or(self.roi), None)
}
}
fn pixel_index(x: i32, y: i32, image_width: i32) -> usize {
((y as usize * image_width as usize) + x as usize) * 3
}
#[cfg(test)]
mod tests {
use super::*;
use opencv::{core::Vec3b, prelude::MatTrait};
fn test_image(width: i32, height: i32) -> Result<Mat> {
let black = Vec3b::from([0, 0, 0]);
let row = vec![black; width as usize];
let rows = vec![row; height as usize];
let refs = rows.iter().map(Vec::as_slice).collect::<Vec<_>>();
Ok(Mat::from_slice_2d(&refs)?)
}
fn set_bgr(image: &mut Mat, x: i32, y: i32, bgr: [u8; 3]) -> Result<()> {
*image.at_2d_mut::<Vec3b>(y, x)? = Vec3b::from(bgr);
Ok(())
}
fn color_template(
first: &str,
offsets: Vec<ColorOffset>,
threshold: f64,
) -> Result<MultiColorTemplate> {
MultiColorTemplate::with_options(
ColorSpec::Hex(first.to_string()),
offsets,
MultiColorTemplateOptions { threshold },
)
}
#[test]
fn parses_hex_color() -> Result<()> {
let color = RgbColor::from_hex("#FF5500")?;
assert_eq!(
color,
RgbColor {
r: 255,
g: 85,
b: 0
}
);
assert!(RgbColor::from_hex("FFF").is_err());
assert!(RgbColor::from_hex("GGGGGG").is_err());
Ok(())
}
#[test]
fn rejects_empty_offsets() {
let template = color_template("1E841E", Vec::new(), 1.0);
assert!(template.is_err());
}
#[test]
fn finds_exact_color_pattern() -> Result<()> {
let mut image = test_image(100, 100)?;
set_bgr(&mut image, 50, 50, [30, 132, 30])?;
set_bgr(&mut image, 51, 50, [30, 132, 30])?;
let template = color_template("1E841E", vec![ColorOffset::new(1, 0, "1E841E")], 1.0)?;
let result = template.find_in_bgr(&image, None, None)?.expect("match");
assert_eq!(result.center(), (50, 50));
assert!(result.is_point());
assert_eq!(result.bbox.width, 1);
assert_eq!(result.bbox.height, 1);
Ok(())
}
#[test]
fn finds_with_offsets() -> Result<()> {
let mut image = test_image(100, 100)?;
set_bgr(&mut image, 50, 50, [30, 132, 30])?;
set_bgr(&mut image, 60, 55, [46, 148, 46])?;
let template = color_template(
"1E841E",
vec![ColorOffset {
dx: 10,
dy: 5,
color: ColorSpec::Hex("2E942E".to_string()),
}],
1.0,
)?;
let result = template.find_in_bgr(&image, None, None)?.expect("match");
assert_eq!(result.center(), (50, 50));
Ok(())
}
#[test]
fn offset_mismatch_does_not_skip_later_candidate_in_same_row() -> Result<()> {
let mut image = test_image(40, 20)?;
set_bgr(&mut image, 10, 10, [30, 132, 30])?;
set_bgr(&mut image, 20, 10, [30, 132, 30])?;
set_bgr(&mut image, 21, 10, [46, 148, 46])?;
let template = color_template("1E841E", vec![ColorOffset::new(1, 0, "2E942E")], 1.0)?;
let result = template.find_in_bgr(&image, None, None)?.expect("match");
assert_eq!(result.center(), (20, 10));
Ok(())
}
#[test]
fn respects_threshold_roi_and_max_count() -> Result<()> {
let mut image = test_image(100, 100)?;
set_bgr(&mut image, 20, 20, [35, 137, 35])?;
set_bgr(&mut image, 21, 20, [35, 137, 35])?;
set_bgr(&mut image, 60, 60, [30, 132, 30])?;
set_bgr(&mut image, 61, 60, [30, 132, 30])?;
set_bgr(&mut image, 70, 70, [30, 132, 30])?;
set_bgr(&mut image, 71, 70, [30, 132, 30])?;
let exact = color_template("1E841E", vec![ColorOffset::new(1, 0, "1E841E")], 1.0)?;
assert!(exact
.find_in_bgr(
&image,
Some(Roi {
x: 0,
y: 0,
width: 40,
height: 40
}),
None
)?
.is_none());
let tolerant = color_template("1E841E", vec![ColorOffset::new(1, 0, "1E841E")], 0.95)?;
let roi_result = tolerant
.find_in_bgr(
&image,
Some(Roi {
x: 0,
y: 0,
width: 40,
height: 40,
}),
None,
)?
.expect("tolerant match");
assert_eq!(roi_result.center(), (20, 20));
let all = exact.find_all_in_bgr(
&image,
Some(Roi {
x: 50,
y: 50,
width: 40,
height: 40,
}),
None,
1,
)?;
assert_eq!(all.len(), 1);
assert_eq!(all[0].center(), (60, 60));
Ok(())
}
}