Update app.py

app.py

@@ -1,7 +1,7 @@
 import os, re, io
 import xml.etree.ElementTree as ET
 import gradio as gr
-from hand import Hand  # your model
+from hand import Hand  # your handwriting model wrapper
 
 # -----------------------------------------------------------------------------
 # Setup
@@ -10,14 +10,13 @@ os.makedirs("img", exist_ok=True)
 hand = Hand()
 
 # -----------------------------------------------------------------------------
-# Helpers
+# SVG / coord helpers
 # -----------------------------------------------------------------------------
 def _parse_viewbox(root):
     """
-    Parse SVG viewBox robustly:
-    - accepts "0 0 600 160" or "0,0,600,160"
-    - falls back to width/height attributes
-    - strips 'px'
+    Robustly parse viewBox:
+    - accepts '0 0 600 160' or '0,0,600,160'
+    - falls back to width/height attrs (with px)
     """
     vb = root.get("viewBox")
     if vb:
@@ -31,16 +30,16 @@ def _parse_viewbox(root):
                 return (x, y, w, h)
             except ValueError:
                 pass
-    # fallback: width/height attributes
+
     def _num(v, default):
         if not v:
             return float(default)
-        v = v.strip().lower().replace("px", "")
-        v = v.replace(",", ".")
+        v = v.strip().lower().replace("px", "").replace(",", ".")
         try:
             return float(v)
         except ValueError:
             return float(default)
+
     w = _num(root.get("width"), 1200.0)
     h = _num(root.get("height"), 400.0)
     return (0.0, 0.0, w, h)
@@ -63,9 +62,13 @@ def _translate_group(elem, dx, dy):
     elem.set("transform", (prev + f" translate({dx},{dy})").strip())
 
 # -----------------------------------------------------------------------------
-# Tokenizer
+# Tokenization
 # -----------------------------------------------------------------------------
 def _tokenize_line(line):
+    """
+    Split into [('text'|'sep_space'|'sep_underscore', value), ...]
+    Preserves exact underscore counts and spaces.
+    """
     tokens, i, n = [], 0, len(line)
     while i < n:
         ch = line[i]
@@ -88,28 +91,36 @@ def _tokenize_line(line):
     return tokens
 
 def _display_text_from_tokens(tokens):
+    """What we actually render with the model (underscores & spaces -> single spaces)."""
     return "".join([v if t == "text" else " " for t, v in tokens])
 
 # -----------------------------------------------------------------------------
-# Rasterize and analyze gaps
+# Rasterization & analysis
 # -----------------------------------------------------------------------------
-def _rasterize_svg(svg_str, scale=2.0):
+def _rasterize_svg(svg_str, scale=3.0):
+    """Rasterize SVG string to RGBA image (higher scale = finer gap detection)."""
     import cairosvg
     from PIL import Image
     png = cairosvg.svg2png(bytestring=svg_str.encode("utf-8"), scale=scale, background_color="none")
     return Image.open(io.BytesIO(png)).convert("RGBA")
 
 def _find_blobs_and_gaps(alpha_img):
+    """
+    Return (blobs, gaps, content_bbox) using per-column alpha.
+    blobs/gaps are lists of x-intervals [start, end) in pixels.
+    """
     w, h = alpha_img.size
     bbox = alpha_img.getbbox()
     if not bbox:
         return [], [], (0, 0, w, h)
     left, top, right, bottom = bbox
+
     def col_has_ink(x):
         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)
@@ -123,13 +134,57 @@ def _find_blobs_and_gaps(alpha_img):
         x += 1
     if in_blob: blobs.append((start, right))
     else: gaps.append((start, right))
+
+    # only interior gaps (exclude margins)
     core_gaps = [(blobs[i][1], blobs[i + 1][0]) for i in range(len(blobs) - 1)]
     return blobs, core_gaps, (left, top, right, bottom)
 
