Enhance Annotator with mask-cutout functionality (#101)

Co-authored-by: jamjamjon <zhangjian@zhuofansoft.com>

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"] }

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",

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<Vec<Release>> {

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<Color>) -> 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<Color>,
 ) -> 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)

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(())

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<u8>,     // 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<Color>,              // For ALL
-    skeleton: Option<Skeleton>,       // For Keypoints
-    colormap256: Option<ColorMap256>, // 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<u8>,         // 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<Color>, // For Masks
+    text_loc: TextLoc,                    // For ALL
+    color: StyleColors,                   // For ALL
+    palette: Vec<Color>,                  // For ALL
+    skeleton: Option<Skeleton>,           // For Keypoints
+    colormap256: Option<ColorMap256>,     // 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,