Update app.py

app.py

@@ -1,10 +1,8 @@
-import os
-import re
+import os, re, io
 import xml.etree.ElementTree as ET
 import gradio as gr
 
-# Your handwriting model wrapper
-from hand import Hand
+from hand import Hand  # your model
 
 # -----------------------------------------------------------------------------
 # Setup
@@ -13,169 +11,224 @@ os.makedirs("img", exist_ok=True)
 hand = Hand()
 
 # -----------------------------------------------------------------------------
-# SVG helpers
+# Small helpers
 # -----------------------------------------------------------------------------
-def _extract_paths(svg_root_or_group):
-    """Return a list of <path> elements (namespace-agnostic)."""
-    return [el for el in svg_root_or_group.iter() if el.tag.endswith('path')]
+def _parse_viewbox(root):
+    vb = root.get("viewBox")
+    if not vb:
+        return (0.0, 0.0, 1200.0, 400.0)
+    x, y, w, h = map(float, vb.split())
+    return (x, y, w, h)
 
-def _bbox_of_paths(paths):
-    """
-    Conservative bbox by scanning numeric coords inside each path 'd'.
-    Not mathematically perfect for curve extrema, but close enough for layout.
-    """
-    xs, ys = [], []
-    num_re = re.compile(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?')
-    for p in paths:
-        d = p.get('d', '')
-        nums = list(map(float, num_re.findall(d)))
-        for i in range(0, len(nums) - 1, 2):
-            xs.append(nums[i])
-            ys.append(nums[i + 1])
-    if not xs or not ys:
-        return (0.0, 0.0, 0.0, 0.0)
-    return (min(xs), min(ys), max(xs), max(ys))
+def _extract_paths(elem):
+    return [e for e in elem.iter() if e.tag.endswith("path")]
 
 def _translate_group(elem, dx, dy):
-    prev = elem.get('transform', '')
-    t = f"translate({dx},{dy})"
-    elem.set('transform', (prev + " " + t).strip())
-
-# -----------------------------------------------------------------------------
-# Visual (raster) measurement for exact spacing
-# -----------------------------------------------------------------------------
-def _tight_width_via_raster(group_elem, stroke_color="#000", stroke_width=2):
-    """
-    Render a tiny SVG containing only this group, rasterize to PNG in-memory,
-    and return the non-transparent width in pixels. Produces pixel-accurate
-    width that's immune to relative path quirks.
-    """
-    import io
-    from PIL import Image
+    prev = elem.get("transform", "")
+    elem.set("transform", (prev + f" translate({dx},{dy})").strip())
+
+# Tokenize a line into runs of text / spaces / underscores (order matters)
+def _tokenize_line(line):
+    tokens = []
+    i, n = 0, len(line)
+    while i < n:
+        ch = line[i]
+        if ch == '_':
+            j = i
+            while j < n and line[j] == '_': j += 1
+            tokens.append(("sep_underscore", line[i:j]))
+            i = j
+        elif ch.isspace():
+            j = i
+            while j < n and line[j].isspace(): j += 1
+            tokens.append(("sep_space", line[i:j]))
+            i = j
+        else:
+            j = i
+            while j < n and (line[j] != '_' and not line[j].isspace()):
+                j += 1
+            tokens.append(("text", line[i:j]))
+            i = j
+    return tokens
+
+# Build the text that the model will actually render (underscores & spaces -> space)
+def _display_text_from_tokens(tokens):
+    out = []
+    for t, v in tokens:
+        if t == "text":
+            out.append(v)
+        else:
+            out.append(" ")  # normalize separators
+    return "".join(out)
+
+# Rasterize an SVG string to RGBA PIL.Image and return (img, vw, vh, vx, vy)
+def _rasterize_svg(svg_str, scale=2.0):
     import cairosvg
-    import xml.etree.ElementTree as ET
-
-    # Build a tiny standalone SVG for the group (children are already paths)
-    svg = ET.Element('svg', {
-        'xmlns': 'http://www.w3.org/2000/svg',
-        'viewBox': '0 0 1200 400'
-    })
-    g = ET.Element('g')
-    for p in list(group_elem):
-        g.append(p)
-    svg.append(g)
-
-    svg_bytes = ET.tostring(svg, encoding='utf-8')
-
-    # Rasterize in memory
-    png_bytes = cairosvg.svg2png(bytestring=svg_bytes, scale=1.0, background_color="none")
-    img = Image.open(io.BytesIO(png_bytes)).convert('RGBA')
-
-    # Non-transparent bounding box
-    alpha = img.split()[-1]
-    bbox = alpha.getbbox()  # (left, top, right, bottom) or None
+    from PIL import Image
+    png = cairosvg.svg2png(bytestring=svg_str.encode("utf-8"), scale=scale, background_color="none")
+    img = Image.open(io.BytesIO(png)).convert("RGBA")
+    return img
+
+# From alpha image, find left-to-right blobs & gaps within the content bbox
+def _find_blobs_and_gaps(alpha_img):
+    w, h = alpha_img.size
+    bbox = alpha_img.getbbox()
     if not bbox:
-        return 0.0
+        return [], [], (0, 0, w, h)  # nothing
     left, top, right, bottom = bbox
-    return float(right - left)
+
+    # columns with any ink
+    def col_has_ink(x):
+        # check a cropped band to speed up
+        for y in range(top, bottom):
+            if alpha_img.getpixel((x, y)) > 0:
+                return True
+        return False
+
+    blobs, gaps = [], []
+    x = left
+    in_blob = col_has_ink(x)
+    start = x
+    while x < right:
+        has = col_has_ink(x)
+        if has != in_blob:
+            if in_blob:
+                blobs.append((start, x))  # [start, x)
+            else:
+                gaps.append((start, x))
+            start = x
+            in_blob = has
+        x += 1
+    # close last run
+    if in_blob:
+        blobs.append((start, right))
+    else:
+        gaps.append((start, right))
+
+    # We only want gaps **between** blobs, not leading/trailing margins:
+    core_gaps = []
+    for i in range(len(blobs) - 1):
+        core_gaps.append((blobs[i][1], blobs[i + 1][0]))
+
+    return blobs, core_gaps, (left, top, right, bottom)
+
+# Map pixel x/y to SVG coords via viewBox
+def _px_to_svg_x(x_px, img_w, vb):
+    vx, vy, vw, vh = vb
+    return vx + (x_px / float(img_w)) * vw
+
+def _px_to_svg_y(y_px, img_h, vb):
+    vx, vy, vw, vh = vb
+    return vy + (y_px / float(img_h)) * vh
+
+# Draw N underscores filling a gap (pixel coords -> converted to SVG coords)
+def _draw_underscores_in_gap(root, gap_px, baseline_px, img_w, img_h, vb,
+                             color, stroke_width, n, pad_px=3, between_px=4):
+    gap_w = max(0, gap_px[1] - gap_px[0] - 2 * pad_px)
+    if gap_w <= 6 or n <= 0:
+        return
+    # fit underscores nicely within the gap
+    # give each underscore ~85% of its slot
+    slot = gap_w / n
+    line_len = max(8, int(slot * 0.85) - between_px)
+    offset = (slot - line_len) / 2.0
+
+    for i in range(n):
+        x0_px = gap_px[0] + pad_px + i * slot + offset
+        x1_px = x0_px + line_len
+        y_px = baseline_px
+
+        x0 = _px_to_svg_x(x0_px, img_w, vb)
+        x1 = _px_to_svg_x(x1_px, img_w, vb)
+        y  = _px_to_svg_y(y_px,  img_h, vb)
+
+        p = ET.Element("path")
+        p.set("d", f"M{x0},{y} L{x1},{y}")
+        p.set("stroke", color)
+        p.set("stroke-width", str(max(1, stroke_width)))
+        p.set("fill", "none")
+        p.set("stroke-linecap", "round")
+        root.append(p)
 
 # -----------------------------------------------------------------------------
-# Segment compositor (key fix for underscores)
+# Render ONE line with correct underscores by analyzing the full-line mask
 # -----------------------------------------------------------------------------
-def _render_part_group(text_part, style, bias, color, stroke_width):
+def render_line_svg_with_underscores(line, style, bias, color, stroke_width):
     """
-    Render one underscore-free 'part' using the Hand model to a temp SVG,
-    extract its path group, and return (group, path_bbox).
+    Returns (line_group, width_estimate). The group contains model strokes + our underscores.
     """
-    # Avoid empty part causing errors: render a single space
-    lines = [text_part if text_part else " "]
+    # 1) Tokenize and build the display line (underscores & spaces -> spaces)
+    tokens = _tokenize_line(line)
+    display_line = _display_text_from_tokens(tokens).replace("/", "-").replace("\\", "-")
+
+    # 2) Ask the model to render this single line to a temp SVG
     hand.write(
-        filename='img/part.tmp.svg',
-        lines=lines,
+        filename="img/line.tmp.svg",
+        lines=[display_line if display_line.strip() else " "],
         biases=[bias],
         styles=[style],
         stroke_colors=[color],
         stroke_widths=[stroke_width]
     )
-    svg = ET.parse('img/part.tmp.svg').getroot()
-    g = ET.Element('g')
-    for p in _extract_paths(svg):
+    root = ET.parse("img/line.tmp.svg").getroot()
+    vb = _parse_viewbox(root)
+
+    # 3) Rasterize the exact SVG we will augment, then find blobs/gaps
+    from PIL import Image
+    img = _rasterize_svg(ET.tostring(root, encoding="unicode"))
+    alpha = img.split()[-1]
+    blobs, gaps, content_bbox = _find_blobs_and_gaps(alpha)
+    img_w, img_h = img.size
+    left, top, right, bottom = content_bbox
+    line_h_px = max(20, bottom - top)
+    baseline_px = bottom - int(0.18 * line_h_px)  # tuck slightly above bottom
+
+    # 4) Build a <g> that contains all original paths
+    g = ET.Element("g")
+    for p in _extract_paths(root):
         g.append(p)
-    bbox = _bbox_of_paths(_extract_paths(g))
-    return g, bbox
 
-def render_line_with_underscores(
-    line,
-    style,
-    bias,
-    color,
-    stroke_width,
-    line_y,
-    x_start=40,
-    gap_left=8,
-    gap_right=8,
-    underscore_len=22
-):
-    """
-    Compose a line with underscores by rendering surrounding segments,
-    measuring their *visual* widths via rasterization, and drawing the
-    underscore at the baseline. Returns (list_of_elements, rightmost_x).
-    """
-    parts = line.split('_')
-
-    part_groups, visual_widths, path_bboxes = [], [], []
-
-    # 1) Render each part to a <g>
-    for part in parts:
-        g, bbox = _render_part_group(part, style, bias, color, stroke_width)
-        part_groups.append(g)
-        path_bboxes.append(bbox)
-
-    # 2) Measure visual widths robustly (pixel-accurate)
-    for g in part_groups:
-        w = _tight_width_via_raster(g, stroke_color=color, stroke_width=stroke_width)
-        if w <= 0.1:  # fallback to path bbox width if needed
-            minx, _, maxx, _ = _bbox_of_paths(_extract_paths(g))
-            w = max(0.0, maxx - minx)
-        visual_widths.append(w)
-
-    # 3) Estimate line metrics for baseline
-    line_height = 40.0
-    for minx, miny, maxx, maxy in path_bboxes:
-        if maxy - miny > 0.0:
-            line_height = max(20.0, maxy - miny)
-            break
-    base_y = line_y + 0.8 * line_height
-
-    # 4) Compose all elements left-to-right using measured widths
-    composed = []
-    cursor_x = x_start
-
-    for i, (g, bbox, visw) in enumerate(zip(part_groups, path_bboxes, visual_widths)):
-        minx, _, maxx, _ = bbox
-
-        # Align this group's left edge at cursor_x
-        _translate_group(g, dx=(cursor_x - minx), dy=line_y)
-        composed.append(g)
-        cursor_x += visw  # advance by visual width
-
-        # Draw underscore between parts
-        if i < len(part_groups) - 1:
-            x0 = cursor_x + gap_left
-            x1 = x0 + underscore_len
-            underscore = ET.Element('path')
-            underscore.set('d', f"M{x0},{base_y} L{x1},{base_y}")
-            underscore.set('stroke', color)
-            underscore.set('stroke-width', str(max(1, stroke_width)))
-            underscore.set('fill', 'none')
-            underscore.set('stroke-linecap', 'round')
-            composed.append(underscore)
-            cursor_x = x1 + gap_right
-
-    return composed, cursor_x
+    # 5) Determine which gaps correspond to underscores (not regular spaces)
+    # Walk tokens: whenever we have TEXT then SEPs then TEXT, consume one gap.
+    gap_idx = 0
+    i = 0
+    text_run_count = sum(1 for t, v in tokens if t == "text" and len(v) > 0)
+    # Defensive: if model merged two words, blobs may be fewer—clamp at min
+    max_gaps_we_can_use = max(0, min(len(gaps), max(0, text_run_count - 1)))
+
+    while i < len(tokens):
+        t, v = tokens[i]
+        if t == "text" and len(v) > 0:
+            # Count any following separators as one logical gap between this and the next text
+            j = i + 1
+            underscore_count = 0
+            saw_sep = False
+            while j < len(tokens) and tokens[j][0].startswith("sep_"):
+                saw_sep = True
+                if tokens[j][0] == "sep_underscore":
+                    underscore_count += len(tokens[j][1])
+                j += 1
+
+            if saw_sep and gap_idx < max_gaps_we_can_use:
+                if underscore_count > 0:
+                    _draw_underscores_in_gap(
+                        g, gaps[gap_idx], baseline_px, img_w, img_h, vb,
+                        color, stroke_width, underscore_count
+                    )
+                gap_idx += 1
+            i = j
+        else:
+            i += 1
+
+    # 6) Estimate width (use content bbox)
+    width_estimate = right - left
+    if width_estimate <= 0:
+        width_estimate = img_w // 2
+
+    return g, width_estimate
 
 # -----------------------------------------------------------------------------
-# Handwriting generation (multi-line + composition)
+# Generate full SVG (multi-line stacking)
 # -----------------------------------------------------------------------------
 def generate_handwriting(
     text,
@@ -185,53 +238,37 @@ def generate_handwriting(
     stroke_width=2,
     multiline=True
 ):
-    """
-    Generate a composed SVG that places underscores geometrically between
-    segments with pixel-accurate spacing.
-    """
     try:
-        lines = text.split('\n') if multiline else [text]
+        lines = text.split("\n") if multiline else [text]
 
-        # Validation + normalize slashes; KEEP underscores intact
         for idx, ln in enumerate(lines):
             if len(ln) > 75:
-                return f"Error: Line {idx + 1} is too long (max 75 characters)"
-            lines[idx] = ln.replace('/', '-').replace('\\', '-')
+                return f"Error: Line {idx+1} is too long (max 75 characters)"
 
-        # Fresh SVG root (no white background -> transparent)
-        svg_root = ET.Element('svg', {
-            'xmlns': 'http://www.w3.org/2000/svg',
-            'viewBox': '0 0 1200 800'
+        # Compose final SVG
+        svg_root = ET.Element("svg", {
+            "xmlns": "http://www.w3.org/2000/svg",
+            "viewBox": "0 0 1200 800"
         })
 
         y0 = 80.0
-        line_gap = 100.0
-        max_x = 0.0
-
-        for i, original_line in enumerate(lines):
-            line_y = y0 + i * line_gap
-            elems, right_x = render_line_with_underscores(
-                original_line,
-                style,
-                bias,
-                color,
-                stroke_width,
-                line_y,
-                x_start=40,
-                gap_left=8,
-                gap_right=8,
-                underscore_len=22
+        line_gap = 110.0
+        max_right = 0.0
+
+        for i, line in enumerate(lines):
+            g, w = render_line_svg_with_underscores(
+                line, style, bias, color, stroke_width
             )
-            for el in elems:
-                svg_root.append(el)
-            max_x = max(max_x, right_x)
+            # place this line group lower on the page
+            _translate_group(g, dx=40, dy=y0 + i * line_gap)
+            svg_root.append(g)
+            max_right = max(max_right, 40 + w)
 
-        # Tighten viewBox to content width + margin
-        width = max(300, int(max_x + 40))
-        height = int(y0 + len(lines) * line_gap)
-        svg_root.set('viewBox', f"0 0 {width} {height}")
+        height = int(y0 + len(lines) * line_gap + 80)
+        width = max(300, int(max_right + 40))
+        svg_root.set("viewBox", f"0 0 {width} {height}")
 
-        svg_content = ET.tostring(svg_root, encoding='unicode')
+        svg_content = ET.tostring(svg_root, encoding="unicode")
         with open("img/output.svg", "w", encoding="utf-8") as f:
             f.write(svg_content)
         return svg_content
@@ -243,7 +280,6 @@ def generate_handwriting(
 # PNG export (transparent)
 # -----------------------------------------------------------------------------
 def export_to_png(svg_content):
-    """Convert SVG to transparent PNG using CairoSVG and Pillow."""
     try:
         import cairosvg
         from PIL import Image
@@ -255,27 +291,16 @@ def export_to_png(svg_content):
         with open(tmp_svg, "w", encoding="utf-8") as f:
             f.write(svg_content)
 
-        # High scale for crisp strokes
-        cairosvg.svg2png(
-            url=tmp_svg,
-            write_to="img/output_temp.png",
-            scale=2.0,
-            background_color="none"
-        )
+        cairosvg.svg2png(url=tmp_svg, write_to="img/output_temp.png", scale=2.0, background_color="none")
 
         img = Image.open("img/output_temp.png")
-        if img.mode != 'RGBA':
-            img = img.convert('RGBA')
+        if img.mode != "RGBA":
+            img = img.convert("RGBA")
 
         # Ensure any near-white is transparent (safety)
         datas = img.getdata()
-        new_data = []
-        for item in datas:
-            if item[0] > 240 and item[1] > 240 and item[2] > 240:
-                new_data.append((255, 255, 255, 0))
-            else:
-                new_data.append(item)
-        img.putdata(new_data)
+        img.putdata([(255, 255, 255, 0) if r > 240 and g > 240 and b > 240 else (r, g, b, a)
+                     for (r, g, b, a) in datas])
 
         out_path = "img/output.png"
         img.save(out_path, "PNG")
@@ -296,8 +321,8 @@ def export_to_png(svg_content):
 # -----------------------------------------------------------------------------
 def generate_handwriting_wrapper(text, style, bias, color, stroke_width, multiline=True):
     svg = generate_handwriting(text, style, bias, color, stroke_width, multiline)
-    png_path = export_to_png(svg)
-    return svg, png_path, "img/output.svg"
+    png = export_to_png(svg)
+    return svg, png, "img/output.svg"
 
 css = """
 .container {max-width: 900px; margin: auto;}
@@ -307,31 +332,29 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("# 🖋️ Handwriting Synthesis (Underscore-safe)")
-    gr.Markdown("Generate realistic handwritten text with **perfect underscore alignment**—no retraining required.")
+    gr.Markdown("# 🖋️ Handwriting Synthesis (Underscore-aware)")
+    gr.Markdown("Underscores are placed only where typed, aligned to real gaps from the model output.")
 
     with gr.Row():
         with gr.Column(scale=2):
             text_input = gr.Textbox(
                 label="Text Input",
-                placeholder="Try: zeb_3asba, user_name, long__double__underscore",
-                lines=5,
-                max_lines=10,
+                placeholder="Try: user_name, zeb_3asba, or zeb aasba (no underscore should be drawn)",
+                lines=5, max_lines=10,
             )
             with gr.Row():
                 with gr.Column(scale=1):
-                    style_select = gr.Slider(minimum=0, maximum=12, step=1, value=9, label="Handwriting Style")
+                    style_select = gr.Slider(0, 12, step=1, value=9, label="Handwriting Style")
                 with gr.Column(scale=1):
-                    bias_slider = gr.Slider(minimum=0.5, maximum=1.0, step=0.05, value=0.75, label="Neatness (Higher = Neater)")
+                    bias_slider = gr.Slider(0.5, 1.0, step=0.05, value=0.75, label="Neatness (Higher = Neater)")
             with gr.Row():
                 with gr.Column(scale=1):
                     color_picker = gr.ColorPicker(label="Ink Color", value="#000000")
                 with gr.Column(scale=1):
-                    stroke_width = gr.Slider(minimum=1, maximum=4, step=0.5, value=2, label="Stroke Width")
+                    stroke_width = gr.Slider(1, 4, step=0.5, value=2, label="Stroke Width")
             with gr.Row():
                 generate_btn = gr.Button("Generate Handwriting", variant="primary")
                 clear_btn = gr.Button("Clear")
-
         with gr.Column(scale=3):
             output_svg = gr.HTML(label="Generated Handwriting (SVG)", elem_classes=["output-container"])
             output_png = gr.Image(type="filepath", label="Generated Handwriting (PNG)", elem_classes=["output-container"])
@@ -341,9 +364,9 @@ with gr.Blocks(css=css) as demo:
 
     gr.Markdown("""
     ### Notes
-    - Underscores are drawn between segments using **pixel-accurate measurements**.
-    - Slashes (/, \\) are normalized to dashes (-) for model stability.
-    - Each line ≤ 75 characters. Transparent PNG export included.
+    - Only underscores you type are drawn; normal spaces remain spaces.
+    - Alignment uses the *actual gaps* between model-drawn word blobs.
+    - Lines ≤ 75 chars. Slashes (/ and \\) normalized to dashes (-).
     """)
 
     generate_btn.click(
@@ -367,8 +390,6 @@ with gr.Blocks(css=css) as demo:
 # -----------------------------------------------------------------------------
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-
-    # Soft check for optional deps used here and in measurer
     missing = []
     try:
         import cairosvg  # noqa
@@ -378,9 +399,6 @@ if __name__ == "__main__":
         from PIL import Image  # noqa
     except ImportError:
         missing.append("pillow")
-
     if missing:
-        print(f"WARNING: Missing packages for underscore alignment & PNG export: {', '.join(missing)}")
-        print("Install with: pip install " + " ".join(missing))
-
+        print("Install:", " ".join(missing))
     demo.launch(server_name="0.0.0.0", server_port=port)