initial commit

pydiffvg/parse_svg.py

@@ -0,0 +1,578 @@
+import torch
+import xml.etree.ElementTree as etree
+import numpy as np
+import diffvg
+import os
+import pydiffvg
+import svgpathtools
+import svgpathtools.parser
+import re
+import warnings
+import cssutils
+import logging
+cssutils.log.setLevel(logging.ERROR)
+
+def remove_namespaces(s):
+    """
+        {...} ... -> ...
+    """
+    return re.sub('{.*}', '', s)
+
+def parse_style(s, defs):
+    style_dict = {}
+    for e in s.split(';'):
+        key_value = e.split(':')
+        if len(key_value) == 2:
+            key = key_value[0].strip()
+            value = key_value[1].strip()
+            if key == 'fill' or key == 'stroke':
+                # Special case: convert colors into tensor in definitions so
+                # that different shapes can share the same color
+                value = parse_color(value, defs)
+            style_dict[key] = value
+    return style_dict
+
+def parse_hex(s):
+    """
+        Hex to tuple
+    """
+    s = s.lstrip('#')
+    if len(s) == 3:
+        s = s[0] + s[0] + s[1] + s[1] + s[2] + s[2]
+    rgb = tuple(int(s[i:i+2], 16) for i in (0, 2, 4))
+    # sRGB to RGB
+    # return torch.pow(torch.tensor([rgb[0] / 255.0, rgb[1] / 255.0, rgb[2] / 255.0]), 2.2)
+    return torch.pow(torch.tensor([rgb[0] / 255.0, rgb[1] / 255.0, rgb[2] / 255.0]), 1.0)
+
+def parse_int(s):
+    """
+        trim alphabets
+    """
+    return int(float(''.join(i for i in s if (not i.isalpha()))))
+
+def parse_color(s, defs):
+    if s is None:
+        return None
+    if isinstance(s, torch.Tensor):
+        return s
+    s = s.lstrip(' ')
+    color = torch.tensor([0.0, 0.0, 0.0, 1.0])
+    if s[0] == '#':
+        color[:3] = parse_hex(s)
+    elif s[:3] == 'url':
+        # url(#id)
+        color = defs[s[4:-1].lstrip('#')]
+    elif s == 'none':
+        color = None
+    elif s[:4] == 'rgb(':
+        rgb = s[4:-1].split(',')
+        color = torch.tensor([int(rgb[0]) / 255.0, int(rgb[1]) / 255.0, int(rgb[2]) / 255.0, 1.0])
+    elif s == 'none':
+        return None
+    else:
+        warnings.warn('Unknown color command ' + s)
+    return color
+
+# https://github.com/mathandy/svgpathtools/blob/7ebc56a831357379ff22216bec07e2c12e8c5bc6/svgpathtools/parser.py
+def _parse_transform_substr(transform_substr):
+    type_str, value_str = transform_substr.split('(')
+    value_str = value_str.replace(',', ' ')
+    values = list(map(float, filter(None, value_str.split(' '))))
+
+    transform = np.identity(3)
+    if 'matrix' in type_str:
+        transform[0:2, 0:3] = np.array([values[0:6:2], values[1:6:2]])
+    elif 'translate' in transform_substr:
+        transform[0, 2] = values[0]
+        if len(values) > 1:
+            transform[1, 2] = values[1]
+    elif 'scale' in transform_substr:
+        x_scale = values[0]
+        y_scale = values[1] if (len(values) > 1) else x_scale
+        transform[0, 0] = x_scale
+        transform[1, 1] = y_scale
+    elif 'rotate' in transform_substr:
+        angle = values[0] * np.pi / 180.0
+        if len(values) == 3:
+            offset = values[1:3]
+        else:
+            offset = (0, 0)
+        tf_offset = np.identity(3)
+        tf_offset[0:2, 2:3] = np.array([[offset[0]], [offset[1]]])
+        tf_rotate = np.identity(3)
+        tf_rotate[0:2, 0:2] = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
+        tf_offset_neg = np.identity(3)
+        tf_offset_neg[0:2, 2:3] = np.array([[-offset[0]], [-offset[1]]])
+
+        transform = tf_offset.dot(tf_rotate).dot(tf_offset_neg)
+    elif 'skewX' in transform_substr:
+        transform[0, 1] = np.tan(values[0] * np.pi / 180.0)
+    elif 'skewY' in transform_substr:
+        transform[1, 0] = np.tan(values[0] * np.pi / 180.0)
+    else:
+        # Return an identity matrix if the type of transform is unknown, and warn the user
+        warnings.warn('Unknown SVG transform type: {0}'.format(type_str))
+    return transform
+
+def parse_transform(transform_str):
+    """
+        Converts a valid SVG transformation string into a 3x3 matrix.
+        If the string is empty or null, this returns a 3x3 identity matrix
+    """
+    if not transform_str:
+        return np.identity(3)
+    elif not isinstance(transform_str, str):
+        raise TypeError('Must provide a string to parse')
+
+    total_transform = np.identity(3)
+    transform_substrs = transform_str.split(')')[:-1]  # Skip the last element, because it should be empty
+    for substr in transform_substrs:
+        total_transform = total_transform.dot(_parse_transform_substr(substr))
+
+    return torch.from_numpy(total_transform).type(torch.float32)
+
+def parse_linear_gradient(node, transform, defs):
+    begin = torch.tensor([0.0, 0.0])
+    end = torch.tensor([0.0, 0.0])
+    offsets = []
+    stop_colors = []
+    # Inherit from parent
+    for key in node.attrib:
+        if remove_namespaces(key) == 'href':
+            value = node.attrib[key]
+            parent = defs[value.lstrip('#')]
+            begin = parent.begin
+            end = parent.end
+            offsets = parent.offsets
+            stop_colors = parent.stop_colors
+
+    for attrib in node.attrib:
+        attrib = remove_namespaces(attrib)
+        if attrib == 'x1':
+            begin[0] = float(node.attrib['x1'])
+        elif attrib == 'y1':
+            begin[1] = float(node.attrib['y1'])
+        elif attrib == 'x2':
+            end[0] = float(node.attrib['x2'])
+        elif attrib == 'y2':
+            end[1] = float(node.attrib['y2'])
+        elif attrib == 'gradientTransform':
+            transform = transform @ parse_transform(node.attrib['gradientTransform'])
+
+    begin = transform @ torch.cat((begin, torch.ones([1])))
+    begin = begin / begin[2]
+    begin = begin[:2]
+    end = transform @ torch.cat((end, torch.ones([1])))
+    end = end / end[2]
+    end = end[:2]
+
+    for child in node:
+        tag = remove_namespaces(child.tag)
+        if tag == 'stop':
+            offset = float(child.attrib['offset'])
+            color = [0.0, 0.0, 0.0, 1.0]
+            if 'stop-color' in child.attrib:
+                c = parse_color(child.attrib['stop-color'], defs)
+                color[:3] = [c[0], c[1], c[2]]
+            if 'stop-opacity' in child.attrib:
+                color[3] = float(child.attrib['stop-opacity'])
+            if 'style' in child.attrib:
+                style = parse_style(child.attrib['style'], defs)
+                if 'stop-color' in style:
+                    c = parse_color(style['stop-color'], defs)
+                    color[:3] = [c[0], c[1], c[2]]
+                if 'stop-opacity' in style:
+                    color[3] = float(style['stop-opacity'])
+            offsets.append(offset)
+            stop_colors.append(color)
+    if isinstance(offsets, list):
+        offsets = torch.tensor(offsets)
+    if isinstance(stop_colors, list):
+        stop_colors = torch.tensor(stop_colors)
+
+    return pydiffvg.LinearGradient(begin, end, offsets, stop_colors)
+
+
+def parse_radial_gradient(node, transform, defs):
+    begin = torch.tensor([0.0, 0.0])
+    end = torch.tensor([0.0, 0.0])
+    center = torch.tensor([0.0, 0.0])
+    radius = torch.tensor([0.0, 0.0])
+    offsets = []
+    stop_colors = []
+    # Inherit from parent
+    for key in node.attrib:
+        if remove_namespaces(key) == 'href':
+            value = node.attrib[key]
+            parent = defs[value.lstrip('#')]
+            begin = parent.begin
+            end = parent.end
+            offsets = parent.offsets
+            stop_colors = parent.stop_colors
+
+    for attrib in node.attrib:
+        attrib = remove_namespaces(attrib)
+        if attrib == 'cx':
+            center[0] = float(node.attrib['cx'])
+        elif attrib == 'cy':
+            center[1] = float(node.attrib['cy'])
+        elif attrib == 'fx':
+            radius[0] = float(node.attrib['fx'])
+        elif attrib == 'fy':
+            radius[1] = float(node.attrib['fy'])
+        elif attrib == 'fr':
+            radius[0] = float(node.attrib['fr'])
+            radius[1] = float(node.attrib['fr'])
+        elif attrib == 'gradientTransform':
+            transform = transform @ parse_transform(node.attrib['gradientTransform'])
+
+    # TODO: this is incorrect
+    center = transform @ torch.cat((center, torch.ones([1])))
+    center = center / center[2]
+    center = center[:2]
+
+    for child in node:
+        tag = remove_namespaces(child.tag)
+        if tag == 'stop':
+            offset = float(child.attrib['offset'])
+            color = [0.0, 0.0, 0.0, 1.0]
+            if 'stop-color' in child.attrib:
+                c = parse_color(child.attrib['stop-color'], defs)
+                color[:3] = [c[0], c[1], c[2]]
+            if 'stop-opacity' in child.attrib:
+                color[3] = float(child.attrib['stop-opacity'])
+            if 'style' in child.attrib:
+                style = parse_style(child.attrib['style'], defs)
+                if 'stop-color' in style:
+                    c = parse_color(style['stop-color'], defs)
+                    color[:3] = [c[0], c[1], c[2]]
+                if 'stop-opacity' in style:
+                    color[3] = float(style['stop-opacity'])
+            offsets.append(offset)
+            stop_colors.append(color)
+    if isinstance(offsets, list):
+        offsets = torch.tensor(offsets)
+    if isinstance(stop_colors, list):
+        stop_colors = torch.tensor(stop_colors)
+
+    return pydiffvg.RadialGradient(begin, end, offsets, stop_colors)
+
+def parse_stylesheet(node, transform, defs):
+    # collect CSS classes
+    sheet = cssutils.parseString(node.text)
+    for rule in sheet:
+        if hasattr(rule, 'selectorText') and hasattr(rule, 'style'):
+            name = rule.selectorText
+            if len(name) >= 2 and name[0] == '.':
+                defs[name[1:]] = parse_style(rule.style.getCssText(), defs)
+    return defs
+
+def parse_defs(node, transform, defs):
+    for child in node:
+        tag = remove_namespaces(child.tag)
+        if tag == 'linearGradient':
+            if 'id' in child.attrib:
+                defs[child.attrib['id']] = parse_linear_gradient(child, transform, defs)
+        elif tag == 'radialGradient':
+            if 'id' in child.attrib:
+                defs[child.attrib['id']] = parse_radial_gradient(child, transform, defs)
+        elif tag == 'style':
+            defs = parse_stylesheet(child, transform, defs)
+    return defs
+
+def parse_common_attrib(node, transform, fill_color, defs):
+    attribs = {}
+    if 'class' in node.attrib:
+        attribs.update(defs[node.attrib['class']])
+    attribs.update(node.attrib)
+
+    name = ''
+    if 'id' in node.attrib:
+        name = node.attrib['id']
+
+    stroke_color = None
+    stroke_width = torch.tensor(0.5)
+    use_even_odd_rule = False
+
+    new_transform = transform
+    if 'transform' in attribs:
+        new_transform = transform @ parse_transform(attribs['transform'])
+    if 'fill' in attribs:
+        fill_color = parse_color(attribs['fill'], defs)
+    fill_opacity = 1.0
+    if 'fill-opacity' in attribs:
+        fill_opacity *= float(attribs['fill-opacity'])
+    if 'opacity' in attribs:
+        fill_opacity *= float(attribs['opacity'])
+    # Ignore opacity if the color is a gradient
+    if isinstance(fill_color, torch.Tensor):
+        fill_color[3] = fill_opacity
+
+    if 'fill-rule' in attribs:
+        if attribs['fill-rule'] == "evenodd":
+            use_even_odd_rule = True
+        elif attribs['fill-rule'] == "nonzero":
+            use_even_odd_rule = False
+        else:
+            warnings.warn('Unknown fill-rule: {}'.format(attribs['fill-rule']))
+
+    if 'stroke' in attribs:
+        stroke_color = parse_color(attribs['stroke'], defs)
+
+    if 'stroke-width' in attribs:
+        stroke_width = attribs['stroke-width']
+        if stroke_width[-2:] == 'px':
+            stroke_width = stroke_width[:-2]
+        stroke_width = torch.tensor(float(stroke_width) / 2.0)
+
+    if 'style' in attribs:
+        style = parse_style(attribs['style'], defs)
+        if 'fill' in style:
+            fill_color = parse_color(style['fill'], defs)
+        fill_opacity = 1.0
+        if 'fill-opacity' in style:
+            fill_opacity *= float(style['fill-opacity'])
+        if 'opacity' in style:
+            fill_opacity *= float(style['opacity'])
+        if 'fill-rule' in style:
+            if style['fill-rule'] == "evenodd":
+                use_even_odd_rule = True
+            elif style['fill-rule'] == "nonzero":
+                use_even_odd_rule = False
+            else:
+                warnings.warn('Unknown fill-rule: {}'.format(style['fill-rule']))
+        # Ignore opacity if the color is a gradient
+        if isinstance(fill_color, torch.Tensor):
+            fill_color[3] = fill_opacity
+        if 'stroke' in style:
+            if style['stroke'] != 'none':
+                stroke_color = parse_color(style['stroke'], defs)
+                # Ignore opacity if the color is a gradient
+                if isinstance(stroke_color, torch.Tensor):
+                    if 'stroke-opacity' in style:
+                        stroke_color[3] = float(style['stroke-opacity'])
+                    if 'opacity' in style:
+                        stroke_color[3] *= float(style['opacity'])
+                if 'stroke-width' in style:
+                    stroke_width = style['stroke-width']
+                    if stroke_width[-2:] == 'px':
+                        stroke_width = stroke_width[:-2]
+                    stroke_width = torch.tensor(float(stroke_width) / 2.0)
+
+        if isinstance(fill_color, pydiffvg.LinearGradient):
+            fill_color.begin = new_transform @ torch.cat((fill_color.begin, torch.ones([1])))
+            fill_color.begin = fill_color.begin / fill_color.begin[2]
+            fill_color.begin = fill_color.begin[:2]
+            fill_color.end = new_transform @ torch.cat((fill_color.end, torch.ones([1])))
+            fill_color.end = fill_color.end / fill_color.end[2]
+            fill_color.end = fill_color.end[:2]
+        if isinstance(stroke_color, pydiffvg.LinearGradient):
+            stroke_color.begin = new_transform @ torch.cat((stroke_color.begin, torch.ones([1])))
+            stroke_color.begin = stroke_color.begin / stroke_color.begin[2]
+            stroke_color.begin = stroke_color.begin[:2]
+            stroke_color.end = new_transform @ torch.cat((stroke_color.end, torch.ones([1])))
+            stroke_color.end = stroke_color.end / stroke_color.end[2]
+            stroke_color.end = stroke_color.end[:2]
+        if 'filter' in style:
+            print('*** WARNING ***: Ignoring filter for path with id "{}"'.format(name))
+
+    return new_transform, fill_color, stroke_color, stroke_width, use_even_odd_rule
+
+def is_shape(tag):
+    return tag == 'path' or tag == 'polygon' or tag == 'line' or tag == 'circle' or tag == 'rect'
+
+def parse_shape(node, transform, fill_color, shapes, shape_groups, defs):
+    tag = remove_namespaces(node.tag)
+    new_transform, new_fill_color, stroke_color, stroke_width, use_even_odd_rule = \
+        parse_common_attrib(node, transform, fill_color, defs)
+    if tag == 'path':
+        d = node.attrib['d']
+        name = ''
+        if 'id' in node.attrib:
+            name = node.attrib['id']
+        force_closing = new_fill_color is not None
+        paths = pydiffvg.from_svg_path(d, new_transform, force_closing)
+        for idx, path in enumerate(paths):
+            assert(path.points.shape[1] == 2)
+            path.stroke_width = stroke_width
+            path.source_id = name
+            path.id = "{}-{}".format(name,idx) if len(paths)>1 else name
+        prev_shapes_size = len(shapes)
+        shapes = shapes + paths
+        shape_ids = torch.tensor(list(range(prev_shapes_size, len(shapes))))
+        shape_groups.append(pydiffvg.ShapeGroup(\
+            shape_ids = shape_ids,
+            fill_color = new_fill_color,
+            stroke_color = stroke_color,
+            use_even_odd_rule = use_even_odd_rule,
+            id = name))
+    elif tag == 'polygon':
+        name = ''
+        if 'id' in node.attrib:
+            name = node.attrib['id']
+        force_closing = new_fill_color is not None
+        pts = node.attrib['points'].strip()
+        pts = pts.split(' ')
+        # import ipdb; ipdb.set_trace()
+        pts = [[float(y) for y in re.split(',| ', x)] for x in pts if x]
+        pts = torch.tensor(pts, dtype=torch.float32).view(-1, 2)
+        polygon = pydiffvg.Polygon(pts, force_closing)
+        polygon.stroke_width = stroke_width
+        shape_ids = torch.tensor([len(shapes)])
+        shapes.append(polygon)
+        shape_groups.append(pydiffvg.ShapeGroup(\
+            shape_ids = shape_ids,
+            fill_color = new_fill_color,
+            stroke_color = stroke_color,
+            use_even_odd_rule = use_even_odd_rule,
+            shape_to_canvas = new_transform,
+            id = name))
+    elif tag == 'line':
+        x1 = float(node.attrib['x1'])
+        y1 = float(node.attrib['y1'])
+        x2 = float(node.attrib['x2'])
+        y2 = float(node.attrib['y2'])
+        p1 = torch.tensor([x1, y1])
+        p2 = torch.tensor([x2, y2])
+        points = torch.stack((p1, p2))
+        line = pydiffvg.Polygon(points, False)
+        line.stroke_width = stroke_width
+        shape_ids = torch.tensor([len(shapes)])
+        shapes.append(line)
+        shape_groups.append(pydiffvg.ShapeGroup(\
+            shape_ids = shape_ids,
+            fill_color = new_fill_color,
+            stroke_color = stroke_color,
+            use_even_odd_rule = use_even_odd_rule,
+            shape_to_canvas = new_transform))
+    elif tag == 'circle':
+        radius = float(node.attrib['r'])
+        cx = float(node.attrib['cx'])
+        cy = float(node.attrib['cy'])
+        name = ''
+        if 'id' in node.attrib:
+            name = node.attrib['id']
+        center = torch.tensor([cx, cy])
+        circle = pydiffvg.Circle(radius = torch.tensor(radius),
+                                 center = center)
+        circle.stroke_width = stroke_width
+        shape_ids = torch.tensor([len(shapes)])
+        shapes.append(circle)
+        shape_groups.append(pydiffvg.ShapeGroup(\
+            shape_ids = shape_ids,
+            fill_color = new_fill_color,
+            stroke_color = stroke_color,
+            use_even_odd_rule = use_even_odd_rule,
+            shape_to_canvas = new_transform))
+    elif tag == 'ellipse':
+        rx = float(node.attrib['rx'])
+        ry = float(node.attrib['ry'])
+        cx = float(node.attrib['cx'])
+        cy = float(node.attrib['cy'])
+        name = ''
+        if 'id' in node.attrib:
+            name = node.attrib['id']
+        center = torch.tensor([cx, cy])
+        circle = pydiffvg.Circle(radius = torch.tensor(radius),
+                                 center = center)
+        circle.stroke_width = stroke_width
+        shape_ids = torch.tensor([len(shapes)])
+        shapes.append(circle)
+        shape_groups.append(pydiffvg.ShapeGroup(\
+            shape_ids = shape_ids,
+            fill_color = new_fill_color,
+            stroke_color = stroke_color,
+            use_even_odd_rule = use_even_odd_rule,
+            shape_to_canvas = new_transform))
+    elif tag == 'rect':
+        x = 0.0
+        y = 0.0
+        if x in node.attrib:
+            x = float(node.attrib['x'])
+        if y in node.attrib:
+            y = float(node.attrib['y'])
+        w = float(node.attrib['width'])
+        h = float(node.attrib['height'])
+        p_min = torch.tensor([x, y])
+        p_max = torch.tensor([x + w, x + h])
+        rect = pydiffvg.Rect(p_min = p_min, p_max = p_max)
+        rect.stroke_width = stroke_width
+        shape_ids = torch.tensor([len(shapes)])
+        shapes.append(rect)
+        shape_groups.append(pydiffvg.ShapeGroup(\
+            shape_ids = shape_ids,
+            fill_color = new_fill_color,
+            stroke_color = stroke_color,
+            use_even_odd_rule = use_even_odd_rule,
+            shape_to_canvas = new_transform))
+    return shapes, shape_groups
+
+def parse_group(node, transform, fill_color, shapes, shape_groups, defs):
+    if 'transform' in node.attrib:
+        transform = transform @ parse_transform(node.attrib['transform'])
+    if 'fill' in node.attrib:
+        fill_color = parse_color(node.attrib['fill'], defs)
+    for child in node:
+        tag = remove_namespaces(child.tag)
+        if is_shape(tag):
+            shapes, shape_groups = parse_shape(\
+                child, transform, fill_color, shapes, shape_groups, defs)
+        elif tag == 'g':
+            shapes, shape_groups = parse_group(\
+                child, transform, fill_color, shapes, shape_groups, defs)
+    return shapes, shape_groups
+
+def parse_scene(node):
+    canvas_width = -1
+    canvas_height = -1
+    defs = {}
+    shapes = []
+    shape_groups = []
+    fill_color = torch.tensor([0.0, 0.0, 0.0, 1.0])
+    transform = torch.eye(3)
+    if 'viewBox' in node.attrib:
+        view_box_array = node.attrib['viewBox'].split()
+        canvas_width = parse_int(view_box_array[2])
+        canvas_height = parse_int(view_box_array[3])
+    else:
+        if 'width' in node.attrib:
+            canvas_width = parse_int(node.attrib['width'])
+        else:
+            print('Warning: Can\'t find canvas width.')
+        if 'height' in node.attrib:
+            canvas_height = parse_int(node.attrib['height'])
+        else:
+            print('Warning: Can\'t find canvas height.')
+    for child in node:
+        tag = remove_namespaces(child.tag)
+        if tag == 'defs':
+            defs = parse_defs(child, transform, defs)
+        elif tag == 'style':
+            defs = parse_stylesheet(child, transform, defs)
+        elif tag == 'linearGradient':
+            if 'id' in child.attrib:
+                defs[child.attrib['id']] = parse_linear_gradient(child, transform, defs)
+        elif tag == 'radialGradient':
+            if 'id' in child.attrib:
+                defs[child.attrib['id']] = parse_radial_gradient(child, transform, defs)
+        elif is_shape(tag):
+            shapes, shape_groups = parse_shape(\
+                child, transform, fill_color, shapes, shape_groups, defs)
+        elif tag == 'g':
+            shapes, shape_groups = parse_group(\
+                child, transform, fill_color, shapes, shape_groups, defs)
+    return canvas_width, canvas_height, shapes, shape_groups
+
+def svg_to_scene(filename):
+    """
+        Load from a SVG file and convert to PyTorch tensors.
+    """
+
+    tree = etree.parse(filename)
+    root = tree.getroot()
+    cwd = os.getcwd()
+    if (os.path.dirname(filename) != ''):
+        os.chdir(os.path.dirname(filename))
+    ret = parse_scene(root)
+    os.chdir(cwd)
+    return ret