From 86533d9c415d712fecf1a0dec0beab5d304f383b Mon Sep 17 00:00:00 2001 From: Jamjamjon <51357717+jamjamjon@users.noreply.github.com> Date: Fri, 23 May 2025 17:41:24 +0800 Subject: [PATCH] Enhance Annotator with mask-cutout functionality (#101) Co-authored-by: jamjamjon --- Cargo.toml | 1 - examples/rmbg/main.rs | 3 ++- src/io/hub.rs | 7 +++---- src/viz/drawable/mask.rs | 45 +++++++++++++++++++++++++++++++++------- src/viz/drawable/mod.rs | 8 +++---- src/viz/styles.rs | 42 ++++++++++++++++++++----------------- 6 files changed, 70 insertions(+), 36 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 8c56a34..321ae5e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,7 +32,6 @@ natord = "1.0.9" geo = "0.30.0" chrono = "0.4.40" regex = "1.11.1" -sha2 = "0.10.8" tempfile = "3.19.1" video-rs = { version = "0.10.3", features = ["ndarray"], optional = true } fast_image_resize = { version = "5.1.2", features = ["image"] } diff --git a/examples/rmbg/main.rs b/examples/rmbg/main.rs index b1456d9..0d685a3 100644 --- a/examples/rmbg/main.rs +++ b/examples/rmbg/main.rs @@ -43,7 +43,8 @@ fn main() -> anyhow::Result<()> { let ys = model.forward(&xs)?; // annotate - let annotator = Annotator::default(); + let annotator = + Annotator::default().with_mask_style(usls::Style::mask().with_mask_cutout(true)); for (x, y) in xs.iter().zip(ys.iter()) { annotator.annotate(x, y)?.save(format!( "{}.jpg", diff --git a/src/io/hub.rs b/src/io/hub.rs index 7d6c750..e29e754 100644 --- a/src/io/hub.rs +++ b/src/io/hub.rs @@ -2,7 +2,6 @@ use anyhow::{Context, Result}; use indicatif::ProgressStyle; use regex::Regex; use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; use std::time::Duration; @@ -461,9 +460,9 @@ impl Hub { } fn cache_file(owner: &str, repo: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(format!("{}-{}", owner, repo)); - format!(".{:x}", hasher.finalize()) + let safe_owner = owner.replace(|c: char| !c.is_ascii_alphanumeric(), "_"); + let safe_repo = repo.replace(|c: char| !c.is_ascii_alphanumeric(), "_"); + format!(".releases_{}_{}.json", safe_owner, safe_repo) } fn get_releases(owner: &str, repo: &str, to: &Dir, ttl: &Duration) -> Result> { diff --git a/src/viz/drawable/mask.rs b/src/viz/drawable/mask.rs index 41b81d6..62ce7ce 100644 --- a/src/viz/drawable/mask.rs +++ b/src/viz/drawable/mask.rs @@ -1,8 +1,8 @@ use crate::Drawable; use anyhow::Result; -use image::{DynamicImage, RgbaImage}; +use image::{DynamicImage, Rgba, RgbaImage}; -use crate::{ColorMap256, DrawContext, Mask, Style}; +use crate::{Color, ColorMap256, DrawContext, Mask, Style}; fn render_mask(mask: &Mask, colormap256: Option<&ColorMap256>) -> DynamicImage { if let Some(colormap256) = colormap256 { @@ -10,12 +10,26 @@ fn render_mask(mask: &Mask, colormap256: Option<&ColorMap256>) -> DynamicImage { let idx = p[0]; image::Rgb(colormap256.data()[idx as usize].rgb().into()) }); - DynamicImage::from(luma) + luma.into() } else { - DynamicImage::from(mask.mask().clone()) + mask.mask().clone().into() } } +fn apply_mask(origin: &RgbaImage, mask: &Mask, background_color: Option) -> DynamicImage { + let bg = background_color.unwrap_or(Color::green()); + imageproc::map::map_colors2(origin, mask.mask(), |src, mask| { + let [r, g, b, _] = src.0; + let mask_alpha = mask.0[0]; + if mask_alpha == 0 { + Rgba(bg.into()) + } else { + Rgba([r, g, b, mask_alpha]) + } + }) + .into() +} + fn best_grid(n: usize) -> (usize, usize) { let mut best_rows = 1; let mut best_cols = n; @@ -38,6 +52,8 @@ fn draw_masks( masks: &[&Mask], colormap256: Option<&ColorMap256>, canvas: &mut RgbaImage, + mask_cutout: bool, + mask_background_color: Option, ) -> Result<()> { let (w, h) = canvas.dimensions(); let n = masks.len() + 1; // +1 for original @@ -58,7 +74,11 @@ fn draw_masks( let x = ((w as i32 - mw as i32) / 2).max(0) as u32; let y = ((h as i32 - mh as i32) / 2).max(0) as u32; - let mask_dyn = render_mask(mask, colormap256); + let mask_dyn = if mask_cutout { + apply_mask(canvas, mask, mask_background_color) + } else { + render_mask(mask, colormap256) + }; image::imageops::overlay(&mut mask_img, &mask_dyn, x as i64, y as i64); let out_x = (col as u32 * w) as i64; @@ -122,7 +142,14 @@ impl Drawable for [Mask] { self.get_global_style(ctx), self.get_id(), ); - draw_masks(&masks_visible, style.colormap256(), canvas) + + draw_masks( + &masks_visible, + style.colormap256(), + canvas, + style.mask_cutout(), + style.mask_background_color().copied(), + ) } } @@ -175,7 +202,11 @@ impl Drawable for Mask { if style.visible() { let (w, h) = canvas.dimensions(); - let mask_dyn = render_mask(self, style.colormap256()); + let mask_dyn = if style.mask_cutout() { + apply_mask(canvas, self, style.mask_background_color().copied()) + } else { + render_mask(self, style.colormap256()) + }; let (mut out, mask_x, mask_y) = if w <= h { (RgbaImage::new(w * 2, h), w as i64, 0) diff --git a/src/viz/drawable/mod.rs b/src/viz/drawable/mod.rs index 8949982..611521f 100644 --- a/src/viz/drawable/mod.rs +++ b/src/viz/drawable/mod.rs @@ -60,8 +60,8 @@ pub trait Drawable { impl Drawable for Y { fn draw(&self, ctx: &DrawContext, canvas: &mut image::RgbaImage) -> anyhow::Result<()> { - if let Some(probs) = self.probs() { - probs.draw(ctx, canvas)?; + if let Some(masks) = self.masks() { + masks.draw(ctx, canvas)?; } if let Some(polygons) = self.polygons() { polygons.draw(ctx, canvas)?; @@ -78,8 +78,8 @@ impl Drawable for Y { if let Some(keypointss) = self.keypointss() { keypointss.draw(ctx, canvas)?; } - if let Some(masks) = self.masks() { - masks.draw(ctx, canvas)?; + if let Some(probs) = self.probs() { + probs.draw(ctx, canvas)?; } Ok(()) diff --git a/src/viz/styles.rs b/src/viz/styles.rs index beecef4..9f58f0b 100644 --- a/src/viz/styles.rs +++ b/src/viz/styles.rs @@ -4,25 +4,27 @@ use crate::{Color, ColorMap256, Skeleton}; #[derive(Debug, Clone, Builder, PartialEq)] pub struct Style { - visible: bool, // For ALL - text_visible: bool, // For ALL - draw_fill: bool, // For ALL - draw_outline: bool, // For ALL - color_fill_alpha: Option, // Alpha for fill - radius: usize, // For Keypoint - text_x_pos: f32, // For Probs - text_y_pos: f32, // For Probs - thickness: usize, // For Hbb - thickness_threshold: f32, // For Hbb - draw_mask_polygons: bool, // For Masks - draw_mask_polygon_largest: bool, // For Masks - draw_mask_hbbs: bool, // For Masks - draw_mask_obbs: bool, // For Masks - text_loc: TextLoc, // For ALL - color: StyleColors, // For ALL - palette: Vec, // For ALL - skeleton: Option, // For Keypoints - colormap256: Option, // For Masks + visible: bool, // For ALL + text_visible: bool, // For ALL + draw_fill: bool, // For ALL + draw_outline: bool, // For ALL + color_fill_alpha: Option, // Alpha for fill + radius: usize, // For Keypoint + text_x_pos: f32, // For Probs + text_y_pos: f32, // For Probs + thickness: usize, // For Hbb + thickness_threshold: f32, // For Hbb + draw_mask_polygons: bool, // For Masks + draw_mask_polygon_largest: bool, // For Masks + draw_mask_hbbs: bool, // For Masks + draw_mask_obbs: bool, // For Masks + mask_cutout: bool, // For Masks + mask_background_color: Option, // For Masks + text_loc: TextLoc, // For ALL + color: StyleColors, // For ALL + palette: Vec, // For ALL + skeleton: Option, // For Keypoints + colormap256: Option, // For Masks decimal_places: usize, #[args(set_pre = "show")] confidence: bool, @@ -45,6 +47,8 @@ impl Default for Style { draw_mask_polygon_largest: false, draw_mask_hbbs: false, draw_mask_obbs: false, + mask_cutout: false, + mask_background_color: None, radius: 3, text_x_pos: 0.05, text_y_pos: 0.05,