File size: 11,258 Bytes
cd47a59
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
import pip

# pip.main(["install", "pandas"]) # Unused and crashes with numpy alias
pip.main(["install", "tqdm"])

import bpy
import sys
import numpy as np
# Numpy 2.0 backward compatibility for Blender 3.6
sys.modules['numpy._core'] = sys.modules.get('numpy.core', np)
if 'numpy.core.multiarray' in sys.modules:
    sys.modules['numpy._core.multiarray'] = sys.modules['numpy.core.multiarray']
# import pandas as pd # Unused
from pathlib import Path
import os
from ast import literal_eval
from tqdm import tqdm
from contextlib import contextmanager
import pathlib

this_script_path = pathlib.Path(__file__).parent.resolve()
W_FACE_AND_COLOR_FILE = this_script_path / "blend" / "smpl_mesh_info.npy"

FORMAT_LDR = "PNG"
COLOR_DEPTH_LDR = 8
SAMPLES = 1
COLOR_MODE = "RGB"

def setup_device(use_id):
    bpy.context.scene.render.engine = "CYCLES"
    bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA"
    bpy.context.scene.cycles.device = "GPU"
    bpy.context.preferences.addons["cycles"].preferences.get_devices()
    print(bpy.context.preferences.addons["cycles"].preferences.compute_device_type)
    for i, device in enumerate(
        bpy.context.preferences.addons["cycles"].preferences.devices
    ):
        if i == use_id or "CPU" in device["name"]:
            device["use"] = True  # Using all devices, include GPU and CPU
        else:
            device["use"] = False  # Using all devices, include GPU and CPU
        print(device["name"], "USE:", bool(device["use"]))

class SingleDataset:
    def __init__(self, smpl_folder, smpl_suffixes=["npy", "npz"]):
        self.smpl_folder = Path(smpl_folder)
        self.out_folder = Path(smpl_folder).parent
        self.smpl_paths = []
        self.bboxes = []
        self.valid_index = []
        folder = self.out_folder
        self.smpl_paths = sorted(
            [
                path for i in smpl_suffixes
                for path in (folder / "smpl_results").glob("*." + i)
            ]
        )
        self.output_paths = [self.out_folder for smpl_path in self.smpl_paths]
        # Skip finished smpl_path. Enable it if want to continue processing only remaining imgs.
        smpl_fns = [
            os.path.splitext(os.path.basename(smpl_path))[0]
            for smpl_path in self.smpl_paths
        ]  # Example smpl_fn:0000
        imgs_output_path = [
            os.path.join(str(self.output_paths[i]), "visualized_imgs", f"{smpl_fn}.png")
            for i, smpl_fn in enumerate(smpl_fns)
        ]
        imgs_already_exist = [os.path.exists(smpl_path) for smpl_path in imgs_output_path]
        imgs_index_to_inference = np.where(np.array(imgs_already_exist) == False)[0]
        smpl_paths_copy = list(self.smpl_paths)
        output_paths_copy = list(self.output_paths)
        self.smpl_paths = [smpl_paths_copy[img_index] for img_index in imgs_index_to_inference]
        self.output_paths = [output_paths_copy[img_index] for img_index in imgs_index_to_inference]
        print(f"finish loading {imgs_index_to_inference.shape[0]} frames data, \
              skip {len(smpl_fns) - imgs_index_to_inference.shape[0]} existing images")


def load_smpl(smpl_path):
    return np.load(smpl_path, allow_pickle=True).item()


@contextmanager
def stdout_redirected(to=os.devnull):
    """
    Redirects stdout to a specified file.

    Usage:
    with stdout_redirected(to=filename):
        print("from Python")
        os.system("echo non-Python applications are also supported")
    """
    fd = sys.stdout.fileno()

    # Save a copy of the original stdout file descriptor
    original_stdout_fd = os.dup(fd)

    # Redirect stdout to the specified file
    with open(to, 'w') as file:
        os.dup2(file.fileno(), fd)

    try:
        yield
    finally:
        # Restore the original stdout
        os.dup2(original_stdout_fd, fd)
        os.close(original_stdout_fd)