+def _column_profile(alpha_img, top, bottom):
+    """Ink count per column."""
+    w, h = alpha_img.size
+    prof = [0] * w
+    for x in range(w):
+        s = 0
+        for y in range(top, bottom):
+            if alpha_img.getpixel((x, y)) > 0:
+                s += 1
+        prof[x] = s
+    return prof
+
+def _synthesize_gap_near(alpha_img, content_bbox, target_x_px, min_w_px=14, search_pct=0.18):
+    """
+    Create a gap near target_x_px by finding a local 'valley' (min ink column).
+    Ensures at least min_w_px width.
+    """
+    left, top, right, bottom = content_bbox
+    prof = _column_profile(alpha_img, top, bottom)
+    half = max(6, int((right - left) * search_pct))
+    x0 = max(left + 1, int(target_x_px - half))
+    x1 = min(right - 1, int(target_x_px + half))
+
+    best_x, best_v = x0, prof[x0]
+    for x in range(x0, x1):
+        v = prof[x]
+        if v < best_v:
+            best_x, best_v = x, v
+
+    half_w = max(7, min_w_px // 2)
+    g0 = max(left + 1, best_x - half_w)
+    g1 = min(right - 1, best_x + half_w)
+    if g1 - g0 < min_w_px:
+        pad = (min_w_px - (g1 - g0)) // 2 + 1
+        g0 = max(left + 1, g0 - pad)
+        g1 = min(right - 1, g1 + pad)
+    return (g0, g1)
+
 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):
+    """
+    Fill the pixel gap with n evenly spaced underscores.
+    """
     gap_w = max(0, gap_px[1] - gap_px[0] - 2 * pad_px)
-    if gap_w <= 6 or n <= 0: return
+    if gap_w <= 6 or n <= 0:
+        return
     slot = gap_w / n
     line_len = max(8, int(slot * 0.85) - between_px)
     offset = (slot - line_len) / 2.0
@@ -139,7 +194,34 @@ def _draw_underscores_in_gap(root, gap_px, baseline_px, img_w, img_h, vb,
         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)
+        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)
+
+def _draw_margin_underscores(root, edge_px, side, baseline_px, img_w, img_h, vb,
+                             color, stroke_width, n, pad_px=6, between_px=4, max_len_px=28):
+    """
+    Draw underscores left of the first blob ('left') or right of the last blob ('right')
+    for leading/trailing underscores like '__init__'.
+    """
+    if n <= 0: return
+    slot = max_len_px + between_px
+    for i in range(n):
+        if side == "left":
+            x1_px = edge_px - pad_px - i * slot
+            x0_px = x1_px - max_len_px
+        else:
+            x0_px = edge_px + pad_px + i * slot
+            x1_px = x0_px + max_len_px
+        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)
@@ -149,11 +231,19 @@ def _draw_underscores_in_gap(root, gap_px, baseline_px, img_w, img_h, vb,
         root.append(p)
 
 # -----------------------------------------------------------------------------
-# Line renderer with underscores
+# Render a single line with correct underscores
 # -----------------------------------------------------------------------------
 def render_line_svg_with_underscores(line, style, bias, color, stroke_width):
+    """
+    Draw the line via the model, then insert underscores:
+    - Use real gaps when the model produced multiple blobs.
+    - If merged (no gap), synthesize a gap near expected split.
+    - Draw leading/trailing underscores at margins.
+    """
     tokens = _tokenize_line(line)
     display_line = _display_text_from_tokens(tokens).replace("/", "-").replace("\\", "-")
+
+    # Render the (spaces-only) display line once
     hand.write(
         filename="img/line.tmp.svg",
         lines=[display_line if display_line.strip() else " "],
@@ -162,69 +252,134 @@ def render_line_svg_with_underscores(line, style, bias, color, stroke_width):
     )
     root = ET.parse("img/line.tmp.svg").getroot()
     vb = _parse_viewbox(root)
-    img = _rasterize_svg(ET.tostring(root, encoding="unicode"))
+
+    # Rasterize and analyze ink
+    img = _rasterize_svg(ET.tostring(root, encoding="unicode"), scale=3.0)
     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)
+
+    # Pack model paths into a <g>
     g = ET.Element("g")
