feat: better text rendering

koharu-app/src/pipeline/engines/bubble_segmentation.rs

@@ -1,25 +1,13 @@
-//! Speech-bubble segmentation. Produces a `Mask { Bubble }` layer where
-//! pixels inside any detected speech balloon are `255` and everything else
-//! is `0`. The renderer consumes this mask to grow per-block text layout
-//! boxes inside the bubble — so English translations wrap into the full
-//! available space instead of being shrunk to fit the detector's
-//! (Japanese-shaped) source bbox.
-//!
-//! We build the mask from the model's **detection bboxes** (one per
-//! detected bubble), not from its per-pixel probability map. The model's
-//! bbox head is consistently accurate while its mask head under-segments
-//! bubble interiors — tight, text-shaped blobs that don't cover the full
-//! balloon. A text rect computed from the under-segmented mask would land
-//! wherever that small blob is (often off-centre in the actual bubble),
-//! which is exactly the "text outside the bubble" artefact. Bbox-filled
-//! rectangles trade a tiny bit of accuracy in the balloon's curved corners
-//! (handled by the renderer's inset) for guaranteed full coverage.
+//! Speech-bubble segmentation. Produces a `Mask { Bubble }` layer where each
+//! detected balloon contour is encoded as a distinct non-zero grayscale ID and
+//! everything else is `0`. The renderer consumes this ID mask for safe-area
+//! layout, placement, and glyph collision checks against the real bubble shape.
 
 use anyhow::Result;
 use async_trait::async_trait;
 use image::{DynamicImage, GrayImage, Luma};
 use koharu_core::{MaskRole, Op};
-use koharu_ml::speech_bubble_segmentation::SpeechBubbleSegmentation;
+use koharu_ml::speech_bubble_segmentation::{SpeechBubbleRegion, SpeechBubbleSegmentation};
 
 use crate::pipeline::artifacts::Artifact;
 use crate::pipeline::engine::{Engine, EngineCtx, EngineInfo};
