context.py 8.54 KB
Newer Older
Dario Seyb's avatar
Dario Seyb committed
1

Dario Seyb's avatar
Dario Seyb committed
2 3
import pythreejs as three
import numpy as np
4 5
from ipywidgets import HTML, Text
from traitlets import link, dlink, Bool
Dario Seyb's avatar
Dario Seyb committed
6
from IPython.display import display
7
from .indexed_attribute import *
Dario Seyb's avatar
Dario Seyb committed
8
from .mesh_helper import calculateFaceNormals, calculatePointNormals
9
from .mesh import *
Dario Seyb's avatar
Dario Seyb committed
10 11

class Context(object):
12 13 14
    """
    This is a module docstring
    """
Dario Seyb's avatar
Dario Seyb committed
15

Dario Seyb's avatar
Dario Seyb committed
16
    def __init__(self, width=600, height=400, background_color = '#dddddd'):
Dario Seyb's avatar
Dario Seyb committed
17 18 19
        self.camera = three.PerspectiveCamera(position=[1,1,1], fov=20,
                children=[three.DirectionalLight(color='#ffffff', position=[-30, 50, 10], intensity=1.0)])

Dario Seyb's avatar
Dario Seyb committed
20
        self.camera.aspect = width/float(height)
Dario Seyb's avatar
Dario Seyb committed
21

Dario Seyb's avatar
Dario Seyb committed
22 23 24
        self.minCorner = np.array([ 100000.0,  100000.0,  100000.0])
        self.maxCorner = np.array([-100000.0, -100000.0, -100000.0])

25
        self.scene = three.Scene(children=[three.AmbientLight(color='#aaaaaa'), self.camera])
Dario Seyb's avatar
Dario Seyb committed
26
        self.scene.background = background_color
Dario Seyb's avatar
Dario Seyb committed
27 28

        self.orbit_controls = three.OrbitControls(controlling=self.camera)
29 30
        self.click_picker = three.Picker(controlling=self.scene, root=self.scene, event='dblclick')

Dario Seyb's avatar
Dario Seyb committed
31
        self.renderer = three.Renderer(camera=self.camera,  background=background_color, background_opacity=1,
32
                                       scene=self.scene, controls=[self.orbit_controls, self.click_picker],
Dario Seyb's avatar
Dario Seyb committed
33
                                       width=width, height=height, antialias=True )
Dario Seyb's avatar
Dario Seyb committed
34

35 36 37 38 39
        def on_picked(change):
            self.orbit_controls.target = change['new']

        self.click_picker.observe(on_picked, names=['point'])

40
    def draw_sphere(self):
Dario Seyb's avatar
Dario Seyb committed
41 42
        mesh_obj = three.Mesh(three.SphereBufferGeometry(20, 16, 16),
                              three.MeshPhysicalMaterial(color='red'),
43 44 45
                              position=[0, 0, 0])
        self.scene.add(mesh_obj)

46 47 48 49 50 51 52 53 54 55 56 57 58 59 60

    def draw(self, obj, shading='flat', z_offset=0.5, texture=None, point_size = 1, perspective = False, line_width = 1):
        obj.prepare_render()

        if isinstance(obj, Mesh):
            return self.draw_faces(obj.vertices, obj.tri_face_indices, obj.normals, obj.colors, obj.uvs, shading, z_offset, texture)
        elif isinstance(obj, EdgeList):
            return self.draw_edges(obj.vertices, obj.edge_indices, obj.colors, obj.uvs, z_offset, texture, line_width)
        elif isinstance(obj, PointList):
            return self.draw_vertices(obj.vertices, obj.colors, obj.uvs, point_size, z_offset, texture, perspective)

        return self
            


Dario Seyb's avatar
Dario Seyb committed
61 62
    def draw_faces(self, vertices, face_indices, normals=None, colors=None, uvs=None,
                   shading='flat', z_offset=0.5, texture=None):
Dario Seyb's avatar
Dario Seyb committed
63

64 65
        assert(len(face_indices) > 0 and len(vertices) > 0)

66

Dario Seyb's avatar
Dario Seyb committed
67
        vertices = np.array(vertices)
68
        face_indices = np.array(face_indices)
Dario Seyb's avatar
Dario Seyb committed
69

Dario Seyb's avatar
Dario Seyb committed
70
        mat = None
71
        # Setup Material
Dario Seyb's avatar
Dario Seyb committed
72 73 74 75 76
        if shading == 'flat' or shading == 'hidden' or shading == 'wireframe':
            mat = three.MeshLambertMaterial(color='#dddddd', colorWrite=shading is not 'hidden')
        elif shading == 'phong':
            mat = three.MeshPhongMaterial(color='#dddddd')