-    for p in _extract_paths(root): g.append(p)
-    gap_idx = 0
-    text_run_count = sum(1 for t, v in tokens if t == "text" and len(v) > 0)
-    max_gaps = max(0, min(len(gaps), text_run_count - 1))
+    for p in _extract_paths(root):
+        g.append(p)
+
+    # Count leading/trailing underscores (for margins)
+    leading_us = 0
+    i = 0
+    while i < len(tokens) and tokens[i][0] != "text":
+        if tokens[i][0] == "sep_underscore":
+            leading_us += len(tokens[i][1])
+        i += 1
+
+    trailing_us = 0
+    j = len(tokens) - 1
+    while j >= 0 and tokens[j][0] != "text":
+        if tokens[j][0] == "sep_underscore":
+            trailing_us += len(tokens[j][1])
+        j -= 1
+
+    # Between-word positions and underscore counts
+    positions = []  # list of underscore_count (can be 0 if just space)
     i = 0
     while i < len(tokens):
         t, v = tokens[i]
         if t == "text" and len(v) > 0:
             j = i + 1
-            underscore_count, saw_sep = 0, False
+            u_count, saw_sep = 0, 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])
+                    u_count += len(tokens[j][1])
                 j += 1
-            if saw_sep and gap_idx < max_gaps:
-                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
+            if saw_sep:
+                positions.append(u_count)
             i = j
         else:
             i += 1
+
+    words = [v for (t, v) in tokens if t == "text" and len(v) > 0]
+    total_chars = sum(len(w) for w in words) or 1
+
+    used_gap_idx = 0
+    seen_chars = 0
+    word_idx = 0
+    for pos_idx, underscores_here in enumerate(positions):
+        if word_idx < len(words):
+            seen_chars += len(words[word_idx])
+            word_idx += 1
+
+        if underscores_here <= 0:
+            continue  # pure space in input -> draw nothing
+
+        # Prefer detected interior gap, else synthesize near expected split
+        if used_gap_idx < len(gaps):
+            gap_px = gaps[used_gap_idx]
+            used_gap_idx += 1
+        else:
+            ratio = min(0.95, max(0.05, seen_chars / float(total_chars)))
+            target_x_px = left + ratio * (right - left)
+            gap_px = _synthesize_gap_near(alpha, content_bbox, target_x_px)
+
+        _draw_underscores_in_gap(
+            g, gap_px, baseline_px, img_w, img_h, vb,
+            color, stroke_width, underscores_here
+        )
+
+    # Margin underscores (leading/trailing)
+    if blobs:
+        first_left = blobs[0][0]
+        last_right = blobs[-1][1]
+        if leading_us:
+            _draw_margin_underscores(
+                g, first_left, "left", baseline_px, img_w, img_h, vb,
+                color, stroke_width, leading_us
+            )
+        if trailing_us:
+            _draw_margin_underscores(
+                g, last_right, "right", baseline_px, img_w, img_h, vb,
+                color, stroke_width, trailing_us
+            )
+
     width_estimate = right - left if right > left else img_w // 2
     return g, width_estimate
 
 # -----------------------------------------------------------------------------
-# Generate handwriting (multi-line)
+# Generate full SVG (multi-line stacking)
 # -----------------------------------------------------------------------------
 def generate_handwriting(text, style, bias=0.75, color="#000000", stroke_width=2, multiline=True):
     try:
         lines = text.split("\n") if multiline else [text]
         for idx, ln in enumerate(lines):
             if len(ln) > 75:
-                return f"Error: Line {idx+1} too long (max 75 chars)"
-        svg_root = ET.Element("svg", {"xmlns": "http://www.w3.org/2000/svg","viewBox": "0 0 1200 800"})
+                return f"Error: Line {idx+1} is too long (max 75 characters)"
+
+        svg_root = ET.Element("svg", {"xmlns": "http://www.w3.org/2000/svg", "viewBox": "0 0 1200 800"})
         y0, line_gap, max_right = 80.0, 110.0, 0.0
+
         for i, line in enumerate(lines):
             g, w = render_line_svg_with_underscores(line, style, bias, color, stroke_width)
             _translate_group(g, dx=40, dy=y0 + i * line_gap)
             svg_root.append(g)
             max_right = max(max_right, 40 + w)
+
         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")
-        with open("img/output.svg", "w", encoding="utf-8") as f: f.write(svg_content)
+        with open("img/output.svg", "w", encoding="utf-8") as f:
+            f.write(svg_content)
         return svg_content
