add utils.py

utils.py

@@ -0,0 +1,216 @@
+import os
+import re
+import numpy as np
+import vtracer
+import svgpathtools
+import cairosvg
+import io
+import cv2
+from lxml import etree
+from scipy.cluster.hierarchy import linkage, fcluster
+from scipy.spatial.distance import cdist 
+from python_tsp.heuristics import solve_tsp_local_search
+from fast_tsp import find_tour
+from svgpathtools import Path
+from tqdm import tqdm
+
+def parse_transform(transform_str):
+    if not transform_str: return np.eye(3)
+    matrix = np.eye(3)
+
+    numbers = r"[-+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?"
+
+    match = re.findall(r"matrix\(" + ",".join([numbers]*6) + r"\)", transform_str)
+    if match:
+        a, b, c, d, e, f = map(float, re.findall(numbers, match[0]))
+        m = np.array([[a, c, e], [b, d, f], [0, 0, 1]])
+        matrix = m @ matrix
+
+    match = re.findall(r"translate\(([^)]+)\)", transform_str)
+    if match:
+        parts = [float(v) for v in re.findall(numbers, match[0])]
+        tx, ty = parts if len(parts) == 2 else (parts[0], 0)
+        m = np.array([[1, 0, tx], [0, 1, ty], [0, 0, 1]])
+        matrix = m @ matrix
+
+    match = re.findall(r"scale\(([^)]+)\)", transform_str)
+    if match:
+        parts = [float(v) for v in re.findall(numbers, match[0])]
+        sx, sy = parts if len(parts) == 2 else (parts[0], parts[0])
+        m = np.array([[sx, 0, 0], [0, sy, 0], [0, 0, 1]])
+        matrix = m @ matrix
+
+    return matrix
+
+def get_global_transform(element):
+    transform = np.eye(3)
+    while element is not None:
+        t = element.get("transform")
+        if t: transform = parse_transform(t) @ transform
+        element = element.getparent()
+    return transform
+
+def get_transformed_paths_and_coords(svg_content, mode):
+    """
+    Parses the SVG, extracts path elements and their global transforms,
+    and returns data structures for sequencing and later rendering.
+    """
+    parser = etree.XMLParser(remove_blank_text=True)
+    if mode == 'file':
+        svg_content = io.BytesIO(svg_content.encode('utf-8'))
+    tree = etree.parse(svg_content, parser)
+    root = tree.getroot()
+
+    width_str = root.get("width")
+    height_str = root.get("height")
+    viewBox = root.get("viewBox")
+    path_elements = root.findall(".//{*}path")
+
+    if not viewBox and width_str and height_str:
+        # Remove potential units like 'px' to get clean numbers
+        width = re.sub(r'[a-zA-Z%]', '', width_str)
+        height = re.sub(r'[a-zA-Z%]', '', height_str)
+        viewBox = f"0 0 {width} {height}"
+        print(f"SVG missing viewBox. Created a default: '{viewBox}'")
+
+    width = root.get("width")
+    height = root.get("height")
+    transformed_coords = []
+    paths_data = []
+
+    for elem in path_elements:
+        d_string = elem.get('d')
+        if not d_string:
+            continue
+        path = svgpathtools.parse_path(d_string)
+
+        transform = get_global_transform(elem)
+
+        start_vec = np.array([[path.start.real], [path.start.imag], [1]])
+        transformed_start = transform @ start_vec
+
+        coord = (transformed_start[0, 0], transformed_start[1, 0])
+        transformed_coords.append(coord)
+
+        paths_data.append({
+            'path': path, 
+            'transform': transform,
+            'element': elem,
+            'coord': coord,  
+        })
+
+    print(f"Extracted {len(transformed_coords)} paths with their elements.")
+    return paths_data, np.array(transformed_coords), width, height, viewBox
+
+def transform_path(path, matrix):
+    """Apply a 3x3 numpy transform to an svgpathtools Path object."""
+    new_segments = []
+    for seg in path:
+        start = np.array([[seg.start.real], [seg.start.imag], [1]])
+        end = np.array([[seg.end.real], [seg.end.imag], [1]])
+        start_t = matrix @ start
+        end_t = matrix @ end
+        seg.start = complex(start_t[0,0], start_t[1,0])
+        seg.end = complex(end_t[0,0], end_t[1,0])
+        new_segments.append(seg)
+    return Path(*new_segments)
+
+def sequence_strokes(paths_data, coords, proximity_threshold=40):
+    """
+    Clusters strokes by proximity and then finds the optimal drawing order
+    """
+    if len(coords) < 2:
+        print("Fewer than 2 strokes, no sequencing needed.")
+        return paths_data
+
+    print("Clustering strokes by proximity...")
+
+    Z = linkage(coords, method='ward')
+    labels = fcluster(Z, t=proximity_threshold, criterion='distance')
+    num_clusters = len(set(labels))
+    print(f"{num_clusters} clusters detected.")
+
+    if num_clusters <= 1:
+        print("All strokes are in a single cluster, no reordering needed.")
+        return paths_data
+
+    clusters = {i: {'paths_data': [], 'coords': []} for i in range(1, num_clusters + 1)}
+    for i, label in enumerate(labels):
+        clusters[label]['paths_data'].append(paths_data[i])
+        clusters[label]['coords'].append(coords[i])
+
+    centroids = [np.mean(c['coords'], axis=0) for c in clusters.values()]
+
+    print("Solving TSP for optimal cluster drawing order...")
+    distance_matrix_float = cdist(centroids, centroids)
+
+    integer_distance_matrix = distance_matrix_float.astype(np.int32).tolist()
+
+    permutation = find_tour(integer_distance_matrix)
+
+    final_sequence = []
+    for cluster_idx in permutation:
+        cluster_label = cluster_idx + 1
+        final_sequence.extend(clusters[cluster_label]['paths_data'])
+
+    print("Final stroke sequence created.")
+    return final_sequence
+
+def serialize_paths(paths_data):
+    serialized = []
+    for data in paths_data:
+        elem = data['element']
+        style = elem.get("style", "")
+        fill_match = re.search(r'fill:\s*([^;]+)', style)
+        stroke_match = re.search(r'stroke:\s*([^;]+)', style)
+
+        serialized.append({
+            "d": elem.get("d"),
+            "transform": elem.get("transform"),
+            "fill": elem.get("fill") or (fill_match.group(1).strip() if fill_match else "#000000"),
+        })
+    return serialized
+
+def process_svg(svg_content, mode):
+    paths_data, coords, width, height, viewBox = get_transformed_paths_and_coords(svg_content, mode)
+
+    for data in paths_data:
+        path = data['path'] 
+        xmin, xmax, ymin, ymax = path.bbox()
+        area = (xmax - xmin) * (ymax - ymin)
+        data['area'] = area
+
+    areas = [p['area'] for p in paths_data]
+    areas_sorted = sorted(areas)
+    median_index = len(areas_sorted) // 2
+    layer_index = int(len(areas_sorted) * 0.99)
+    areaThreshold = areas_sorted[median_index]
+    layerThreshold = areas_sorted[layer_index]
+
+    fill_strokes = []
+    detail_strokes = []
+    layer_strokes = []
+
+    for data in paths_data:
+        if data['area'] >= layerThreshold:
+            layer_strokes.append(data)
+        elif data['area'] >= areaThreshold:
+            fill_strokes.append(data)
+        else:
+            detail_strokes.append(data)
+
+    ordered_fills = sequence_strokes(fill_strokes, np.array([p['coord'] for p in fill_strokes]))
+    ordered_details = sequence_strokes(detail_strokes, np.array([p['coord'] for p in detail_strokes]))
+
+    return {
+        "width": width,
+        "height": height,
+        "viewbox": viewBox,
+        "layers": serialize_paths(layer_strokes),
+        "fills": serialize_paths(ordered_fills),
+        "details": serialize_paths(ordered_details)
+    }
+
+        
+
+