Dario Seyb's avatar
Dario Seyb committed
77
        mat.wireframe = shading is 'wireframe'
78
        mat.map = texture
Dario Seyb's avatar
Dario Seyb committed
79 80 81 82 83 84

        if z_offset is not 0:
            mat.polygonOffset = True
            mat.polygonOffsetFactor = z_offset
            mat.polygonOffsetUnits = 0.1

Dario Seyb's avatar
Dario Seyb committed
85
        if normals is 'vertex':
Dario Seyb's avatar
Dario Seyb committed
86 87
            normals = calculatePointNormals(vertices, face_indices)
        elif normals is 'face':
88 89
            normals = calculateFaceNormals(vertices, face_indices)

90 91 92 93 94
        # Resolve the given attributes, it's ok if they are None
        resolved_attribs = resolve_attributes(face_indices, [normals, colors, uvs])
        resolved_normals    = resolved_attribs[0]
        resolved_colors     = resolved_attribs[1]
        resolved_uvs        = resolved_attribs[2]
Dario Seyb's avatar
Dario Seyb committed
95

96 97 98 99 100 101 102 103
        # De-index our vertices
        vertices, face_indices = stretch_vertices(vertices, face_indices)

        # Create attributes dictionary, always use at least position and index
        attributes=dict(
            position = three.BufferAttribute(np.asarray(vertices, dtype=np.float32), normalized=False),
            index =    three.BufferAttribute(np.asarray(face_indices, dtype=np.uint16).ravel(), normalized=False),
        )
Dario Seyb's avatar
Dario Seyb committed
104

105 106
        if resolved_colors is not None:
            mat.vertexColors = 'VertexColors'
Dario Seyb's avatar
Dario Seyb committed
107
            mat.color = 'white'
108
            attributes['color'] = three.BufferAttribute(np.asarray(resolved_colors, dtype=np.float32))
Dario Seyb's avatar
Dario Seyb committed
109

110 111 112 113 114
        if resolved_normals is not None:
            attributes['normal'] = three.BufferAttribute(np.asarray(resolved_normals, dtype=np.float32))

        if resolved_uvs is not None:
            attributes['uv'] = three.BufferAttribute(np.asarray(resolved_uvs, dtype=np.float32))
Dario Seyb's avatar
Dario Seyb committed
115 116 117

        mesh_geom = three.BufferGeometry(attributes=attributes)

Dario Seyb's avatar
Dario Seyb committed
118
        mesh_obj = three.Mesh(geometry=mesh_geom, material=mat)
Dario Seyb's avatar
Dario Seyb committed
119 120
        self.minCorner = np.minimum(self.minCorner, vertices.min(axis=0))
        self.maxCorner = np.maximum(self.maxCorner, vertices.max(axis=0))
Dario Seyb's avatar
Dario Seyb committed
121

Dario Seyb's avatar
Dario Seyb committed
122 123 124
        self.scene.add(mesh_obj)
        return self

125
    def draw_edges(self, vertices, edge_indices=None, colors=None, uvs=None, z_offset=0, texture=None, linewidth=1):
Dario Seyb's avatar
Dario Seyb committed
126

Dario Seyb's avatar
Dario Seyb committed
127
        mat = three.LineBasicMaterial(color='#000000')
128
        mat.linewidth = linewidth
129
        mat.map = texture
Dario Seyb's avatar
Dario Seyb committed
130

Dario Seyb's avatar
Dario Seyb committed
131 132 133 134
        if z_offset is not 0:
            mat.polygonOffset = True
            mat.polygonOffsetFactor = z_offset
            mat.polygonOffsetUnits = 0.1
Dario Seyb's avatar
Dario Seyb committed
135

136 137
        if edge_indices is None:
            edge_indices = zip(np.arange(0, len(vertices)-1), np.arange(1, len(vertices)))
Dario Seyb's avatar
Dario Seyb committed
138

Dario Seyb's avatar
Dario Seyb committed
139 140
        assert(len(edge_indices) > 0 and len(vertices) > 0)

Dario Seyb's avatar
Dario Seyb committed
141
        vertices = np.array(vertices)
Dario Seyb's avatar
Dario Seyb committed
142
        edge_indices = np.array(edge_indices)
Dario Seyb's avatar
Dario Seyb committed
143