def rendering_pipeline(dataset, ref_img_path):
    scene = bpy.context.scene
    scene.render.image_settings.file_format = str(FORMAT_LDR)
    scene.render.image_settings.color_depth = str(COLOR_DEPTH_LDR)
    scene.render.image_settings.color_mode = str(COLOR_MODE)
    scene.render.resolution_percentage = 100
    scene.render.use_persistent_data = True
    scene.cycles.use_denoising = False

    camera = bpy.data.objects["Camera"]
    camera.data.clip_start = 0.05
    camera.data.clip_end = 1e12
    camera.data.cycles.samples = SAMPLES
    scene.cycles.samples = SAMPLES
    bpy.context.scene.view_layers["ViewLayer"].use_pass_z = True
    render_layers = bpy.context.scene.view_layers
    if "mesh_collection" not in bpy.data.collections.keys():
        mesh_collection = bpy.data.collections.new("mesh_collection")
        bpy.context.scene.collection.children.link(mesh_collection)
    mesh_collection = bpy.data.collections.get("mesh_collection")
    mat_semantic = bpy.data.materials.get("Semantic")
    mat_normal = bpy.data.materials.get("Normal")

    result_dict = np.load(W_FACE_AND_COLOR_FILE, allow_pickle=True).item()
    faces = result_dict["faces"]
    verts_color = result_dict["verts_color"][:, ::-1].astype(np.float32) / 255.0

    result_dict_list = []

    processed = []
    for path in tqdm(dataset.smpl_paths, total=len(dataset.smpl_paths), desc="Loading smpls into RAM"):
        result = load_smpl(path)
        processed.append(result)

    for smpl in tqdm(processed, total=len(dataset.smpl_paths), desc="Loading smpls into RAM"):
        result_dict_list.append(smpl)

    for smpl_path, output_path, result_dict in tqdm(
        zip(dataset.smpl_paths, dataset.output_paths, result_dict_list),
        total=len(dataset.output_paths),
        desc="Rendering Images",
        miniters=10,
    ):
        render_path = output_path
        smpl_fn, _ = os.path.splitext(os.path.basename(smpl_path))
        smpl_fn = smpl_fn.split(".")[0]
        ori_img = bpy.data.images.load(ref_img_path)
        ori_img.name = "ori_img.png"
        ori_img.colorspace_settings.name = "Raw"
        bpy.data.scenes["Scene"].node_tree.nodes["Image"].image = ori_img

        cam_t = result_dict["cam_t"][0]
        verts = result_dict["verts"][0] + cam_t
        img_size = result_dict["render_res"].astype(np.int32)
        camera.data.sensor_width = img_size.max()
        camera.data.lens = result_dict["scaled_focal_length"]
        scene.render.resolution_x = img_size[0]
        scene.render.resolution_y = img_size[1]

        # make mesh
        new_mesh = bpy.data.meshes.new("smpl_mesh")
        new_mesh.from_pydata(verts, edges=[], faces=faces)
        # make object from mesh
        new_object = bpy.data.objects.new("new_object", new_mesh)
        # make collection
        new_object.rotation_euler[0] = -np.pi / 2

        for f in new_object.data.polygons:
            f.use_smooth = True

        # add object to scene collection
        mesh_collection.objects.link(new_object)

        # Render Normal Map and Depth Map
        bpy.data.scenes["Scene"].node_tree.nodes["Depth Output"].base_path = (
            os.path.join(render_path, "depth")
        )
        bpy.data.scenes["Scene"].node_tree.nodes["Visualize Output"].base_path = (
            os.path.join(render_path, "visualized_imgs")
        )
        bpy.data.scenes["Scene"].node_tree.nodes["Mask Output"].base_path = (
            os.path.join(render_path, "mask")
        )

        new_object.data.materials.append(mat_normal)

        scene.view_settings.view_transform = "Raw"
        bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[0].default_value = (0, 0, 0, 1)
        output_name = f"{smpl_fn}.png"
        depth_path = os.path.join(render_path, "depth", output_name)
        normal_path = os.path.join(render_path, "normal", output_name)
        vis_path = os.path.join(render_path, "visualized_imgs", output_name)
        mask_path = os.path.join(render_path, "mask", output_name)
        semantic_path = os.path.join(render_path, "semantic_map", output_name)
        scene.render.filepath = normal_path
        for layer in render_layers:
            # some condition
            layer.use = layer.name == "ViewLayer"
        bpy.context.scene.render.film_transparent = True
        with stdout_redirected():
            bpy.ops.render.render(write_still=True)

        if os.path.isfile(depth_path):
            os.remove(depth_path)
        if os.path.isfile(os.path.join(render_path, "depth", f"{0:04d}.png")):
            os.rename(os.path.join(render_path, "depth", f"{0:04d}.png"), depth_path)

        if os.path.isfile(vis_path):
            os.remove(vis_path)
        if os.path.isfile(os.path.join(render_path, "visualized_imgs", f"{0:04d}.png")):
            os.rename(
                os.path.join(render_path, "visualized_imgs", f"{0:04d}.png"), vis_path
            )

        if os.path.isfile(mask_path):
            os.remove(mask_path)
        if os.path.isfile(os.path.join(render_path, "mask", f"{0:04d}.png")):
            os.rename(os.path.join(render_path, "mask", f"{0:04d}.png"), mask_path)

        # This is to reference the vertex color layer later
        vertex_colors_name = "vert_colors"
        # Here the color layer is made on the mesh
        new_mesh.vertex_colors.new(name=vertex_colors_name)
        color_layer = new_mesh.vertex_colors[vertex_colors_name]

        # We loop over all the polygons
        for poly in new_mesh.polygons:
            # We get the polygon index and the corresponding mesh index
            for vert_i_poly, vert_i_mesh in enumerate(poly.vertices):
                # We get the loop index from the polygon index
                vert_i_loop = poly.loop_indices[vert_i_poly]
                # We set the color for the vertex
                rgb = verts_color[vert_i_mesh].tolist()
                rgb.append(1)
                color_layer.data[vert_i_loop].color = rgb
        del rgb
        mat_semantic.node_tree.nodes["Color Attribute"].layer_name = "vert_colors"
        new_object.data.materials[0] = mat_semantic
        scene.view_settings.view_transform = "Standard"
        scene.render.filepath = semantic_path

        for layer in render_layers:
            # some condition
            layer.use = layer.name == "VertexColor Layer"
        bpy.context.scene.render.film_transparent = False
        with stdout_redirected():
            bpy.ops.render.render(write_still=True)

        bpy.data.objects.remove(new_object, do_unlink=True)
        bpy.data.meshes.remove(new_mesh)
        bpy.data.images.remove(ori_img)
        if os.path.isfile(os.path.join(render_path, "visualized_imgs", f"{0:04d}.png")):
            os.remove(os.path.join(render_path, "visualized_imgs", f"{0:04d}.png"))
        del new_mesh
        del ori_img
        del new_object
    del result_dict_list


if __name__ == "__main__":
    import sys

    argv = sys.argv
    print(f"Rendering:")
    try:
        argv.index("--device")
    except:
        print("Use Only CPU for Rendering")
    else:
        setup_device(int(argv[argv.index("--device") + 1]))

    smpl_folder = os.path.abspath(argv[argv.index("--driving_path") + 1])
    ref_img_path = os.path.abspath(argv[argv.index("--reference_path") + 1])
    rendering_pipeline(SingleDataset(smpl_folder), ref_img_path)