+
     except Exception as e:
         return f"Error: {str(e)}"
 
 # -----------------------------------------------------------------------------
-# Export to PNG
+# PNG export (transparent)
 # -----------------------------------------------------------------------------
 def export_to_png(svg_content):
     try:
@@ -232,19 +387,26 @@ def export_to_png(svg_content):
         from PIL import Image
         if not svg_content or svg_content.startswith("Error:"):
             return None
+
         tmp_svg = "img/temp.svg"
-        with open(tmp_svg, "w", encoding="utf-8") as f: f.write(svg_content)
+        with open(tmp_svg, "w", encoding="utf-8") as f:
+            f.write(svg_content)
+
         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")
+
+        # Make near-white fully transparent (safety)
         datas = img.getdata()
-        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])
+        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")
         try: os.remove("img/output_temp.png")
         except: pass
         return out_path
+
     except Exception as e:
         print(f"Error converting to PNG: {str(e)}")
         return None
@@ -257,27 +419,60 @@ def generate_handwriting_wrapper(text, style, bias, color, stroke_width, multili
     png = export_to_png(svg)
     return svg, png, "img/output.svg"
 
-css = ".container {max-width:900px;margin:auto;} .output-container{min-height:300px;}"
+css = """
+.container {max-width: 900px; margin: auto;}
+.output-container {min-height: 300px;}
+"""
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# 🖋️ Handwriting Synthesis (Underscore-aware)")
+    gr.Markdown("Underscores are inserted only where typed, aligned to real or synthesized gaps.")
+
     with gr.Row():
         with gr.Column(scale=2):
-            text_input = gr.Textbox(label="Text Input", placeholder="Try: user_name vs zeb aasba", lines=5)
-            style_select = gr.Slider(0, 12, step=1, value=9, label="Handwriting Style")
-            bias_slider = gr.Slider(0.5, 1.0, step=0.05, value=0.75, label="Neatness")
-            color_picker = gr.ColorPicker(label="Ink Color", value="#000000")
-            stroke_width = gr.Slider(1, 4, step=0.5, value=2, label="Stroke Width")
+            text_input = gr.Textbox(
+                label="Text Input",
+                placeholder="Try: user_name, zeb_3asba, long__name, __init__",
+                lines=5
+            )
+            with gr.Row():
+                style_select = gr.Slider(0, 12, step=1, value=9, label="Handwriting Style")
+                bias_slider = gr.Slider(0.5, 1.0, step=0.05, value=0.75, label="Neatness (Higher = Neater)")
+            with gr.Row():
+                color_picker = gr.ColorPicker(label="Ink Color", value="#000000")
+                stroke_width = gr.Slider(1, 4, step=0.5, value=2, label="Stroke Width")
             generate_btn = gr.Button("Generate Handwriting", variant="primary")
         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"])
             download_svg_file = gr.File(label="Download SVG")
             download_png_file = gr.File(label="Download PNG")
-    generate_btn.click(fn=generate_handwriting_wrapper,
-                       inputs=[text_input, style_select, bias_slider, color_picker, stroke_width],
-                       outputs=[output_svg, output_png, download_svg_file]
-    ).then(fn=lambda p: p, inputs=[output_png], outputs=[download_png_file])
 
+    generate_btn.click(
+        fn=generate_handwriting_wrapper,
+        inputs=[text_input, style_select, bias_slider, color_picker, stroke_width],
+        outputs=[output_svg, output_png, download_svg_file]
+    ).then(
+        fn=lambda p: p,
+        inputs=[output_png],
+        outputs=[download_png_file]
+    )
+
+# -----------------------------------------------------------------------------
+# Main
+# -----------------------------------------------------------------------------
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    # Make sure dependencies for rasterization/export are present
+    missing = []
+    try:
+        import cairosvg  # noqa
+    except ImportError:
+        missing.append("cairosvg")
+    try:
+        from PIL import Image  # noqa
+    except ImportError:
+        missing.append("pillow")
+    if missing:
+        print("Install:", " ".join(missing))
+    port = int(os.environ.get("PORT", 7860))
+    demo.launch(server_name="0.0.0.0", server_port=port)