@@ -33,36 +21,9 @@ impl Engine for Model {
         let image = load_source_image(ctx.scene, ctx.page, ctx.blobs)?;
         let result = self.0.inference(&image)?;
 
-        let w = result.image_width;
-        let h = result.image_height;
-        let mut mask = GrayImage::from_pixel(w, h, Luma([0u8]));
-
-        // Each detected bubble gets a unique non-zero grayscale ID in the
-        // mask. The renderer reads the ID under a text seed to recover
-        // exactly which bubble bbox that seed belongs to — no flood fill,
-        // no CC merging, no partition heuristics. Overlapping bboxes stay
-        // separable because smaller bubbles overwrite larger ones (painted
-        // last, after sorting by descending area), so a text seed inside
-        // a nested or embedded small bubble reads the small bubble's ID.
-        //
-        // Cap at 255 IDs; typical manga pages have well under 20 bubbles.
         let mut regions: Vec<_> = result.regions.iter().collect();
         regions.sort_by_key(|region| std::cmp::Reverse(region.area));
-        for (i, region) in regions.iter().take(255).enumerate() {
-            let id = (i + 1) as u8;
-            let x0 = region.bbox[0].floor().max(0.0) as u32;
-            let y0 = region.bbox[1].floor().max(0.0) as u32;
-            let x1 = (region.bbox[2].ceil().max(0.0) as u32).min(w);
-            let y1 = (region.bbox[3].ceil().max(0.0) as u32).min(h);
-            if x1 <= x0 || y1 <= y0 {
-                continue;
-            }
-            for y in y0..y1 {
-                for x in x0..x1 {
-                    mask.put_pixel(x, y, Luma([id]));
-                }
-            }
-        }
+        let mask = paint_bubble_id_mask(result.image_width, result.image_height, &regions);
 
         let blob = ctx.blobs.put_webp(&DynamicImage::ImageLuma8(mask))?;
         Ok(vec![upsert_mask_blob(
@@ -74,6 +35,35 @@ impl Engine for Model {
     }
 }
 
+fn paint_bubble_id_mask(width: u32, height: u32, regions: &[&SpeechBubbleRegion]) -> GrayImage {
+    let mut mask = GrayImage::from_pixel(width, height, Luma([0u8]));
+
+    // Cap at 255 IDs; typical manga pages have well under 20 bubbles.
+    // Larger bubbles are painted first so smaller overlapping contours keep
+    // their own ID when painted later.
+    for (i, region) in regions.iter().take(255).enumerate() {
+        if region.mask.is_empty() {
+            continue;
+        }
+        let id = (i + 1) as u8;
+        let src_width = region.mask.width as usize;
+        let max_x = region.mask.width.min(width.saturating_sub(region.mask.x));
+        let max_y = region.mask.height.min(height.saturating_sub(region.mask.y));
+        for local_y in 0..max_y {
+            let src_row = local_y as usize * src_width;
+            let y = region.mask.y + local_y;
+            for local_x in 0..max_x {
+                if region.mask.pixels[src_row + local_x as usize] == 0 {
+                    continue;
+                }
+                mask.put_pixel(region.mask.x + local_x, y, Luma([id]));
+            }
+        }
+    }
+
+    mask
+}
+
 inventory::submit! {
     EngineInfo {
         id: "speech-bubble-segmentation",
@@ -86,3 +76,35 @@ inventory::submit! {
         }),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use koharu_ml::speech_bubble_segmentation::{SpeechBubbleRegion, SpeechBubbleRegionMask};
+
+    #[test]
+    fn id_mask_paints_region_contours_not_bboxes() {
+        let region = SpeechBubbleRegion {
+            label_id: 0,
+            label: "bubble".to_string(),
+            score: 0.9,
+            bbox: [1.0, 1.0, 4.0, 4.0],
+            area: 3,
+            mask: SpeechBubbleRegionMask {
+                x: 1,
+                y: 1,
+                width: 3,
+                height: 3,
+                pixels: vec![0, 255, 0, 255, 255, 0, 0, 0, 0],
+            },
+        };
+
+        let mask = paint_bubble_id_mask(6, 6, &[&region]);
+
+        assert_eq!(mask.get_pixel(1, 1).0[0], 0);
+        assert_eq!(mask.get_pixel(2, 1).0[0], 1);
+        assert_eq!(mask.get_pixel(1, 2).0[0], 1);
+        assert_eq!(mask.get_pixel(2, 2).0[0], 1);
+        assert_eq!(mask.get_pixel(3, 2).0[0], 0);
+    }
+}

koharu-app/src/renderer.rs

@@ -252,7 +252,6 @@ impl Renderer {
             .text_align
             .map(core_align_to_renderer)
             .unwrap_or(RendererTextAlign::Center);
-        let seed_box = resolved_box.seed_box;
         let layout_box = resolved_box.layout_box;
 
         let mut layout_builder = TextLayout::new(&font, None)
@@ -285,7 +284,7 @@ impl Renderer {
                 },
             )?;
             let transform = centred_sprite_transform(
-                seed_box,
+                layout_box,
                 rendered.width(),
                 rendered.height(),
                 block.transform.rotation_deg,
@@ -300,7 +299,6 @@ impl Renderer {
             let candidate = fit_rendered_with_mask_collision(
                 &layout_builder,
                 translation,
-                writing_mode,
                 layout_box,
                 style.font_size,
                 min_font_size,
@@ -329,12 +327,6 @@ impl Renderer {
 
         let candidate = render_candidate(&layout)?;
 
-        // Place the sprite centred on the *seed* (detector's original
-        // text bbox). The seed is always positioned where the source
-        // language placed the text — inside the bubble body, never on
-        // the tail — so anchoring here keeps translations in the body
-        // even when the bubble bbox extends into the tail area.
-        //
         Ok(Some(RenderedBlock {
             node_id: block.node_id,
             sprite: DynamicImage::ImageRgba8(candidate.image),
@@ -370,8 +362,6 @@ impl Renderer {
 // ---------------------------------------------------------------------------
 
 const MASK_COLLISION_ALPHA_THRESHOLD: u8 = 8;
-const COLLISION_SQUEEZE_FACTOR: f32 = 0.90;
-const COLLISION_SQUEEZE_ATTEMPTS: usize = 3;
 const FIT_EPSILON: f32 = 0.5;
 
 struct RenderedTextCandidate {
@@ -379,9 +369,9 @@ struct RenderedTextCandidate {
     transform: Transform,
 }
 
-struct CollisionFitAttempt {
-    valid: Option<RenderedTextCandidate>,
-    fallback: Option<RenderedTextCandidate>,
+struct MaskCollisionAttempt {
+    candidate: RenderedTextCandidate,
+    valid: bool,
 }
 
 fn min_font_size_for_image(image_width: u32, image_height: u32) -> f32 {
@@ -457,7 +447,6 @@ fn fit_font_size<'a>(
 fn fit_rendered_with_mask_collision<'a, F>(
     layout_builder: &TextLayout<'a>,
     text: &str,
-    writing_mode: WritingMode,
     layout_box: LayoutBox,
     explicit_size: Option<f32>,
     min_size: f32,
@@ -469,54 +458,61 @@ fn fit_rendered_with_mask_collision<'a, F>(
 where
     F: FnMut(&LayoutRun<'a>) -> Result<RenderedTextCandidate>,
 {
-    let upper = explicit_size.unwrap_or(max_size).max(1.0).round() as i32;
-    let lower = min_size.max(1.0).round() as i32;
-    let min_size = if explicit_size.is_some() {
-        lower.min(upper)
-    } else {
-        lower
-    };
-    let max_size = upper.max(min_size);
+    if let Some(size) = explicit_size {
+        let attempt = render_mask_collision_attempt(
+            layout_builder,
+            text,
+            layout_box,
+            size.max(1.0),
+            mask,
+            bubble_id,
+            render_candidate,
+        )?;
+        return Ok(attempt.candidate);
+    }
+
+    let min_size = min_size.max(1.0).round() as i32;
+    let max_size = (max_size.max(1.0).round() as i32).max(min_size);
 
-    let min_attempt = try_mask_collision_size(
+    if let Some(candidate) = try_mask_collision_size(
+        layout_builder,
+        text,
+        layout_box,
+        max_size as f32,
+        mask,
+        bubble_id,
+        render_candidate,
+    )? {
+        return Ok(candidate);
+    }
+
+    let min_attempt = render_mask_collision_attempt(
         layout_builder,
         text,
-        writing_mode,
         layout_box,
         min_size as f32,
         mask,
         bubble_id,
         render_candidate,
     )?;
-    let Some(mut best) = min_attempt.valid else {
-        if let Some(fallback) = min_attempt.fallback {
-            return Ok(fallback);
-        }
-        return render_collision_fallback(
-            layout_builder,
-            text,
-            writing_mode,
-            layout_box,
-            min_size as f32,
-            render_candidate,
-        );
-    };
+    if !min_attempt.valid {
+        return Ok(min_attempt.candidate);
+    }
+    let mut best = min_attempt.candidate;
 
     let mut lo = min_size + 1;
-    let mut hi = max_size;
+    let mut hi = max_size - 1;
     while lo <= hi {
         let mid = lo + (hi - lo) / 2;
-        let attempt = try_mask_collision_size(
+        if let Some(candidate) = try_mask_collision_size(
             layout_builder,
             text,
-            writing_mode,
             layout_box,
             mid as f32,
             mask,
             bubble_id,
             render_candidate,
-        )?;
-        if let Some(candidate) = attempt.valid {
+        )? {
             best = candidate;
             lo = mid + 1;
         } else {
@@ -531,120 +527,71 @@ where
 fn try_mask_collision_size<'a, F>(
     layout_builder: &TextLayout<'a>,
     text: &str,
-    writing_mode: WritingMode,
     layout_box: LayoutBox,
     font_size: f32,
     mask: &GrayImage,
     bubble_id: u8,
     render_candidate: &mut F,
-) -> Result<CollisionFitAttempt>
+) -> Result<Option<RenderedTextCandidate>>
 where
     F: FnMut(&LayoutRun<'a>) -> Result<RenderedTextCandidate>,
 {
-    let mut fallback = None;
-    for main_extent in collision_squeeze_extents(layout_box, writing_mode) {
-        let layout = run_collision_layout_at(
-            layout_builder,
-            text,
-            writing_mode,
-            layout_box,
-            font_size,
-            main_extent,
-        )?;
-        if !layout_fits_collision_attempt(&layout, writing_mode, layout_box, main_extent) {
-            continue;
-        }
-
-        let candidate = render_candidate(&layout)?;
-        if sprite_collides_with_bubble_mask(&candidate.image, &candidate.transform, mask, bubble_id)
-        {
-            fallback = Some(candidate);
-            continue;
-        }
-        return Ok(CollisionFitAttempt {
-            valid: Some(candidate),
-            fallback,
-        });
+    let layout = run_collision_layout_at(layout_builder, text, layout_box, font_size)?;
+    let fits_layout_box = layout_fits_collision_attempt(&layout, layout_box);
+    if !fits_layout_box {
+        return Ok(None);
     }
 
-    Ok(CollisionFitAttempt {
-        valid: None,
-        fallback,
-    })
+    let candidate = render_candidate(&layout)?;
+    if sprite_collides_with_bubble_mask(&candidate.image, &candidate.transform, mask, bubble_id) {
+        return Ok(None);
+    }
+    Ok(Some(candidate))
 }
 
-fn render_collision_fallback<'a, F>(
+#[allow(clippy::too_many_arguments)]
+fn render_mask_collision_attempt<'a, F>(
     layout_builder: &TextLayout<'a>,
     text: &str,
-    writing_mode: WritingMode,
     layout_box: LayoutBox,
     font_size: f32,
+    mask: &GrayImage,
+    bubble_id: u8,
     render_candidate: &mut F,
-) -> Result<RenderedTextCandidate>
+) -> Result<MaskCollisionAttempt>
 where
     F: FnMut(&LayoutRun<'a>) -> Result<RenderedTextCandidate>,
 {
-    let layout = run_collision_layout_at(
-        layout_builder,
-        text,
-        writing_mode,
-        layout_box,
-        font_size,
-        primary_collision_extent(layout_box, writing_mode),
-    )?;
-    render_candidate(&layout)
+    let layout = run_collision_layout_at(layout_builder, text, layout_box, font_size)?;
+    let fits_layout_box = layout_fits_collision_attempt(&layout, layout_box);
+    let candidate = render_candidate(&layout)?;
+    let valid = fits_layout_box
+        && !sprite_collides_with_bubble_mask(
+            &candidate.image,
+            &candidate.transform,
+            mask,
+            bubble_id,
+        );
+    Ok(MaskCollisionAttempt { candidate, valid })
 }
 
 fn run_collision_layout_at<'a>(
     layout_builder: &TextLayout<'a>,
     text: &str,
-    writing_mode: WritingMode,
     layout_box: LayoutBox,
     font_size: f32,
-    main_extent: f32,
 ) -> Result<LayoutRun<'a>> {
-    let (max_width, max_height) = match writing_mode {
-        WritingMode::Horizontal => (main_extent, layout_box.height),
-        WritingMode::VerticalRl => (layout_box.width, main_extent),
-    };
     layout_builder
         .clone()
         .with_font_size(font_size.max(1.0))
-        .with_max_width(max_width.max(1.0))
-        .with_max_height(max_height.max(1.0))
+        .with_max_width(layout_box.width.max(1.0))
+        .with_max_height(layout_box.height.max(1.0))
         .run(text)
 }
 
-fn layout_fits_collision_attempt(
-    layout: &LayoutRun<'_>,
-    writing_mode: WritingMode,
-    layout_box: LayoutBox,
-    main_extent: f32,
-) -> bool {
-    let fits_box = layout.width <= layout_box.width + FIT_EPSILON
-        && layout.height <= layout_box.height + FIT_EPSILON;
-    let fits_main = match writing_mode {
-        WritingMode::Horizontal => layout.width <= main_extent + FIT_EPSILON,
-        WritingMode::VerticalRl => layout.height <= main_extent + FIT_EPSILON,
-    };
-    fits_box && fits_main
-}
-
-fn collision_squeeze_extents(layout_box: LayoutBox, writing_mode: WritingMode) -> Vec<f32> {
-    let mut extents = Vec::with_capacity(COLLISION_SQUEEZE_ATTEMPTS);
-    let mut extent = primary_collision_extent(layout_box, writing_mode);
-    for _ in 0..COLLISION_SQUEEZE_ATTEMPTS {
-        extents.push(extent.max(1.0));
-        extent *= COLLISION_SQUEEZE_FACTOR;
-    }
-    extents
-}
-
-fn primary_collision_extent(layout_box: LayoutBox, writing_mode: WritingMode) -> f32 {
-    match writing_mode {
-        WritingMode::Horizontal => layout_box.width,
-        WritingMode::VerticalRl => layout_box.height,
-    }
+fn layout_fits_collision_attempt(layout: &LayoutRun<'_>, layout_box: LayoutBox) -> bool {
+    layout.width <= layout_box.width + FIT_EPSILON
+        && layout.height <= layout_box.height + FIT_EPSILON
 }
 
 fn sprite_collides_with_bubble_mask(
@@ -722,9 +669,9 @@ fn resolve_layout_boxes(
         .into_iter()
         .map(|(seed_box, bubble_match)| match bubble_match {
             // Connected bubbles can contain multiple independently detected
-            // text blocks. Expanding all of them to the same bubble bbox makes
-            // their layouts collide, so shared bubbles fall back to each
-            // block's original detector box.
+            // text blocks. Expanding all of them to the same safe area makes
+            // their layouts collide, so shared bubbles keep each block's
+            // original detector box.
             Some(matched) if counts.get(&matched.id).copied().unwrap_or(0) == 1 => {
                 ResolvedLayoutBox {
                     seed_box,
@@ -934,8 +881,6 @@ fn centred_sprite_transform(
     sprite_height: u32,
     rotation_deg: f32,
 ) -> Transform {
-    // Centering on the anchor_box (usually the detector's seed_box)
-    // maintains original positioning and avoids tail drift.
     let sprite_w = sprite_width as f32;
     let sprite_h = sprite_height as f32;
     let cx = anchor_box.x + anchor_box.width * 0.5;
@@ -1062,7 +1007,53 @@ mod tests {
     }
 
     #[test]
-    fn shared_bubble_falls_back_to_seed_boxes() {
+    fn mask_collision_fit_renders_min_size_when_no_safe_size_exists() -> Result<()> {
+        let font = any_system_font();
+        let layout_builder = TextLayout::new(&font, None);
+        let layout_box = LayoutBox {
+            x: 0.0,
+            y: 0.0,
+            width: 24.0,
+            height: 12.0,
+        };
+        let mask = GrayImage::from_pixel(64, 64, Luma([0u8]));
+        let mut rendered_sizes = Vec::new();
+        let mut render_candidate = |layout: &LayoutRun<'_>| -> Result<RenderedTextCandidate> {
+            rendered_sizes.push(layout.font_size);
+            let width = layout.width.ceil().max(1.0) as u32;
+            let height = layout.height.ceil().max(1.0) as u32;
+            Ok(RenderedTextCandidate {
+                image: RgbaImage::from_pixel(width, height, Rgba([0, 0, 0, 255])),
+                transform: Transform {
+                    x: 0.0,
+                    y: 0.0,
+                    width: width as f32,
+                    height: height as f32,
+                    rotation_deg: 0.0,
+                },
+            })
+        };
+
+        let candidate = fit_rendered_with_mask_collision(
+            &layout_builder,
+            "overflowing text",
+            layout_box,
+            None,
+            12.0,
+            18.0,
+            &mask,
+            1,
+            &mut render_candidate,
+        )?;
+
+        assert_eq!(rendered_sizes.last().copied(), Some(12.0));
+        assert!(candidate.image.width() >= 1);
+        assert!(candidate.image.height() >= 1);
+        Ok(())
+    }
+
+    #[test]
+    fn shared_bubble_keeps_seed_boxes_to_avoid_overlap() {
         let mut mask = GrayImage::from_pixel(200, 200, Luma([0u8]));
         paint_rect(&mut mask, 10, 10, 190, 190, 1);
         let index = BubbleIndex::new(mask);
@@ -1139,25 +1130,6 @@ mod tests {
         ));
     }
 
-    #[test]
-    fn collision_squeeze_extents_retry_the_primary_axis() {
-        let layout_box = LayoutBox {
-            x: 0.0,
-            y: 0.0,
-            width: 100.0,
-            height: 50.0,
-        };
-
-        assert_eq!(
-            collision_squeeze_extents(layout_box, WritingMode::Horizontal),
-            vec![100.0, 90.0, 81.0]
-        );
-        assert_eq!(
-            collision_squeeze_extents(layout_box, WritingMode::VerticalRl),
-            vec![50.0, 45.0, 40.5]
-        );
-    }
-
     fn block(x: f32, y: f32, width: f32, height: f32, translation: &str) -> RenderBlockInput {
         RenderBlockInput {
             node_id: NodeId::new(),
@@ -1185,6 +1157,50 @@ mod tests {
         }
     }
 
+    fn any_system_font() -> Font {
+        let mut book = FontBook::new();
+        let preferred = [
+            "Yu Gothic",
+            "MS Gothic",
+            "Noto Sans CJK JP",
+            "Noto Sans",
+            "Arial",
+            "DejaVu Sans",
+            "Liberation Sans",
+        ];
+
+        for name in preferred {
+            if let Some(post_script_name) = book
+                .all_families()
+                .into_iter()
+                .find(|face| {
+                    face.post_script_name == name
+                        || face
+                            .families
+                            .iter()
+                            .any(|(family, _)| family.as_str() == name)
+                })
+                .map(|face| face.post_script_name)
+                .filter(|post_script_name| !post_script_name.is_empty())
+                && let Ok(font) = book.query(&post_script_name)
+            {
+                return font;
+            }
+        }
+
+        if let Some(face) = book
+            .all_families()
+            .into_iter()
+            .find(|face| !face.post_script_name.is_empty())
+        {
+            return book
+                .query(&face.post_script_name)
+                .expect("failed to load first system font");
+        }
+
+        panic!("no system font available for tests");
+    }
+
     #[test]
     fn centred_sprite_transform_anchors_to_provided_box_center() {
         let anchor = LayoutBox {

koharu-ml/src/speech_bubble_segmentation/mod.rs

@@ -68,6 +68,31 @@ pub struct SpeechBubbleSegmentationResult {
     pub probability_map: ProbabilityMap,
 }
 
+#[derive(Debug, Clone)]
+pub struct SpeechBubbleRegionMask {
+    pub x: u32,
+    pub y: u32,
+    pub width: u32,
+    pub height: u32,
+    pub pixels: Vec<u8>,
+}
+
+impl SpeechBubbleRegionMask {
+    pub fn empty(x: u32, y: u32) -> Self {
+        Self {
+            x,
+            y,
+            width: 0,
+            height: 0,
+            pixels: Vec::new(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.width == 0 || self.height == 0 || self.pixels.is_empty()
+    }
+}
+
 #[derive(Debug, Clone, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct SpeechBubbleRegion {
@@ -76,6 +101,8 @@ pub struct SpeechBubbleRegion {
     pub score: f32,
     pub bbox: [f32; 4],
     pub area: u32,
+    #[serde(skip_serializing)]
+    pub mask: SpeechBubbleRegionMask,
 }
 
 #[derive(Debug, Clone)]
@@ -332,18 +359,22 @@ fn postprocess(
 
     let mut regions = Vec::with_capacity(raw_regions.len());
     for (region, mask) in raw_regions.iter().zip(mask_probabilities.iter()) {
-        let area = merge_mask_into_probability_map(
+        let (area, region_mask) = extract_region_contour_mask(
             &mut probability_map,
             mask,
             region.bbox,
             config.mask_threshold,
-        );
+        )?;
+        if area == 0 {
+            continue;
+        }
         regions.push(SpeechBubbleRegion {
             label_id: region.label_id,
             label: region.label.clone(),
             score: region.score,
             bbox: region.bbox,
             area,
+            mask: region_mask,
         });
     }
 
@@ -519,16 +550,20 @@ fn mask_crop_window(
     (top, left, bottom, right)
 }
 
-fn merge_mask_into_probability_map(
+fn extract_region_contour_mask(
     probability_map: &mut ProbabilityMap,
     mask: &[f32],
     bbox: [f32; 4],
     threshold: f32,
-) -> u32 {
+) -> Result<(u32, SpeechBubbleRegionMask)> {
     let width = probability_map.width as usize;
     let height = probability_map.height as usize;
     if mask.len() != width * height {
-        return 0;
+        bail!(
+            "speech bubble mask length {} does not match image area {}",
+            mask.len(),
+            width * height
+        );
     }
 
     let x1 = bbox[0].floor().clamp(0.0, probability_map.width as f32) as usize;
@@ -536,29 +571,46 @@ fn merge_mask_into_probability_map(
     let x2 = bbox[2].ceil().clamp(0.0, probability_map.width as f32) as usize;
     let y2 = bbox[3].ceil().clamp(0.0, probability_map.height as f32) as usize;
     if x2 <= x1 || y2 <= y1 {
-        return 0;
+        return Ok((0, SpeechBubbleRegionMask::empty(x1 as u32, y1 as u32)));
     }
 
+    let mask_width = x2 - x1;
+    let mask_height = y2 - y1;
+    let mut pixels = vec![0u8; mask_width * mask_height];
     let mut area = 0u32;
     for y in y1..y2.min(height) {
         let row_offset = y * width;
+        let local_row_offset = (y - y1) * mask_width;
         for x in x1..x2.min(width) {
             let idx = row_offset + x;
             let value = mask[idx];
             if value >= threshold {
                 area += 1;
+                pixels[local_row_offset + (x - x1)] = u8::MAX;
             }
             if value > probability_map.values[idx] {
                 probability_map.values[idx] = value;
             }
         }
     }
-    area
+    Ok((
+        area,
+        SpeechBubbleRegionMask {
+            x: x1 as u32,
+            y: y1 as u32,
+            width: mask_width as u32,
+            height: mask_height as u32,
+            pixels,
+        },
+    ))
 }
 
 #[cfg(test)]
 mod tests {
-    use super::{PreparedInput, map_bbox_to_original, mask_crop_window};
+    use super::{
+        PreparedInput, extract_region_contour_mask, map_bbox_to_original, mask_crop_window,
+    };
+    use crate::probability_map::ProbabilityMap;
     use candle_core::{DType, Device, Tensor};
 
     #[test]
@@ -587,4 +639,27 @@ mod tests {
         let (top, left, bottom, right) = mask_crop_window(1000, 500, 160, 160);
         assert_eq!((top, left, bottom, right), (40, 0, 120, 160));
     }
+
+    #[test]
+    fn extract_region_contour_mask_keeps_thresholded_shape() -> anyhow::Result<()> {
+        let mut probability_map = ProbabilityMap::zeros(6, 5);
+        let mut mask = vec![0.0f32; 6 * 5];
+        mask[1 + 1 * 6] = 0.9;
+        mask[2 + 1 * 6] = 0.8;
+        mask[2 + 2 * 6] = 0.7;
+        mask[4 + 3 * 6] = 0.4;
+
+        let (area, region_mask) =
+            extract_region_contour_mask(&mut probability_map, &mask, [1.0, 1.0, 5.0, 4.0], 0.5)?;
+
+        assert_eq!(area, 3);
+        assert_eq!((region_mask.x, region_mask.y), (1, 1));
+        assert_eq!((region_mask.width, region_mask.height), (4, 3));
+        assert_eq!(region_mask.pixels[0], u8::MAX);
+        assert_eq!(region_mask.pixels[1], u8::MAX);
+        assert_eq!(region_mask.pixels[5], u8::MAX);
+        assert_eq!(region_mask.pixels[11], 0);
+        assert_eq!(probability_map.values[4 + 3 * 6], 0.4);
+        Ok(())
+    }
 }

koharu-renderer/src/layout.rs

@@ -19,6 +19,8 @@ pub use crate::segment::{LineBreakSuffix, hyphenation_lang_from_tag};
 pub use crate::shape::{PositionedGlyph, ShapedRun, ShapingOptions, TextShaper};
 
 const HYPHENATION_MIN_WORD_LEN: usize = 8;
+const LINE_BREAK_HYPHEN_PENALTY: f32 = 2_000.0;
+const LINE_BREAK_OVERFLOW_MULTIPLIER: f32 = 10_000.0;
 
 /// Writing mode for text layout.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -65,6 +67,34 @@ pub struct LayoutRun<'a> {
     pub font_size: f32,
 }
 
+#[derive(Clone)]
+struct LineRun<'a> {
+    shaped: ShapedRun<'a>,
+    level: unicode_bidi::Level,
+}
+
+#[derive(Clone)]
+struct ShapedBreakSuffix<'a> {
+    runs: Vec<LineRun<'a>>,
+    advance: f32,
+}
+
+#[derive(Clone)]
+struct ShapedSegment<'a> {
+    range: Range<usize>,
+    next_offset: usize,
+    is_mandatory: bool,
+    runs: Vec<LineRun<'a>>,
+    advance: f32,
+    break_suffix: Option<ShapedBreakSuffix<'a>>,
+}
+
+#[derive(Clone, Copy, Debug)]
+struct LineBreakMeasure {
+    advance: f32,
+    break_suffix_advance: f32,
+}
+
 #[derive(Clone)]
 pub struct TextLayout<'a> {
     writing_mode: WritingMode,
@@ -236,92 +266,9 @@ impl<'a> TextLayout<'a> {
         .unwrap_or(f32::INFINITY);
         let max_extent_finite = max_extent.is_finite() && max_extent > 0.0;
 
-        let segments = line_breaker.line_segments(text);
-
         let mut fonts: Vec<&Font> = Vec::with_capacity(1 + self.fallback_fonts.len());
         fonts.push(self.font);
         fonts.extend(self.fallback_fonts.iter());
-        let mut lines: Vec<LayoutLine<'a>> = Vec::new();
-
-        // Temporary storage for runs in the current line to be reordered
-        struct LineRun<'a> {
-            shaped: ShapedRun<'a>,
-            level: unicode_bidi::Level,
-        }
-        struct ShapedBreakSuffix<'a> {
-            runs: Vec<LineRun<'a>>,
-            advance: f32,
-        }
-        let mut current_line_runs: Vec<LineRun<'a>> = Vec::new();
-        let mut line_offset = 0usize;
-        let mut current_advance = 0.0f32;
-        let mut current_break_suffix: Option<ShapedBreakSuffix<'a>> = None;
-
-        let finalize_current_line = |runs: &mut Vec<LineRun<'a>>,
-                                     offset: &mut usize,
-                                     visible_end: usize,
-                                     next_offset: usize,
-                                     lines: &mut Vec<LayoutLine<'a>>,
-                                     break_suffix: Option<ShapedBreakSuffix<'a>>,
-                                     force_push: bool| {
-            if runs.is_empty() && !force_push {
-                *offset = next_offset;
-                return;
-            }
-
-            if let Some(mut suffix) = break_suffix {
-                runs.append(&mut suffix.runs);
-            }
-
-            let levels: Vec<unicode_bidi::Level> = runs.iter().map(|r| r.level).collect();
-            let visual_indices = reorder_visual(&levels);
-
-            let mut line = LayoutLine {
-                range: *offset..visible_end,
-                direction: if self.writing_mode.is_vertical() {
-                    harfrust::Direction::TopToBottom
-                } else {
-                    bidi_info
-                        .paragraphs
-                        .iter()
-                        .find(|p| *offset >= p.range.start && *offset <= p.range.end)
-                        .map(|p| {
-                            if p.level.is_rtl() {
-                                harfrust::Direction::RightToLeft
-                            } else {
-                                harfrust::Direction::LeftToRight
-                            }
-                        })
-                        .unwrap_or(harfrust::Direction::LeftToRight)
-                },
-                ..Default::default()
-            };
-
-            let mut pen_x = 0.0f32;
-            let mut pen_y = 0.0f32;
-
-            for idx in visual_indices {
-                let run = &mut runs[idx];
-                for glyph in std::mem::take(&mut run.shaped.glyphs) {
-                    line.glyphs.push(glyph);
-                }
-                if self.writing_mode.is_vertical() {
-                    pen_y -= run.shaped.y_advance;
-                } else {
-                    pen_x += run.shaped.x_advance;
-                }
-            }
-
-            line.advance = if self.writing_mode.is_vertical() {
-                pen_y.abs()
-            } else {
-                pen_x
-            };
-
-            lines.push(line);
-            runs.clear();
-            *offset = next_offset;
-        };
 
         let shape_break_suffix = |suffix: LineBreakSuffix,
                                   level: unicode_bidi::Level,
@@ -340,14 +287,15 @@ impl<'a> TextLayout<'a> {
                 for glyph in &mut shaped.glyphs {
                     glyph.cluster += cluster as u32;
                 }
-                advance += shaped.x_advance;
+                advance += shaped.x_advance.abs();
                 runs.push(LineRun { shaped, level });
             }
 
             Ok(ShapedBreakSuffix { runs, advance })
         };
 
-        for segment in segments {
+        let mut shaped_segments = Vec::new();
+        for segment in line_breaker.line_segments(text) {
             let segment_text = &text[segment.range.clone()];
 
             let mut segment_runs = Vec::new();
@@ -397,9 +345,9 @@ impl<'a> TextLayout<'a> {
                         }
 
                         segment_advance += if self.writing_mode.is_vertical() {
-                            shaped.y_advance
+                            shaped.y_advance.abs()
                         } else {
-                            shaped.x_advance
+                            shaped.x_advance.abs()
                         };
 
                         segment_runs.push(LineRun { shaped, level });
@@ -414,62 +362,45 @@ impl<'a> TextLayout<'a> {
                 None
             };
 
-            let would_overflow = if self.writing_mode.is_vertical() {
-                current_advance.abs() + segment_advance.abs() > max_extent
-            } else {
-                let break_suffix_advance = current_break_suffix
-                    .as_ref()
-                    .map_or(0.0, |suffix| suffix.advance);
-                current_advance + break_suffix_advance + segment_advance > max_extent
-            };
-
-            if would_overflow && !current_line_runs.is_empty() {
-                finalize_current_line(
-                    &mut current_line_runs,
-                    &mut line_offset,
-                    segment.range.start,
-                    segment.range.start,
-                    &mut lines,
-                    current_break_suffix.take(),
-                    false,
-                );
-                current_advance = 0.0;
-            }
+            shaped_segments.push(ShapedSegment {
+                range: segment.range,
+                next_offset: segment.next_offset,
+                is_mandatory: segment.is_mandatory,
+                runs: segment_runs,
+                advance: segment_advance,
+                break_suffix: segment_break_suffix,
+            });
+        }
 
-            current_line_runs.extend(segment_runs);
-            current_advance += segment_advance;
-            current_break_suffix = segment_break_suffix;
-
-            if segment.is_mandatory {
-                // Consecutive mandatory breaks (e.g. "A\n\nB") will drop the empty line:
-                // when segment.is_mandatory and current_line_runs is empty, finalize_current_line
-                // early-returns without pushing a LayoutLine. Consider special-casing mandatory
-                // breaks to push an empty LayoutLine (advance=0, glyphs empty, range=offset..visible_end)
-                // so blank lines are preserved.
-                finalize_current_line(
-                    &mut current_line_runs,
-                    &mut line_offset,
-                    segment.range.end,
-                    segment.next_offset,
-                    &mut lines,
-                    None,
-                    true,
-                );
-                current_advance = 0.0;
-                current_break_suffix = None;
+        let mut lines: Vec<LayoutLine<'a>> = Vec::new();
+        let mut line_offset = 0usize;
+        let mut paragraph_start = 0usize;
+        for (index, segment) in shaped_segments.iter().enumerate() {
+            if !segment.is_mandatory {
+                continue;
             }
+            self.append_balanced_segment_lines(
+                &shaped_segments[paragraph_start..=index],
+                &mut line_offset,
+                segment.next_offset,
+                true,
+                max_extent,
+                &bidi_info,
+                &mut lines,
+            );
+            paragraph_start = index + 1;
+        }
+        if paragraph_start < shaped_segments.len() {
+            self.append_balanced_segment_lines(
+                &shaped_segments[paragraph_start..],
+                &mut line_offset,
+                text.len(),
+                false,
+                max_extent,
+                &bidi_info,
+                &mut lines,
+            );
         }
-
-        // Finalize last line
-        finalize_current_line(
-            &mut current_line_runs,
-            &mut line_offset,
-            text.len(),
-            text.len(),
-            &mut lines,
-            None,
-            false, // Don't force push a final empty line if the text didn't end with a break
-        );
 
         // Baselines depend only on line index and metrics. For vertical text we compute absolute X
         // positions within the layout bounds (0..width) so the renderer can draw from the left.
@@ -603,6 +534,153 @@ impl<'a> TextLayout<'a> {
         })
     }
 
+    #[allow(clippy::too_many_arguments)]
+    fn append_balanced_segment_lines(
+        &self,
+        segments: &[ShapedSegment<'a>],
+        line_offset: &mut usize,
+        final_next_offset: usize,
+        force_final_line: bool,
+        max_extent: f32,
+        bidi_info: &BidiInfo<'_>,
+        lines: &mut Vec<LayoutLine<'a>>,
+    ) {
+        if segments.is_empty() {
+            if force_final_line {
+                *line_offset = self.push_layout_line(
+                    Vec::new(),
+                    *line_offset,
+                    *line_offset,
+                    final_next_offset,
+                    None,
+                    true,
+                    bidi_info,
+                    lines,
+                );
+            }
+            return;
+        }
+
+        let break_indices = if max_extent.is_finite() && max_extent > 0.0 {
+            let measures = segments
+                .iter()
+                .map(|segment| LineBreakMeasure {
+                    advance: segment.advance,
+                    break_suffix_advance: segment
+                        .break_suffix
+                        .as_ref()
+                        .map_or(0.0, |suffix| suffix.advance),
+                })
+                .collect::<Vec<_>>();
+            optimal_line_breaks(&measures, max_extent)
+        } else {
+            vec![segments.len()]
+        };
+
+        let mut start = 0usize;
+        for end in break_indices {
+            if end <= start || end > segments.len() {
+                continue;
+            }
+            let final_line = end == segments.len();
+            let visible_end = segments[end - 1].range.end;
+            let next_offset = if final_line {
+                final_next_offset
+            } else {
+                segments[end].range.start
+            };
+            let break_suffix = if final_line {
+                None
+            } else {
+                segments[end - 1].break_suffix.clone()
+            };
+            let runs = segments[start..end]
+                .iter()
+                .flat_map(|segment| segment.runs.iter().cloned())
+                .collect::<Vec<_>>();
+            *line_offset = self.push_layout_line(
+                runs,
+                *line_offset,
+                visible_end,
+                next_offset,
+                break_suffix,
+                force_final_line && final_line,
+                bidi_info,
+                lines,
+            );
+            start = end;
+        }
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn push_layout_line(
+        &self,
+        mut runs: Vec<LineRun<'a>>,
+        offset: usize,
+        visible_end: usize,
+        next_offset: usize,
+        break_suffix: Option<ShapedBreakSuffix<'a>>,
+        force_push: bool,
+        bidi_info: &BidiInfo<'_>,
+        lines: &mut Vec<LayoutLine<'a>>,
+    ) -> usize {
+        if runs.is_empty() && !force_push {
+            return next_offset;
+        }
+
+        if let Some(mut suffix) = break_suffix {
+            runs.append(&mut suffix.runs);
+        }
+
+        let levels: Vec<unicode_bidi::Level> = runs.iter().map(|r| r.level).collect();
+        let visual_indices = reorder_visual(&levels);
+
+        let mut line = LayoutLine {
+            range: offset..visible_end,
+            direction: if self.writing_mode.is_vertical() {
+                harfrust::Direction::TopToBottom
+            } else {
+                bidi_info
+                    .paragraphs
+                    .iter()
+                    .find(|p| offset >= p.range.start && offset <= p.range.end)
+                    .map(|p| {
+                        if p.level.is_rtl() {
+                            harfrust::Direction::RightToLeft
+                        } else {
+                            harfrust::Direction::LeftToRight
+                        }
+                    })
+                    .unwrap_or(harfrust::Direction::LeftToRight)
+            },
+            ..Default::default()
+        };
+
+        let mut pen_x = 0.0f32;
+        let mut pen_y = 0.0f32;
+
+        for idx in visual_indices {
+            let run = &mut runs[idx];
+            for glyph in std::mem::take(&mut run.shaped.glyphs) {
+                line.glyphs.push(glyph);
+            }
+            if self.writing_mode.is_vertical() {
+                pen_y -= run.shaped.y_advance;
+            } else {
+                pen_x += run.shaped.x_advance;
+            }
+        }
+
+        line.advance = if self.writing_mode.is_vertical() {
+            pen_y.abs()
+        } else {
+            pen_x
+        };
+
+        lines.push(line);
+        next_offset
+    }
+
     fn ink_bounds(&self, font_size: f32, lines: &[LayoutLine<'a>]) -> Option<(f32, f32, f32, f32)> {
         let mut metrics_cache = HashMap::new();
 
@@ -699,6 +777,74 @@ impl<'a> TextLayout<'a> {
     }
 }
 
+fn optimal_line_breaks(segments: &[LineBreakMeasure], max_extent: f32) -> Vec<usize> {
+    let len = segments.len();
+    if len == 0 {
+        return Vec::new();
+    }
+    if !max_extent.is_finite() || max_extent <= 0.0 {
+        return vec![len];
+    }
+
+    let mut dp = vec![f32::INFINITY; len + 1];
+    let mut prev = vec![None; len + 1];
+    dp[0] = 0.0;
+
+    for start in 0..len {
+        if !dp[start].is_finite() {
+            continue;
+        }
+        let mut advance = 0.0f32;
+        for end in start + 1..=len {
+            advance += segments[end - 1].advance;
+            let suffix_advance = if end < len {
+                segments[end - 1].break_suffix_advance
+            } else {
+                0.0
+            };
+            let line_advance = advance + suffix_advance;
+            let is_single_segment = end == start + 1;
+            if line_advance > max_extent && !is_single_segment {
+                break;
+            }
+
+            let mut cost = dp[start] + line_break_badness(line_advance, max_extent);
+            if end < len && suffix_advance > 0.0 {
+                cost += LINE_BREAK_HYPHEN_PENALTY;
+            }
+
+            if cost < dp[end] {
+                dp[end] = cost;
+                prev[end] = Some(start);
+            }
+        }
+    }
+
+    if !dp[len].is_finite() {
+        return vec![len];
+    }
+
+    let mut breaks = Vec::new();
+    let mut index = len;
+    while index > 0 {
+        breaks.push(index);
+        let Some(previous) = prev[index] else {
+            return vec![len];
+        };
+        index = previous;
+    }
+    breaks.reverse();
+    breaks
+}
+
+fn line_break_badness(line_advance: f32, max_extent: f32) -> f32 {
+    if line_advance <= max_extent {
+        (max_extent - line_advance).powi(3)
+    } else {
+        (line_advance - max_extent).powi(3) * LINE_BREAK_OVERFLOW_MULTIPLIER
+    }
+}
+
 fn centered_x_offset(x_min: f32, x_max: f32) -> f32 {
     -((x_min + x_max) * 0.5)
 }
@@ -894,6 +1040,19 @@ mod tests {
         );
     }
 
+    #[test]
+    fn optimal_line_breaks_balance_ragged_lines() {
+        let segments = vec![
+            LineBreakMeasure {
+                advance: 30.0,
+                break_suffix_advance: 0.0,
+            };
+            7
+        ];
+
+        assert_eq!(optimal_line_breaks(&segments, 100.0), vec![2, 4, 7]);
+    }
+
     #[test]
     fn layout_baselines_horizontal_follow_font_metrics() -> anyhow::Result<()> {
         let font = any_system_font();