automatic zoom

app.py

@@ -418,6 +418,9 @@ from transformers import (
 )
 import re
 import xml.etree.ElementTree as ET
+import numpy as np
+import xml.etree.ElementTree as ET
+from svgpathtools import parse_path
 
 
 # ======================
@@ -548,10 +551,72 @@ def recolor_svg(svg_text: str, predicted_ids: set[str]) -> str:
 
     return ET.tostring(root, encoding="unicode")
 
-def set_viewbox(svg_text: str, viewbox: tuple[float, float, float, float]) -> str:
+def compute_viewbox_from_ids(svg_text: str, ids: set[str], margin_ratio: float = 0.08):
+    """
+    Compute a tight viewBox around the given country ids based on their path geometry.
+    Returns (min_x, min_y, width, height)
+    """
+    root = ET.fromstring(svg_text)
+
+    xmin, ymin = np.inf, np.inf
+    xmax, ymax = -np.inf, -np.inf
+
+    # Find all path/polygon with id in ids
+    for el in root.iter():
+        el_id = el.attrib.get("id")
+        if not el_id or el_id not in ids:
+            continue
+
+        tag = el.tag.split("}")[-1].lower()
+
+        if tag == "path":
+            d = el.attrib.get("d")
+            if not d:
+                continue
+            p = parse_path(d)
+            # svgpathtools bbox: (xmin, xmax, ymin, ymax)
+            bxmin, bxmax, bymin, bymax = p.bbox()
+            xmin = min(xmin, bxmin)
+            xmax = max(xmax, bxmax)
+            ymin = min(ymin, bymin)
+            ymax = max(ymax, bymax)
+
+        elif tag == "polygon":
+            pts = el.attrib.get("points", "").strip()
+            if not pts:
+                continue
+            coords = []
+            for chunk in pts.replace(",", " ").split():
+                coords.append(float(chunk))
+            xs = coords[0::2]
+            ys = coords[1::2]
+            xmin = min(xmin, min(xs))
+            xmax = max(xmax, max(xs))
+            ymin = min(ymin, min(ys))
+            ymax = max(ymax, max(ys))
+
+    if not np.isfinite(xmin):
+        # Didn't find those ids
+        return None
+
+    w = xmax - xmin
+    h = ymax - ymin
+
+    # Add a margin
+    mx = w * margin_ratio
+    my = h * margin_ratio
+
+    xmin -= mx
+    ymin -= my
+    w += 2 * mx
+    h += 2 * my
+
+    return (float(xmin), float(ymin), float(w), float(h))
+
+
+def set_viewbox(svg_text: str, viewbox):
     root = ET.fromstring(svg_text)
     root.attrib["viewBox"] = " ".join(str(x) for x in viewbox)
-    # Ensure it scales nicely
     root.attrib["preserveAspectRatio"] = "xMidYMid meet"
     return ET.tostring(root, encoding="unicode")
 

requirements.txt

@@ -1,4 +1,6 @@
 transformers>=4.30.0
 torch>=2.0.0
 pandas>=1.5.0
-gradio>=4.0
+gradio>=4.0
+svgpathtools
+numpy