change python scripts so modules which register with blender have a register() functi...
[blender.git] / release / scripts / op / 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 LICENCE BLOCK *****
20 # --------------------------------------------------------------------------
21
22 # <pep8 compliant>
23
24 # History
25 #
26 # Originally written by Campbell Barton aka ideasman42
27 #
28 # 2009-11-01: * 2.5 port by Keith "Wahooney" Boshoff
29 #              * Replaced old method with my own, speed is similar (about 0.001 sec on Suzanne)
30 #               but results are far more accurate
31 #
32
33 import bpy
34 import math
35 import time
36
37 from Mathutils import Vector
38 from bpy.props import *
39
40
41 def applyVertexDirt(me, blur_iterations, blur_strength, clamp_dirt, clamp_clean, dirt_only):
42 ##    Window.WaitCursor(1)
43
44     #BPyMesh.meshCalcNormals(me)
45
46     vert_tone = [0.0] * len(me.verts)
47
48     min_tone = 180.0
49     max_tone = 0.0
50
51     # create lookup table for each vertex's connected vertices (via edges)
52     con = []
53
54     con = [[] for i in range(len(me.verts))]
55
56     # add connected verts
57     for e in me.edges:
58         con[e.verts[0]].append(e.verts[1])
59         con[e.verts[1]].append(e.verts[0])
60
61     for v in me.verts:
62         vec = Vector()
63         no = v.normal
64         co = v.co
65
66         # get the direction of the vectors between the vertex and it's connected vertices
67         for c in con[v.index]:
68             vec += Vector(me.verts[c].co - co).normalize()
69
70         # normalize the vector by dividing by the number of connected verts
71         vec /= len(con[v.index])
72
73         # angle is the acos of the dot product between vert and connected verts normals
74         ang = math.acos(no.dot(vec))
75
76         # enforce min/max
77         ang = max(clamp_dirt, ang)
78
79         if not dirt_only:
80             ang = min(clamp_clean, ang)
81
82         vert_tone[v.index] = ang
83
84     # blur tones
85     for i in range(blur_iterations):
86         # backup the original tones
87         orig_vert_tone = list(vert_tone)
88
89         # use connected verts look up for blurring
90         for j, c in enumerate(con):
91             for v in c:
92                 vert_tone[j] += blur_strength * orig_vert_tone[v]
93
94             vert_tone[j] /= len(c) * blur_strength + 1
95
96     min_tone = min(vert_tone)
97     max_tone = max(vert_tone)
98
99     # debug information
100     # print(min_tone * 2 * math.pi)
101     # print(max_tone * 2 * math.pi)
102     # print(clamp_clean)
103     # print(clamp_dirt)
104
105     tone_range = max_tone - min_tone
106
107     if not tone_range:
108         return
109
110     active_col_layer = None
111
112     if len(me.vertex_colors):
113         for lay in me.vertex_colors:
114             if lay.active:
115                 active_col_layer = lay.data
116     else:
117         bpy.ops.mesh.vertex_color_add()
118         me.vertex_colors[0].active = True
119         active_col_layer = me.vertex_colors[0].data
120
121     if not active_col_layer:
122         return('CANCELLED', )
123
124     for i, f in enumerate(me.faces):
125         if not me.use_paint_mask or f.selected:
126
127             f_col = active_col_layer[i]
128
129             f_col = [f_col.color1, f_col.color2, f_col.color3, f_col.color4]
130
131             for j, v in enumerate(f.verts):
132                 col = f_col[j]
133                 tone = vert_tone[me.verts[v].index]
134                 tone = (tone - min_tone) / tone_range
135
136                 if dirt_only:
137                     tone = min(tone, 0.5)
138                     tone *= 2
139
140                 col[0] = tone * col[0]
141                 col[1] = tone * col[1]
142                 col[2] = tone * col[2]
143
144 ##    Window.WaitCursor(0)
145
146
147 class VertexPaintDirt(bpy.types.Operator):
148
149     bl_idname = "mesh.vertex_paint_dirt"
150     bl_label = "Dirty Vertex Colors"
151     bl_register = True
152     bl_undo = True
153
154     blur_strength = FloatProperty(name="Blur Strength", description="Blur strength per iteration", default=1.0, min=0.01, max=1.0)
155     blur_iterations = IntProperty(name="Blur Iterations", description="Number times to blur the colors. (higher blurs more)", default=1, min=0, max=40)
156     clean_angle = FloatProperty(name="Highlight Angle", description="Less then 90 limits the angle used in the tonal range", default=180.0, min=0.0, max=180.0)
157     dirt_angle = FloatProperty(name="Dirt Angle", description="Less then 90 limits the angle used in the tonal range", default=0.0, min=0.0, max=180.0)
158     dirt_only = BoolProperty(name="Dirt Only", description="Dont calculate cleans for convex areas", default=False)
159
160     def execute(self, context):
161         obj = context.object
162
163         if not obj or obj.type != 'MESH':
164             print('Error, no active mesh object, aborting')
165             return('CANCELLED',)
166
167         mesh = obj.data
168
169         t = time.time()
170
171         applyVertexDirt(mesh, self.properties.blur_iterations, self.properties.blur_strength, math.radians(self.properties.dirt_angle), math.radians(self.properties.clean_angle), self.properties.dirt_only)
172
173         print('Dirt calculated in %.6f' % (time.time() - t))
174
175         return('FINISHED',)
176
177
178 def register():
179     bpy.types.register(VertexPaintDirt)
180
181 def unregister():
182     bpy.types.unregister(VertexPaintDirt)
183