Code cleanup: unused python vars & imports
[blender.git] / release / scripts / startup / bl_operators / vertexpaint_dirt.py
1 # ***** BEGIN GPL LICENSE BLOCK *****
2 #
3 # Script copyright (C) Campbell J Barton
4 #
5 # This program is free software; you can redistribute it and/or
6 # modify it under the terms of the GNU General Public License
7 # as published by the Free Software Foundation; either version 2
8 # of the License, or (at your option) any later version.
9 #
10 # This program is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 # GNU General Public License for more details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with this program; if not, write to the Free Software Foundation,
17 # Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
18 #
19 # ***** END GPL LICENSE BLOCK *****
20 # --------------------------------------------------------------------------
21
22 # <pep8 compliant>
23
24 # Contributor(s): Keith "Wahooney" Boshoff, Campbell Barton
25
26
27 def applyVertexDirt(me, blur_iterations, blur_strength, clamp_dirt, clamp_clean, dirt_only):
28     from mathutils import Vector
29     from math import acos
30
31     vert_tone = [0.0] * len(me.vertices)
32
33     min_tone = 180.0
34     max_tone = 0.0
35
36     # create lookup table for each vertex's connected vertices (via edges)
37     con = []
38
39     con = [[] for i in range(len(me.vertices))]
40
41     # add connected verts
42     for e in me.edges:
43         con[e.vertices[0]].append(e.vertices[1])
44         con[e.vertices[1]].append(e.vertices[0])
45
46     for i, v in enumerate(me.vertices):
47         vec = Vector()
48         no = v.normal
49         co = v.co
50
51         # get the direction of the vectors between the vertex and it's connected vertices
52         for c in con[i]:
53             vec += (me.vertices[c].co - co).normalized()
54
55         # normalize the vector by dividing by the number of connected verts
56         tot_con = len(con[i])
57
58         if tot_con == 0:
59             continue
60
61         vec /= tot_con
62
63         # angle is the acos() of the dot product between vert and connected verts normals
64         ang = acos(no.dot(vec))
65
66         # enforce min/max
67         ang = max(clamp_dirt, ang)
68
69         if not dirt_only:
70             ang = min(clamp_clean, ang)
71
72         vert_tone[i] = ang
73
74     # blur tones
75     for i in range(blur_iterations):
76         # backup the original tones
77         orig_vert_tone = list(vert_tone)
78
79         # use connected verts look up for blurring
80         for j, c in enumerate(con):
81             for v in c:
82                 vert_tone[j] += blur_strength * orig_vert_tone[v]
83
84             vert_tone[j] /= len(c) * blur_strength + 1
85
86     min_tone = min(vert_tone)
87     max_tone = max(vert_tone)
88
89     # debug information
90     # print(min_tone * 2 * math.pi)
91     # print(max_tone * 2 * math.pi)
92     # print(clamp_clean)
93     # print(clamp_dirt)
94
95     tone_range = max_tone - min_tone
96
97     if not tone_range:
98         return {'CANCELLED'}
99
100     active_col_layer = None
101
102     if me.vertex_colors:
103         for lay in me.vertex_colors:
104             if lay.active:
105                 active_col_layer = lay.data
106     else:
107         bpy.ops.mesh.vertex_color_add()
108         me.vertex_colors[0].active = True
109         active_col_layer = me.vertex_colors[0].data
110
111     if not active_col_layer:
112         return
113
114     use_paint_mask = me.use_paint_mask
115
116     for i, p in enumerate(me.polygons):
117         if not use_paint_mask or p.select:
118             for loop_index in p.loop_indices:
119                 loop = me.loops[loop_index]
120                 v = loop.vertex_index
121                 col = active_col_layer[loop_index].color
122                 tone = vert_tone[v]
123                 tone = (tone - min_tone) / tone_range
124
125                 if dirt_only:
126                     tone = min(tone, 0.5)
127                     tone *= 2.0
128
129                 col[0] = tone * col[0]
130                 col[1] = tone * col[1]
131                 col[2] = tone * col[2]
132     me.update()
133     return {'FINISHED'}
134
135
136 import bpy
137 from bpy.types import Operator
138 from bpy.props import FloatProperty, IntProperty, BoolProperty
139 from math import pi
140
141
142 class VertexPaintDirt(Operator):
143     bl_idname = "paint.vertex_color_dirt"
144     bl_label = "Dirty Vertex Colors"
145     bl_options = {'REGISTER', 'UNDO'}
146
147     blur_strength = FloatProperty(
148             name="Blur Strength",
149             description="Blur strength per iteration",
150             min=0.01, max=1.0,
151             default=1.0,
152             )
153     blur_iterations = IntProperty(
154             name="Blur Iterations",
155             description="Number of times to blur the colors (higher blurs more)",
156             min=0, max=40,
157             default=1,
158             )
159     clean_angle = FloatProperty(
160             name="Highlight Angle",
161             description="Less than 90 limits the angle used in the tonal range",
162             min=0.0, max=pi,
163             default=pi,
164             unit="ROTATION",
165             )
166     dirt_angle = FloatProperty(
167             name="Dirt Angle",
168             description="Less than 90 limits the angle used in the tonal range",
169             min=0.0, max=pi,
170             default=0.0,
171             unit="ROTATION",
172             )
173     dirt_only = BoolProperty(
174             name="Dirt Only",
175             description="Don't calculate cleans for convex areas",
176             default=False,
177             )
178
179     @classmethod
180     def poll(cls, context):
181         obj = context.object
182         return (obj and obj.type == 'MESH')
183
184     def execute(self, context):
185         import time
186
187         obj = context.object
188         mesh = obj.data
189
190         t = time.time()
191
192         ret = applyVertexDirt(mesh, self.blur_iterations, self.blur_strength, self.dirt_angle, self.clean_angle, self.dirt_only)
193
194         print('Dirt calculated in %.6f' % (time.time() - t))
195
196         return ret