FlatCAM/VisPyVisuals.py

from vispy.visuals import CompoundVisual, LineVisual, MeshVisual
from vispy.scene.visuals import create_visual_node
from vispy.gloo import set_state
from vispy.geometry.triangulation import Triangulation
from vispy.color import Color
from shapely.geometry import Polygon, LineString, LinearRing
import threading
import numpy as np

try:
    from shapely.ops import triangulate
    import Polygon as gpc
except:
    pass


# Add clear_data method to LineVisual
def clear_data(self):
    self._bounds = None
    self._pos = None
    self._changed['pos'] = True
    self.update()

LineVisual.clear_data = clear_data


class ShapeGroup(object):
    def __init__(self, collection):
        """
        Represents group of shapes in collection
        :param collection: ShapeCollection
            Collection to work with
        """
        self._collection = collection
        self._indexes = []
        self._visible = True

    def add(self, shape, color=None, face_color=None, visible=True, update=False, layer=1):
        """
        Adds shape to collection and store index in group
        :param shape: shapely.geometry
            Shapely geometry object
        :param color: str, tuple
            Line/edge color
        :param face_color: str, tuple
            Polygon face color
        :param visible: bool
            Shape visibility
        :param update: bool
            Set True to redraw collection
        :param layer: int
            Layer number. 0 - lowest.
        """
        self._indexes.append(self._collection.add(shape, color, face_color, visible, update, layer))

    def clear(self, update=False):
        """
        Removes group shapes from collection, clear indexes
        :param update: bool
            Set True to redraw collection
        """
        for i in self._indexes:
            self._collection.remove(i, False)

        del self._indexes[:]

        if update:
            self._collection.redraw()

    def redraw(self):
        """
        Redraws shape collection
        """
        self._collection.redraw()

    @property
    def visible(self):
        """
        Visibility of group
        :return: bool
        """
        return self._visible

    @visible.setter
    def visible(self, value):
        """
        Visibility of group
        :param value: bool
        """
        self._visible = value
        for i in self._indexes:
            self._collection.data[i]['visible'] = value

        self._collection.redraw()


