-
Notifications
You must be signed in to change notification settings - Fork 454
/
FlowRendererUtility.py
145 lines (117 loc) · 7.65 KB
/
FlowRendererUtility.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""Provides functionality to render an optical flow image."""
import os
from typing import Dict, List, Union
import bpy
import numpy as np
from blenderproc.python.utility.BlenderUtility import load_image
from blenderproc.python.renderer import RendererUtility
from blenderproc.python.utility.Utility import Utility, UndoAfterExecution
from blenderproc.python.writer.WriterUtility import _WriterUtility
def render_optical_flow(output_dir: str = None, temp_dir: str = None, get_forward_flow: bool = True,
get_backward_flow: bool = True, blender_image_coordinate_style: bool = False,
forward_flow_output_file_prefix: str = "forward_flow_",
forward_flow_output_key: str = "forward_flow",
backward_flow_output_file_prefix: str = "backward_flow_",
backward_flow_output_key: str = "backward_flow", return_data: bool = True,
verbose: bool = False) -> \
Dict[str, Union[np.ndarray, List[np.ndarray]]]:
""" Renders the optical flow (forward and backward) for all frames.
:param output_dir: The directory to write images to.
:param temp_dir: The directory to write intermediate data to.
:param get_forward_flow: Whether to render forward optical flow.
:param get_backward_flow: Whether to render backward optical flow.
:param blender_image_coordinate_style: Whether to specify the image coordinate system at the bottom left
(blender default; True) or top left (standard convention; False).
:param forward_flow_output_file_prefix: The file prefix that should be used when writing forward flow to a file.
:param forward_flow_output_key: The key which should be used for storing forward optical flow values.
:param backward_flow_output_file_prefix: The file prefix that should be used when writing backward flow to a file.
:param backward_flow_output_key: The key which should be used for storing backward optical flow values.
:param return_data: Whether to load and return generated data.
:param verbose: If True, more details about the rendering process are printed.
:return: dict of lists of raw renderer outputs. Keys can be 'forward_flow', 'backward_flow'
"""
if get_forward_flow is False and get_backward_flow is False:
raise RuntimeError("Take the FlowRenderer Module out of the config if both forward and "
"backward flow are set to False!")
if output_dir is None:
output_dir = Utility.get_temporary_directory()
if temp_dir is None:
temp_dir = Utility.get_temporary_directory()
with UndoAfterExecution():
RendererUtility.render_init()
# the amount of samples must be one and there can not be any noise threshold
RendererUtility.set_max_amount_of_samples(1)
RendererUtility.set_noise_threshold(0)
RendererUtility.set_denoiser(None)
RendererUtility.set_light_bounces(1, 0, 0, 1, 0, 8, 0)
_FlowRendererUtility.output_vector_field(get_forward_flow, get_backward_flow, output_dir)
# only need to render once; both fwd and bwd flow will be saved
temporary_fwd_flow_file_path = os.path.join(temp_dir, 'fwd_flow_')
temporary_bwd_flow_file_path = os.path.join(temp_dir, 'bwd_flow_')
print(f"Rendering {bpy.context.scene.frame_end - bpy.context.scene.frame_start} frames of optical flow...")
RendererUtility.render(temp_dir, "img_flow_temp_ignore_me_", None, load_keys=set(), verbose=verbose)
# After rendering: convert to optical flow or calculate hsv visualization, if desired
for frame in range(bpy.context.scene.frame_start, bpy.context.scene.frame_end):
# temporarily save respective vector fields
if get_forward_flow:
file_path = temporary_fwd_flow_file_path + f"{frame:04d}" + ".exr"
fwd_flow_field = load_image(file_path, num_channels=4).astype(np.float32)
if not blender_image_coordinate_style:
fwd_flow_field[:, :, 1] = fwd_flow_field[:, :, 1] * -1
file_name = os.path.join(output_dir, forward_flow_output_file_prefix) + f"{frame:04d}"
forward_flow = fwd_flow_field * -1 # invert forward flow to point at next frame
np.save(file_name + '.npy', forward_flow[:, :, :2])
if get_backward_flow:
file_path = temporary_bwd_flow_file_path + f"{frame:04d}" + ".exr"
bwd_flow_field = load_image(file_path, num_channels=4).astype(np.float32)
if not blender_image_coordinate_style:
bwd_flow_field[:, :, 1] = bwd_flow_field[:, :, 1] * -1
file_name = os.path.join(output_dir, backward_flow_output_file_prefix) + f"{frame:04d}"
np.save(file_name + '.npy', bwd_flow_field[:, :, :2])
load_keys = set()
# register desired outputs
if get_forward_flow:
Utility.register_output(output_dir, forward_flow_output_file_prefix, forward_flow_output_key, '.npy', '2.0.0')
load_keys.add(forward_flow_output_key)
if get_backward_flow:
Utility.register_output(output_dir, backward_flow_output_file_prefix, backward_flow_output_key, '.npy', '2.0.0')
load_keys.add(backward_flow_output_key)
return _WriterUtility.load_registered_outputs(load_keys) if return_data else {}
class _FlowRendererUtility():
@staticmethod
def output_vector_field(forward_flow: bool, backward_flow: bool, output_dir: str):
""" Configures compositor to output speed vectors.
:param forward_flow: Whether to render forward optical flow.
:param backward_flow: Whether to render backward optical flow.
:param output_dir: The directory to write images to.
"""
# Flow settings (is called "vector" in blender)
bpy.context.scene.render.use_compositing = True
bpy.context.scene.use_nodes = True
bpy.context.view_layer.use_pass_vector = True
# Adapt compositor to output vector field
tree = bpy.context.scene.node_tree
links = tree.links
# Use existing render layer
render_layer_node = tree.nodes.get('Render Layers')
separate_rgba = tree.nodes.new('CompositorNodeSepRGBA')
links.new(render_layer_node.outputs['Vector'], separate_rgba.inputs['Image'])
if forward_flow:
combine_fwd_flow = tree.nodes.new('CompositorNodeCombRGBA')
links.new(separate_rgba.outputs['B'], combine_fwd_flow.inputs['R'])
links.new(separate_rgba.outputs['A'], combine_fwd_flow.inputs['G'])
fwd_flow_output_file = tree.nodes.new('CompositorNodeOutputFile')
fwd_flow_output_file.base_path = output_dir
fwd_flow_output_file.format.file_format = "OPEN_EXR"
fwd_flow_output_file.file_slots.values()[0].path = "fwd_flow_"
links.new(combine_fwd_flow.outputs['Image'], fwd_flow_output_file.inputs['Image'])
if backward_flow:
# actually need to split - otherwise the A channel of the image is getting weird, no idea why
combine_bwd_flow = tree.nodes.new('CompositorNodeCombRGBA')
links.new(separate_rgba.outputs['R'], combine_bwd_flow.inputs['R'])
links.new(separate_rgba.outputs['G'], combine_bwd_flow.inputs['G'])
bwd_flow_output_file = tree.nodes.new('CompositorNodeOutputFile')
bwd_flow_output_file.base_path = output_dir
bwd_flow_output_file.format.file_format = "OPEN_EXR"
bwd_flow_output_file.file_slots.values()[0].path = "bwd_flow_"
links.new(combine_bwd_flow.outputs['Image'], bwd_flow_output_file.inputs['Image'])