144 145 146 147 148 149 150 151 152 153 154 155 156
        resolved_attribs = resolve_attributes(edge_indices, [colors, uvs])
        resolved_colors = resolved_attribs[0]
        resolved_uvs    = resolved_attribs[1]

        vertices, edge_indices = stretch_vertices(vertices, edge_indices)

        attributes=dict(
            position = three.BufferAttribute(np.asarray(vertices, dtype=np.float32), normalized=False),
            index =    three.BufferAttribute(np.asarray(edge_indices, dtype=np.uint16).ravel(), normalized=False),
        )

        if resolved_colors is not None:
            mat.vertexColors = 'VertexColors'
Dario Seyb's avatar
Dario Seyb committed
157
            mat.color = 'white'
158
            attributes['color'] = three.BufferAttribute(np.asarray(resolved_colors, dtype=np.float32))
Dario Seyb's avatar
Dario Seyb committed
159

160 161
        if resolved_uvs is not None:
            attributes['uv'] = three.BufferAttribute(np.asarray(resolved_uvs, dtype=np.float32))
Dario Seyb's avatar
Dario Seyb committed
162

Dario Seyb's avatar
Dario Seyb committed
163
        geom = three.BufferGeometry(attributes = attributes)
Dario Seyb's avatar
Dario Seyb committed
164

Dario Seyb's avatar
Dario Seyb committed
165
        mesh_obj = three.LineSegments(geometry=geom, material=mat)
Dario Seyb's avatar
Dario Seyb committed
166

Dario Seyb's avatar
Dario Seyb committed
167 168
        self.minCorner = np.minimum(self.minCorner, vertices.min(axis=0))
        self.maxCorner = np.maximum(self.maxCorner, vertices.max(axis=0))
Dario Seyb's avatar
Dario Seyb committed
169 170 171 172

        self.scene.add(mesh_obj)
        return self

173
    def draw_vertices(self, vertices, colors=None, uvs=None, point_size=1, z_offset=0, texture=None, perspective=False):
Dario Seyb's avatar
Dario Seyb committed
174

Dario Seyb's avatar
Dario Seyb committed
175
        vertices = np.array(vertices)
Dario Seyb's avatar
Dario Seyb committed
176 177 178 179 180
        matColor = colors

        if colors is None:
            matColor = 'black'
        elif hasattr(colors, '__len__') and (not isinstance(colors, str)):
181
            matColor = '#ffffff'
Dario Seyb's avatar
Dario Seyb committed
182 183 184 185 186 187 188 189
        else:
            colors = None

        attributes = dict(
            position = three.BufferAttribute(np.asarray(vertices, dtype=np.float32), normalized=False),
        )

        mat = three.PointsMaterial(color=matColor, sizeAttenuation=perspective, size=point_size)
190
        mat.map = texture
Dario Seyb's avatar
Dario Seyb committed
191 192 193 194 195 196 197 198 199

        if z_offset is not 0:
            mat.polygonOffset = True
            mat.polygonOffsetFactor = z_offset
            mat.polygonOffsetUnits = 0.1

        if colors is not None:
            attributes['color'] = three.BufferAttribute(np.asarray(colors, dtype=np.float32))
            mat.vertexColors = 'VertexColors'
200 201 202
        
        if uvs is not None:
            attributes['uv'] = three.BufferAttribute(np.asarray(uvs, dtype=np.float32))
Dario Seyb's avatar
Dario Seyb committed
203

Dario Seyb's avatar
Dario Seyb committed
204 205
        geom = three.BufferGeometry(attributes=attributes)
        mesh_obj = three.Points(geometry=geom, material=mat)
Dario Seyb's avatar
Dario Seyb committed
206

Dario Seyb's avatar
Dario Seyb committed
207 208
        self.minCorner = np.minimum(self.minCorner, vertices.min(axis=0))
        self.maxCorner = np.maximum(self.maxCorner, vertices.max(axis=0))
Dario Seyb's avatar
Dario Seyb committed
209

Dario Seyb's avatar
Dario Seyb committed
210 211 212 213
        self.scene.add(mesh_obj)
        return self

    def display(self):
Dario Seyb's avatar
Dario Seyb committed
214 215 216
        center = (self.minCorner + self.maxCorner) * 0.5
        distance = max(np.linalg.norm(center-self.minCorner), np.linalg.norm(center-self.maxCorner))
        self.camera.position = tuple( np.array(self.camera.position) * distance * 5)
Dario Seyb's avatar
Dario Seyb committed
217 218
        self.orbit_controls.target = tuple(center)

Dario Seyb's avatar
Dario Seyb committed
219
        display(self.renderer)