class ShapeCollectionVisual(CompoundVisual):
    def __init__(self, line_width=1, triangulation='gpc', layers=3, **kwargs):
        """
        Represents collection of shapes to draw on VisPy scene
        :param line_width: float
            Width of lines/edges
        :param triangulation: str
            Triangulation method used for polygons translation
            'vispy' - VisPy lib triangulation
            'gpc' - Polygon2 lib
        :param layers: int
            Layers count
            Each layer adds 2 visuals on VisPy scene. Be careful: more layers cause less fps
        :param kwargs:
        """
        self.data = {}
        self.last_key = -1
        self.lock = threading.Lock()

        self._meshes = [MeshVisual() for _ in range(0, layers)]
        self._lines = [LineVisual(antialias=True) for _ in range(0, layers)]

        self._line_width = line_width
        self._triangulation = triangulation

        visuals = [self._lines[i / 2] if i % 2 else self._meshes[i / 2] for i in range(0, layers * 2)]

        CompoundVisual.__init__(self, visuals, **kwargs)

        for m in self._meshes:
            m.set_gl_state(polygon_offset_fill=True, polygon_offset=(1, 1), cull_face=False)

        for l in self._lines:
            l.set_gl_state(blend=True)

        self.freeze()

    def add(self, shape, color=None, face_color=None, visible=True, update=False, layer=1):
        """
        Adds shape to collection
        :param shape: shapely.geometry
            Shapely geometry object
        :param color: str, tuple
            Line/edge color
        :param face_color: str, tuple
            Polygon face color
        :param visible: bool
            Shape visibility
        :param update: bool
            Set True to redraw collection
        :param layer: int
            Layer number. 0 - lowest.
        :return: int
            Index of shape
        """
        self.lock.acquire(True)
        self.last_key += 1
        key = self.last_key
        self.lock.release()

        self.data[key] = {'geometry': shape, 'color': color, 'face_color': face_color,
                                    'visible': visible, 'layer': layer}

        # TODO: Make in separate process -> data[key] (dict)
        self.update_shape_buffers(key)

        if update:
            self._update()

        return key

    def update_shape_buffers(self, key):
        """
        Translates Shapely geometry to internal buffers for speedup redraws
        :param key: int
            Shape index
        """
        mesh_vertices = []                      # Vertices for mesh
        mesh_tris = []                          # Faces for mesh
        mesh_colors = []                        # Face colors
        line_pts = []                           # Vertices for line
        line_colors = []                        # Line color

        geo, color, face_color = self.data[key]['geometry'], self.data[key]['color'], self.data[key]['face_color']

        if geo is not None and not geo.is_empty:
            simple = geo.simplify(0.01)         # Simplified shape
            pts = []                            # Shape line points
            tri_pts = []                        # Mesh vertices
            tri_tris = []                       # Mesh faces

            if type(geo) == LineString:
                # Prepare lines
                pts = self._linestring_to_segments(np.asarray(simple)).tolist()

            elif type(geo) == LinearRing:
                # Prepare lines
                pts = self._linearring_to_segments(np.asarray(simple)).tolist()

            elif type(geo) == Polygon:
                # Prepare polygon faces
                if face_color is not None:

                    if self._triangulation == 'vispy':
                        # VisPy triangulation
                        # Concatenated arrays of external & internal line rings
                        vertices = self._open_ring(np.asarray(simple.exterior))
                        edges = self._generate_edges(len(vertices))

                        for ints in simple.interiors:
                            v = self._open_ring(np.asarray(ints))
                            edges = np.append(edges, self._generate_edges(len(v)) + len(vertices), 0)
                            vertices = np.append(vertices, v, 0)

                        tri = Triangulation(vertices, edges)
                        tri.triangulate()
                        tri_pts, tri_tris = tri.pts.tolist(), tri.tris.tolist()

                    elif self._triangulation == 'gpc':

                        # GPC triangulation
                        p = gpc.Polygon(np.asarray(simple.exterior))

                        # Exclude all internal rings from polygon
                        for ints in simple.interiors:
                            q = gpc.Polygon(np.asarray(ints))
                            p -= q

                        # Triangulate polygon
                        for strip in p.triStrip():
                            # Generate tris indexes for triangle strip [[0, 1, 2], [1, 2, 3], [2, 3, 4], ... ]
                            ti = [[x + y + len(tri_pts) for x in range(0, 3)] for y in range(0, len(strip) - 2)]

                            # Append vertices & tris
                            tri_tris += ti
                            tri_pts += strip

                # Prepare polygon edges
                if color is not None:
                    pts = self._linearring_to_segments(np.asarray(simple.exterior)).tolist()
                    for ints in simple.interiors:
                        pts += self._linearring_to_segments(np.asarray(ints)).tolist()

            # Appending data for mesh
            if len(tri_pts) > 0 and len(tri_tris) > 0:
                mesh_tris += tri_tris
                mesh_vertices += tri_pts
                mesh_colors += [Color(face_color).rgba] * len(tri_tris)

            # Appending data for line
            if len(pts) > 0:
                line_pts += pts
                line_colors += [Color(color).rgba] * len(pts)

        # Store buffers
        self.data[key]['line_pts'] = line_pts
        self.data[key]['line_colors'] = line_colors
        self.data[key]['mesh_vertices'] = mesh_vertices
        self.data[key]['mesh_tris'] = mesh_tris
        self.data[key]['mesh_colors'] = mesh_colors

    def remove(self, key, update=False):
        """
        Removes shape from collection
        :param key: int
            Shape index to remove
        :param update:
            Set True to redraw collection
        """
        self.data.pop(key)
        if update:
            self._update()

    def clear(self, update=False):
        """
        Removes all shapes from colleciton
        :param update: bool
            Set True to redraw collection
        """
        self.data.clear()
        if update:
            self._update()

    def _update(self):
        """
        Merges internal buffers, sets data to visuals, redraws collection on scene
        """
        mesh_vertices = [[] for _ in range(0, len(self._meshes))]       # Vertices for mesh
        mesh_tris = [[] for _ in range(0, len(self._meshes))]           # Faces for mesh
        mesh_colors = [[] for _ in range(0, len(self._meshes))]         # Face colors
        line_pts = [[] for _ in range(0, len(self._lines))]             # Vertices for line
        line_colors = [[] for _ in range(0, len(self._lines))]          # Line color

        # Merge shapes buffers
        for data in self.data.values():
            if data['visible']:
                try:
                    line_pts[data['layer']] += data['line_pts']
                    line_colors[data['layer']] += data['line_colors']
                    mesh_tris[data['layer']] += [[x + len(mesh_vertices[data['layer']])
                                                  for x in y] for y in data['mesh_tris']]
                    mesh_vertices[data['layer']] += data['mesh_vertices']
                    mesh_colors[data['layer']] += data['mesh_colors']
                except Exception as e:
                    print "Data error", e

        # Updating meshes
        for i, mesh in enumerate(self._meshes):
            if len(mesh_vertices[i]) > 0:
                set_state(polygon_offset_fill=False)
                mesh.set_data(np.asarray(mesh_vertices[i]), np.asarray(mesh_tris[i], dtype=np.uint32),
                              face_colors=np.asarray(mesh_colors[i]))
            else:
                mesh.set_data()

            mesh._bounds_changed()

        # Updating lines
        for i, line in enumerate(self._lines):
            if len(line_pts[i]) > 0:
                line.set_data(np.asarray(line_pts[i]), np.asarray(line_colors[i]), self._line_width, 'segments')
            else:
                line.clear_data()

            line._bounds_changed()

        self._bounds_changed()

    def redraw(self):
        """
        Redraws collection
        """
        self._update()

    @staticmethod
    def _open_ring(vertices):
        """
        Make lines ring open
        :param vertices: numpy.array
            Array of lines vertices
        :return: numpy.array
            Opened line strip
        """
        return vertices[:-1] if not np.any(vertices[0] != vertices[-1]) else vertices

    @staticmethod
    def _generate_edges(count):
        """
        Generates edges indexes in form: [[0, 1], [1, 2], [2, 3], ... ]
        :param count: int
            Edges count
        :return: numpy.array
            Edges
        """
        edges = np.empty((count, 2), dtype=np.uint32)
        edges[:, 0] = np.arange(count)
        edges[:, 1] = edges[:, 0] + 1
        edges[-1, 1] = 0
        return edges

    @staticmethod
    def _linearring_to_segments(arr):
        # Close linear ring
        """
        Translates linear ring to line segments
        :param arr: numpy.array
            Array of linear ring vertices
        :return: numpy.array
            Line segments
        """
        if np.any(arr[0] != arr[-1]):
            arr = np.concatenate([arr, arr[:1]], axis=0)

        return ShapeCollection._linestring_to_segments(arr)

    @staticmethod
    def _linestring_to_segments(arr):
        """
        Translates line strip to segments
        :param arr: numpy.array
            Array of line strip vertices
        :return: numpy.array
            Line segments
        """
        return np.asarray(np.repeat(arr, 2, axis=0)[1:-1])


ShapeCollection = create_visual_node(ShapeCollectionVisual)