a8e4743ae4b54927e65d3d3df56eebfbf50ee4e8
[blender.git] / release / scripts / freestyle / modules / freestyle / utils.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 """
20 This module contains helper functions used for Freestyle style module
21 writing.
22 """
23
24 __all__ = (
25     "ContextFunctions",
26     "bound",
27     "bounding_box",
28     "find_matching_vertex",
29     "getCurrentScene",
30     "get_chain_length",
31     "get_test_stroke",
32     "integrate",
33     "iter_distance_along_stroke",
34     "iter_distance_from_camera",
35     "iter_distance_from_object",
36     "iter_material_value",
37     "iter_t2d_along_stroke",
38     "pairwise",
39     "phase_to_direction",
40     "rgb_to_bw",
41     "stroke_curvature",
42     "stroke_normal",
43     "tripplewise",
44     )
45
46
47 # module members
48 from _freestyle import (
49     ContextFunctions,
50     getCurrentScene,
51     integrate,
52     )
53
54 # constructs for helper functions in Python
55 from freestyle.types import (
56     Interface0DIterator,
57     Stroke,
58     StrokeVertexIterator,
59     )
60
61 from mathutils import Vector
62 from functools import lru_cache, namedtuple
63 from math import cos, sin, pi
64 from itertools import tee
65
66
67 # -- real utility functions  -- #
68
69 def rgb_to_bw(r, g, b):
70     """Method to convert rgb to a bw intensity value."""
71     return 0.35 * r + 0.45 * g + 0.2 * b
72
73
74 def bound(lower, x, higher):
75     """Returns x bounded by a maximum and minimum value. Equivalent to:
76     return min(max(x, lower), higher)
77     """
78     # this is about 50% quicker than min(max(x, lower), higher)
79     return (lower if x <= lower else higher if x >= higher else x)
80
81
82 def bounding_box(stroke):
83     """
84     Returns the maximum and minimum coordinates (the bounding box) of the stroke's vertices
85     """
86     x, y = zip(*(svert.point for svert in stroke))
87     return (Vector((min(x), min(y))), Vector((max(x), max(y))))
88
89 def get_dashed_pattern(linestyle):
90     """Extracts the dashed pattern from the various UI options """
91     pattern = []
92     if linestyle.dash1 > 0 and linestyle.gap1 > 0:
93         pattern.append(linestyle.dash1)
94         pattern.append(linestyle.gap1)
95     if linestyle.dash2 > 0 and linestyle.gap2 > 0:
96         pattern.append(linestyle.dash2)
97         pattern.append(linestyle.gap2)
98     if linestyle.dash3 > 0 and linestyle.gap3 > 0:
99         pattern.append(linestyle.dash3)
100         pattern.append(linestyle.gap3)
101     return pattern
102
103 # -- General helper functions -- #
104
105
106 @lru_cache(maxsize=32)
107 def phase_to_direction(length):
108     """
109     Returns a list of tuples each containing:
110     - the phase
111     - a Vector with the values of the cosine and sine of 2pi * phase  (the direction)
112     """
113     results = list()
114     for i in range(length):
115         phase = i / (length - 1)
116         results.append((phase, Vector((cos(2 * pi * phase), sin(2 * pi * phase)))))
117     return results
118
119 # A named tuple primitive used for storing data that has an upper and
120 # lower bound (e.g., thickness, range and certain values)
121 BoundedProperty = namedtuple("BoundedProperty", ["min", "max", "delta"])
122
123 # -- helper functions for chaining -- #
124
125 def get_chain_length(ve, orientation):
126     """Returns the 2d length of a given ViewEdge."""
127     from freestyle.chainingiterators import pyChainSilhouetteGenericIterator
128     length = 0.0
129     # setup iterator
130     _it = pyChainSilhouetteGenericIterator(False, False)
131     _it.begin = ve
132     _it.current_edge = ve
133     _it.orientation = orientation
134     _it.init()
135
136     # run iterator till end of chain
137     while not (_it.is_end):
138         length += _it.object.length_2d
139         if (_it.is_begin):
140             # _it has looped back to the beginning;
141             # break to prevent infinite loop
142             break
143         _it.increment()
144
145     # reset iterator
146     _it.begin = ve
147     _it.current_edge = ve
148     _it.orientation = orientation
149
150     # run iterator till begin of chain
151     if not _it.is_begin:
152         _it.decrement()
153         while not (_it.is_end or _it.is_begin):
154             length += _it.object.length_2d
155             _it.decrement()
156
157     return length
158
159
160 def find_matching_vertex(id, it):
161     """Finds the matching vertex, or returns None."""
162     return next((ve for ve in it if ve.id == id), None)
163
164 # -- helper functions for iterating -- #
165
166 def pairwise(iterable, types={Stroke, StrokeVertexIterator}):
167     """Yields a tuple containing the previous and current object """
168     # use .incremented() for types that support it
169     if type(iterable) in types:
170         it = iter(iterable)
171         return zip(it, it.incremented())
172     else:
173         a, b = tee(iterable)
174         next(b, None)
175         return zip(a, b)
176
177
178 def tripplewise(iterable):
179     """Yields a tuple containing the current object and its immediate neighbors """
180     a, b, c = tee(iterable)
181     next(b, None)
182     next(c, None)
183     return zip(a, b, c)
184
185
186 def iter_t2d_along_stroke(stroke):
187     """Yields the progress along the stroke."""
188     total = stroke.length_2d
189     distance = 0.0
190     # yield for the comparison from the first vertex to itself
191     yield 0.0
192     for prev, svert in pairwise(stroke):
193         distance += (prev.point - svert.point).length
194         yield min(distance / total, 1.0) if total != 0.0 else 0.0
195
196
197 def iter_distance_from_camera(stroke, range_min, range_max, normfac):
198     """
199     Yields the distance to the camera relative to the maximum
200     possible distance for every stroke vertex, constrained by
201     given minimum and maximum values.
202     """
203     for svert in stroke:
204         # length in the camera coordinate
205         distance = svert.point_3d.length
206         if range_min < distance < range_max:
207             yield (svert, (distance - range_min) / normfac)
208         else:
209             yield (svert, 0.0) if range_min > distance else (svert, 1.0)
210
211
212 def iter_distance_from_object(stroke, location, range_min, range_max, normfac):
213     """
214     yields the distance to the given object relative to the maximum
215     possible distance for every stroke vertex, constrained by
216     given minimum and maximum values.
217     """
218     for svert in stroke:
219         distance = (svert.point_3d - location).length  # in the camera coordinate
220         if range_min < distance < range_max:
221             yield (svert, (distance - range_min) / normfac)
222         else:
223             yield (svert, 0.0) if distance < range_min else (svert, 1.0)
224
225
226 def iter_material_value(stroke, func, attribute):
227     "Yields a specific material attribute from the vertex' underlying material."
228     it = Interface0DIterator(stroke)
229     for svert in it:
230         material = func(it)
231         # main
232         if attribute == 'LINE':
233             value = rgb_to_bw(*material.line[0:3])
234         elif attribute == 'DIFF':
235             value = rgb_to_bw(*material.diffuse[0:3])
236         elif attribute == 'SPEC':
237             value = rgb_to_bw(*material.specular[0:3])
238         # line seperate
239         elif attribute == 'LINE_R':
240             value = material.line[0]
241         elif attribute == 'LINE_G':
242             value = material.line[1]
243         elif attribute == 'LINE_B':
244             value = material.line[2]
245         elif attribute == 'LINE_A':
246             value = material.line[3]
247         # diffuse seperate
248         elif attribute == 'DIFF_R':
249             value = material.diffuse[0]
250         elif attribute == 'DIFF_G':
251             value = material.diffuse[1]
252         elif attribute == 'DIFF_B':
253             value = material.diffuse[2]
254         elif attribute == 'ALPHA':
255             value = material.diffuse[3]
256         # specular seperate
257         elif attribute == 'SPEC_R':
258             value = material.specular[0]
259         elif attribute == 'SPEC_G':
260             value = material.specular[1]
261         elif attribute == 'SPEC_B':
262             value = material.specular[2]
263         elif attribute == 'SPEC_HARDNESS':
264             value = material.shininess
265         else:
266             raise ValueError("unexpected material attribute: " + attribute)
267         yield (svert, value)
268
269 def iter_distance_along_stroke(stroke):
270     "Yields the absolute distance along the stroke up to the current vertex."
271     distance = 0.0
272     # the positions need to be copied, because they are changed in the calling function
273     points = tuple(svert.point.copy() for svert in stroke)
274     yield distance
275     for prev, curr in pairwise(points):
276         distance += (prev - curr).length
277         yield distance
278
279 # -- mathmatical operations -- #
280
281
282 def stroke_curvature(it):
283     """
284     Compute the 2D curvature at the stroke vertex pointed by the iterator 'it'.
285     K = 1 / R
286     where R is the radius of the circle going through the current vertex and its neighbors
287     """
288     for _ in it:
289         if (it.is_begin or it.is_end):
290             yield 0.0
291             continue
292         else:
293             it.decrement()
294             prev, current, succ = it.object.point.copy(), next(it).point.copy(), next(it).point.copy()
295             # return the iterator in an unchanged state
296             it.decrement()
297
298         ab = (current - prev)
299         bc = (succ - current)
300         ac = (prev - succ)
301
302         a, b, c = ab.length, bc.length, ac.length
303
304         try:
305             area = 0.5 * ab.cross(ac)
306             K = (4 * area) / (a * b * c)
307         except ZeroDivisionError:
308             K = 0.0
309
310         yield abs(K)
311
312 def stroke_normal(stroke):
313     """
314     Compute the 2D normal at the stroke vertex pointed by the iterator
315     'it'.  It is noted that Normal2DF0D computes normals based on
316     underlying FEdges instead, which is inappropriate for strokes when
317     they have already been modified by stroke geometry modifiers.
318
319     The returned normals are dynamic: they update when the
320     vertex position (and therefore the vertex normal) changes.
321     for use in geometry modifiers it is advised to 
322     cast this generator function to a tuple or list
323     """
324     n = len(stroke) - 1
325
326     for i, svert in enumerate(stroke):
327         if i == 0:
328             e = stroke[i + 1].point - svert.point
329             yield Vector((e[1], -e[0])).normalized()
330         elif i == n:
331             e = svert.point - stroke[i - 1].point
332             yield Vector((e[1], -e[0])).normalized()
333         else:
334             e1 = stroke[i + 1].point - svert.point
335             e2 = svert.point - stroke[i - 1].point
336             n1 = Vector((e1[1], -e1[0])).normalized()
337             n2 = Vector((e2[1], -e2[0])).normalized()
338             yield (n1 + n2).normalized()
339
340 def get_test_stroke():
341     """Returns a static stroke object for testing """
342     from freestyle.types import Stroke, Interface0DIterator, StrokeVertexIterator, SVertex, Id, StrokeVertex
343     # points for our fake stroke
344     points = (Vector((1.0, 5.0, 3.0)), Vector((1.0, 2.0, 9.0)),
345               Vector((6.0, 2.0, 3.0)), Vector((7.0, 2.0, 3.0)), 
346               Vector((2.0, 6.0, 3.0)), Vector((2.0, 8.0, 3.0)))
347     ids = (Id(0, 0), Id(1, 1), Id(2, 2), Id(3, 3), Id(4, 4), Id(5, 5))
348
349     stroke = Stroke()
350     it = iter(stroke)
351
352     for svert in map(SVertex, points, ids):
353         stroke.insert_vertex(StrokeVertex(svert), it)
354         it = iter(stroke)
355
356     stroke.update_length()
357     return stroke