8000 Add ambient occlusion rasterization in main · chr-wei/tinyrenderer_python@d896ce4 · GitHub
[go: up one dir, main page]
Skip to content

Commit d896ce4

Browse files
author
Christian Weihsbach
committed
Add ambient occlusion rasterization in main
1 parent dcbbd4c commit d896ce4

File tree

3 files changed

+165
-51
lines changed

3 files changed

+165
-51
lines changed

‎geom.py‎

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -49,38 +49,38 @@ def __init__(self, *args, shape: tuple = None): # pylint: disable=unused-argumen
4949

5050
def __add__(self, other):
5151
if type(self) == type(other):
52-
(elems, _) = matadd(self.get_field_values(), self._shape,
52+
(elems, _) = mat_add(self.get_field_values(), self._shape,
5353
other.get_field_values(), other._shape)
5454
return type(self)(*elems)
5555
else:
5656
raise TypeError
5757

5858
def __sub__(self, other):
5959
if type(self) == type(other):
60-
(elems, _) = matsub(self.get_field_values(), self._shape,
60+
(elems, _) = mat_sub(self.get_field_values(), self._shape,
6161
other.get_field_values(), other._shape)
6262
return type(self)(*elems)
6363

6464
raise TypeError
6565

6666
def __mul__(self, other):
6767
if isinstance(other, (float, int)):
68-
(elems, _) = compmul(self.get_field_values(), self._shape, other)
68+
(elems, _) = comp_mul(self.get_field_values(), self._shape, other)
6969
return type(self)(*elems)
7070

7171
# All other cases should already have been handled in instance classes
7272
raise TypeError
7373

7474
def __rmul__(self, other):
7575
if isinstance(other, (float, int)):
76-
(elems, _) = compmul(self.get_field_values(), self._shape, other)
76+
(elems, _) = comp_mul(self.get_field_values(), self._shape, other)
7777
return type(self)(*elems)
7878

7979
raise TypeError
8080

8181
def __truediv__(self, other):
8282
if isinstance(other, (float, int)):
83-
(elems, _) = compdiv(self.get_field_values(), self._shape, other)
83+
(elems, _) = comp_div(self.get_field_values(), self._shape, other)
8484
return type(self)(*elems)
8585

8686
raise TypeError
@@ -162,19 +162,21 @@ def tr(self): # pylint: disable=invalid-name
162162
class MixinVector(MixinAlgebra):
163163
"""Mixin providing additional functionalty for vectors based on typing.NamedTuple."""
164164
def __mul__(self, other):
165-
if isinstance(self, MixinVector) and isinstance(other, MixinVector) and \
166-
self._shape[0] < other._shape[0]:
167-
# Calc scalar product
168-
(elems, _) = matmul(self.get_field_values(), self._shape,
169-
other.get_field_values(), other._shape)
165+
if isinstance(self, MixinVector) and isinstance(other, MixinVector):
166+
if self._shape[0] > other._shape[0]:
167+
raise ShapeMissmatchException
168+
else:
169+
# Calc scalar product
170+
(elems, _) = matmul(self.get_field_values(), self._shape,
171+
other.get_field_values(), other._shape)
170172
return elems[0]
171173

172174
# Fallback to MixinAlgebra __mul__
173175
return super().__mul__(other)
174176

175177
def __floordiv__(self, other):
176178
if isinstance(other, (float, int)):
177-
(elems, _) = compfloor(self.get_field_values(), self._shape, other)
179+
(elems, _) = comp_floor(self.get_field_values(), self._shape, other)
178180
return type(self)(elems, shape = self._shape)
179181

180182
return ValueError
@@ -439,7 +441,7 @@ def unpack_nested_iterable_to_list(it_er: Iterable):
439441
it_er = list(chain.from_iterable(it_er))
440442
return it_er
441443

442-
def compmul(mat_0: list, shape_0: tuple, factor: float):
444+
def comp_mul(mat_0: list, shape_0: tuple, factor: float):
443445
"""Performing componentwise multiplication with factor c."""
444446
(rows_0, cols_0) = shape_0
445447

@@ -450,7 +452,7 @@ def compmul(mat_0: list, shape_0: tuple, factor: float):
450452
# Return coefficients and shape tuple
451453
return [e * factor for e in mat_0], shape_0
452454

453-
def compdiv(mat_0: list, shape_0: tuple, divisor: float):
455+
def comp_div(mat_0: list, shape_0: tuple, divisor: float):
454456
"""Performing componentwise real division by divisor."""
455457
(rows_0, cols_0) = shape_0
456458

@@ -461,7 +463,7 @@ def compdiv(mat_0: list, shape_0: tuple, divisor: float):
461463
# Return coefficients and shape tuple
462464
return [e / divisor for e in mat_0], shape_0
463465

464-
def compfloor(mat_0: list, shape_0: tuple, divisor: float):
466+
def comp_floor(mat_0: list, shape_0: tuple, divisor: float):
465467
"""Performing componentwise floor division."""
466468
(rows_0, cols_0) = shape_0
467469

@@ -477,7 +479,7 @@ def vect_norm(all_elems: list):
477479
squared = [elem**2 for elem in all_elems]
478480
return sqrt(reduce(operator.add, squared))
479481

480-
def matadd(mat_0: list, shape_0: tuple, mat_1: list, shape_1: tuple):
482+
def mat_add(mat_0: list, shape_0: tuple, mat_1: list, shape_1: tuple):
481483
"""Performing componentwise addition."""
482484
(rows_0, cols_0) = shape_0
483485
(rows_1, cols_1) = shape_1
@@ -491,10 +493,10 @@ def matadd(mat_0: list, shape_0: tuple, mat_1: list, shape_1: tuple):
491493
# Return coefficients and shape tuple
492494
return map(operator.add, mat_0, mat_1), shape_0
493495

494-
def matsub(mat_0: list, shape_0: tuple, mat_1: list, shape_1: tuple):
496+
def mat_sub(mat_0: list, shape_0: tuple, mat_1: list, shape_1: tuple):
495497
"""Performing componentwise substraction."""
496498
mat_1 = [e * -1 for e in mat_1]
497-
return matadd(mat_0, shape_0, mat_1, shape_1)
499+
return mat_add(mat_0, shape_0, mat_1, shape_1)
498500

499501
def is_row_vect(shape: tuple):
500502
"""Returning true if vector shape is row space e.g. shape = (1,4)"""

‎main.py‎

Lines changed: 61 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
"""A tiny shader fork written in Python 3"""
2-
3-
from progressbar import progressbar
2+
from math import pi, sin, cos
3+
import progressbar
44

55
from tiny_image import TinyImage
66
import our_gl as gl
7-
from geom import ScreenCoords, Vector3D
7+
from geom import ScreenCoords, Vector3D, Vector2D, Point2D
88
from model import ModelStorage, NormalMapType
99
from tiny_shaders import FlatShader, GouraudShader, GouraudShaderSegregated, \
1010
DiffuseGouraudShader, GlobalNormalmapShader, SpecularmapShader, \
11-
TangentNormalmapShader, DepthShader, SpecularShadowShader
11+
TangentNormalmapShader, DepthShader, SpecularShadowShader, \
12+
ZShader, AmbientOcclusionShader
1213

1314
if __name__ == "__main__":
1415
# Model property selection
@@ -38,14 +39,18 @@
3839
# Image property selection
3940
IMG_PROP_SET = 1
4041
if IMG_PROP_SET == 0:
42+
(w, h) = (200, 200)
43+
elif IMG_PROP_SET == 1:
44+
(w, h) = (800, 800)
45+
elif IMG_PROP_SET == 2:
4146
(w, h) = (2000, 2000)
4247
else:
43-
(w, h) = (800, 800)
48+
raise ValueError
4449

4550
image = TinyImage(w, h)
4651

4752
# View property selection
48-
VIEW_PROP_SET = 0
53+
VIEW_PROP_SET = 1
4954
if VIEW_PROP_SET == 0:
5055
EYE = Vector3D(0, 0, 4) # Lookat camera 'EYE' position
5156
CENTER = Vector3D(0, 0, 0) # Lookat 'CENTER'. 'EYE' looks at CENTER
@@ -106,30 +111,44 @@
106111
shadow_image = None
107112
M_sb = None
108113

109-
SHADER_PROP_SET = 6
114+
SHADER_PROP_SET = 8
110115
if SHADER_PROP_SET == 0:
116+
AO_RUN = False
111117
WRITE_SHADOW_BUFFER = False
112-
shader = GouraudShader(mdl, LIGHT_DIR, M_sc)
118+
shader = FlatShader(mdl, LIGHT_DIR, M_sc)
113119
elif SHADER_PROP_SET == 1:
120+
AO_RUN = False
114121
WRITE_SHADOW_BUFFER = False
115-
shader = GouraudShaderSegregated(mdl, LIGHT_DIR, M_sc, 4)
122+
shader = GouraudShader(mdl, LIGHT_DIR, M_sc)
116123
elif SHADER_PROP_SET == 2:
124+
AO_RUN = False
117125
WRITE_SHADOW_BUFFER = False
118-
shader = DiffuseGouraudShader(mdl, LIGHT_DIR, M_sc)
126+
shader = GouraudShaderSegregated(mdl, LIGHT_DIR, M_sc, 4)
119127
elif SHADER_PROP_SET == 3:
128+
AO_RUN = False
120129
WRITE_SHADOW_BUFFER = False
121-
shader = GlobalNormalmapShader(mdl, LIGHT_DIR, M_pe, M_sc, M_pe_IT)
130+
shader = DiffuseGouraudShader(mdl, LIGHT_DIR, M_sc)
122131
elif SHADER_PROP_SET == 4:
132+
AO_RUN = False
123133
WRITE_SHADOW_BUFFER = False
124-
shader = SpecularmapShader(mdl, LIGHT_DIR, M_pe, M_sc, M_pe_IT)
134+
shader = GlobalNormalmapShader(mdl, LIGHT_DIR, M_pe, M_sc, M_pe_IT)
125135
elif SHADER_PROP_SET == 5:
136+
AO_RUN = False
126137
WRITE_SHADOW_BUFFER = False
127-
shader = TangentNormalmapShader(mdl, LIGHT_DIR, M_pe, M_pe_IT, M_viewport)
138+
shader = SpecularmapShader(mdl, LIGHT_DIR, M_pe, M_sc, M_pe_IT)
128139
elif SHADER_PROP_SET == 6:
140+
AO_RUN = False
141+
WRITE_SHADOW_BUFFER = False
142+
shader = TangentNormalmapShader(mdl, LIGHT_DIR, M_pe, M_pe_IT, M_viewport)
143+
elif SHADER_PROP_SET == 7:
129144
WRITE_SHADOW_BUFFER = True
130145
shader = SpecularShadowShader(mdl, LIGHT_DIR, M_pe, M_sc, M_pe_IT, None, None)
146+
elif SHADER_PROP_SET == 8:
147+
AO_RUN = True
148+
WRITE_SHADOW_BUFFER = False
149+
shader = ZShader(mdl, M_sc)
131150
else:
132-
shader = FlatShader(mdl, LIGHT_DIR, M_sc)
151+
raise ValueError
133152

134153
# Shadow buffer run
135154
if WRITE_SHADOW_BUFFER:
@@ -147,21 +166,23 @@
147166
# Apply data to normal shader
148167
shader.uniform_M_sb = M_sb
149168
shader.shadow_buffer = shadow_buffer
150-
169+
151170
print("Saving shadow buffer ...")
152-
for face_idx in progressbar(range(mdl.get_face_count())):
171+
for face_idx in progressbar.progressbar(range(mdl.get_face_count())):
153172
for face_vert_idx in range(3):
154173
# Get transformed vertex and prepare internal shader data
155174
vert = depth_shader.vertex(face_idx, face_vert_idx)
156175
screen_coords = screen_coords.set_col(face_vert_idx, vert)
157176

158177
# Rasterize image (z heigth in dir of light). Shadow buffer is filles as well
159-
shadow_image = gl.draw_triangle(screen_coords, depth_shader, shadow_buffer, shadow_image)
178+
shadow_image = gl.draw_triangle(screen_coords, depth_shader, shadow_buffer,
179+
shadow_image)
180+
160181
shadow_image.save_to_disk("renders/shadow_buffer.png")
161182

162183
# Normal shader run
163184
print("Drawing triangles ...")
164-
for face_idx in progressbar(range(mdl.get_face_count())):
185+
for face_idx in progressbar.progressbar(range(mdl.get_face_count())):
165186
for face_vert_idx in range(3):
166187
# Get transformed vertex and prepare internal shader data
167188
vert = shader.vertex(face_idx, face_vert_idx)
@@ -170,4 +191,26 @@
170191
# Rasterize triangle
171192
image = gl.draw_triangle(screen_coords, shader, zbuffer, image)
172193

194+
# AO run
195+
if AO_RUN:
196+
print("Applying ambient occlusion ...")
197+
with progressbar.ProgressBar(max_value = w*h) as bar:
198+
for image_x in range(0, w):
199+
for image_y in range(0, h):
200+
if zbuffer[image_x][image_y] > -float('Inf'):
201+
summed_ang = 0
202+
RAY_NUM = 8
203+
for ray_angle in [2*pi*ray / RAY_NUM for ray in range(RAY_NUM)]:
204+
max_elevation = AmbientOcclusionShader.max_elevation_angle(
205+
zbuffer, w, h,
206+
Point2D(image_x, image_y),
207+
Vector2D(cos(ray_angle), sin(ray_angle)))
208+
209+
summed_ang = summed_ang + pi/2 - max_elevation
210+
211+
ao_intensity = summed_ang / (pi / 2 * RAY_NUM)
212+
ao_intensity = ao_intensity**100
213+
image.set(image_x, image_y, Vector3D(255, 255, 255) * ao_intensity // 1)
214+
bar.update(image_x * w + image_y)
215+
173216
image.save_to_disk(OUTPUT_FILENAME)

‎tiny_shaders.py‎

Lines changed: 85 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -472,19 +472,88 @@ def fragment(self, bary: Barycentric):
472472
# Do not discard pixel and return color
473473
return (False, color)
474474

475-
def get_max_elevation_angle(pt_amb: Point2D, z_buffer: list, sweep_dir: Vector2D):
476-
"""Returns max elevation angle ray-casted from starting point pt_amb."""
477-
(im_w, im_h) = len(z_buffer)
478-
fact = 1 / max(pt_amb)
479-
vect_amb = Vector3D(pt_amb.x, pt_amb.x, z_buffer[pt_amb.x][pt_amb.y])
480-
sweep_proj = vect_amb
481-
max_angle = 0
482-
483-
while 0 <= sweep_proj.x < im_w and 0 <= sweep_proj.y < im_h:
484-
sweep_proj += fact * Vector3D(sweep_dir.x, sweep_dir.y, 0)
485-
z_height = z_buffer[int(sweep_proj.x // 1)][int(sweep_proj.y // 1)]
486-
sweep = Vector3D(sweep_proj.x, sweep_proj.y, z_height)
487-
alpha = math.acos((sweep - vect_amb).normalize() * sweep_proj.normalize())
488-
max_angle = max(alpha, max_angle)
489-
490-
return max_angle
475+
class ZShader(gl.Shader):
476+
"""Shader used to save shadow buffer."""
477+
mdl: ModelStorage
478+
479+
uniform_M: Matrix4D
480+
481+
def __init__(self, mdl, M):
482+
self.mdl = mdl
483+
self.uniform_M = M # pylint: disable=invalid-name
484+
485+
def vertex(self, face_idx: int, vert_idx: int):
486+
vert = self.mdl.get_vertex(face_idx, vert_idx) # Read the vertex
487+
return transform_vertex_to_screen(vert, self.uniform_M)
488+
489+
def fragment(self, bary: Barycentric):
490+
# Do not return color. This is shader is used passively to create
491+
# generate a screen space shadow buffer
492+
return (True, None)
493+
494+
495+
class AmbientOcclusionShader(gl.Shader):
496+
"""Shader used to compute ambient occlusion local illumination."""
497+
mdl: ModelStorage
498+
499+
# Vertices are stored col-wise
500+
varying_vert = Matrix3D(9*[0])
501+
502+
uniform_M: Matrix4D
503+
uniform_precalc_zbuffer: list
504+
uniform_zbuffer_width: int
505+
uniform_zbuffer_height: int
506+
507+
def __init__(self, mdl, M, percalc_zbuffer, zbuffer_width, zbuffer_height):
508+
self.mdl = mdl
509+
self.uniform_M = M # pylint: disable=invalid-name
510+
self.uniform_precalc_zbuffer = percalc_zbuffer
511+
self.uniform_zbuffer_width = zbuffer_width
512+
self.uniform_zbuffer_height = zbuffer_height
513+
514+
def vertex(self, face_idx: int, vert_idx: int):
515+
vert = self.mdl.get_vertex(face_idx, vert_idx) # Read the vertex
516+
vert = transform_vertex_to_screen(vert, self.uniform_M)
517+
self.varying_vert = self.varying_vert.set_col(vert_idx, vert)
518+
return vert
519+
520+
def fragment(self, bary: Barycentric):
521+
(x_sc, y_sc, _) = self.varying_vert * bary // 1
522+
523+
summed_ang = 0
524+
RAY_NUM = 8
525+
for ray_angle in [2*math.pi*ray / RAY_NUM for ray in range(RAY_NUM)]:
526+
max_elevation = self.max_elevation_angle(
527+
self.uniform_precalc_zbuffer,
528+
self.uniform_zbuffer_width, self.uniform_zbuffer_height,
529+
Point2D(x_sc, y_sc),
530+
Vector2D(math.cos(ray_angle), math.sin(ray_angle)))
531+
532+
summed_ang = summed_ang + (math.pi / 2 - max_elevation)
533+
534+
ao_intensity = summed_ang / (math.pi/2 * RAY_NUM)
535+
ao_intensity = ao_intensity**10
536+
color = Vector3D(255, 255, 255) * ao_intensity // 1
537+
return (False, color)
538+
539+
@staticmethod
540+
def max_elevation_angle(zbuffer: list, image_width, image_height,
541+
pt_amb: Point2D, sweep_dir: Vector2D):
542+
"""Returns max elevation angle ray-casted from starting point pt_amb."""
543+
pt_amb_z = zbuffer[pt_amb.x][pt_amb.y]
544+
sweep = Vector2D(pt_amb)
545+
if abs(sweep_dir.x) > abs(sweep_dir.y):
546+
max_sweep_comp = abs(sweep_dir.x)
547+
else:
548+
max_sweep_comp = abs(sweep_dir.y)
549+
sweep_delta = 1.0 / max_sweep_comp * sweep_dir
550+
max_tan = 0
551+
while True:
552+
z_height = zbuffer[int(sweep.x)][int(sweep.y)]
553+
if z_height > pt_amb_z:
554+
tan_val = (z_height - pt_amb_z) / sweep.abs()
555+
max_tan = max(tan_val, max_tan)
556+
sweep += sweep_delta
557+
if not 0 <= sweep.x < (image_width) or not 0 <= sweep.y < (image_height):
558+
break
559+
return math.atan(max_tan)

0 commit comments

Comments
 (0)
0