File size: 5,692 Bytes
084ab29
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import multiprocessing
import tempfile
from multiprocessing import Process
from typing import Optional

from tqdm import tqdm


def render_hint_images(model_path, fov, pls, power=500., geo_smooth=True, output_folder: Optional[str] = None,
                       env_map: Optional[str] = None, resolution=512, use_gpu=False):
    import bpy
    import numpy as np

    from bpy_helper.camera import create_camera
    from bpy_helper.light import set_env_light, create_point_light
    from bpy_helper.material import create_white_diffuse_material, create_specular_ggx_material
    from bpy_helper.scene import reset_scene, import_3d_model
    from bpy_helper.utils import stdout_redirected

    def configure_blender():
        # Set the render resolution
        bpy.context.scene.render.resolution_x = resolution
        bpy.context.scene.render.resolution_y = resolution
        bpy.context.scene.render.engine = 'CYCLES'
        bpy.context.scene.cycles.samples = 512
        if use_gpu:
            bpy.context.preferences.addons["cycles"].preferences.get_devices()
            bpy.context.scene.cycles.device = 'GPU'
            bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'

        # Enable the alpha channel for GT mask
        bpy.context.scene.render.film_transparent = True
        bpy.context.scene.render.image_settings.color_mode = 'RGBA'

    def render_rgb_and_hint(output_path):
        MAT_DICT = {
            '_diffuse': create_white_diffuse_material(),
            '_ggx0.05': create_specular_ggx_material(0.05),
            '_ggx0.13': create_specular_ggx_material(0.13),
            '_ggx0.34': create_specular_ggx_material(0.34),
        }

        # render
        for mat_name, mat in MAT_DICT.items():
            bpy.context.scene.view_layers["ViewLayer"].material_override = mat

            # and png
            bpy.context.scene.render.image_settings.file_format = 'PNG'
            bpy.context.scene.render.filepath = f'{output_path}{mat_name}.png'
            bpy.ops.render.render(animation=False, write_still=True)

    # Render hints
    reset_scene()
    import_3d_model(model_path)
    if geo_smooth:
        for obj in bpy.data.objects:
            if obj.type == 'MESH':
                obj.modifiers.new("Smooth", type="SMOOTH")
                smooth_modifier = obj.modifiers["Smooth"]
                smooth_modifier.factor = 0.5
                smooth_modifier.iterations = 150
    configure_blender()

    c2w = np.array([
        [-1, 0, 0, 0],
        [0, 0, 1, 0],
        [0, 1, 0, 0],
        [0, 0, 0, 0]
    ])
    camera = create_camera(c2w, fov)
    bpy.context.scene.camera = camera
    if output_folder is None:
        output_folder = tempfile.mkdtemp()
    for i in tqdm(range(len(pls)), desc='Rendering Hints'):
        if env_map:
            z_angle = i / len(pls) * np.pi * 2.
            set_env_light(env_map, rotation_euler=[0, 0, z_angle])
        else:
            pl_pos = pls[i]
            _point_light = create_point_light(pl_pos, power)

        with stdout_redirected():
            render_rgb_and_hint(output_folder + f'/hint{i:02d}')

    return output_folder


def render_bg_images(fov, pls, output_folder: Optional[str] = None, env_map: Optional[str] = None, resolution=512):
    import bpy
    import numpy as np

    from bpy_helper.camera import create_camera
    from bpy_helper.light import set_env_light
    from bpy_helper.scene import reset_scene
    from bpy_helper.utils import stdout_redirected

    def configure_blender():
        # Set the render resolution
        bpy.context.scene.render.resolution_x = resolution
        bpy.context.scene.render.resolution_y = resolution
        bpy.context.scene.render.engine = 'CYCLES'
        bpy.context.scene.cycles.samples = 512

        # Enable the alpha channel for GT mask
        bpy.context.scene.render.film_transparent = False
        bpy.context.scene.render.image_settings.color_mode = 'RGB'

    def render_env_bg(output_path):
        bpy.context.scene.view_layers["ViewLayer"].material_override = None
        bpy.context.scene.render.image_settings.file_format = 'PNG'
        bpy.context.scene.render.filepath = f'{output_path}.png'
        bpy.ops.render.render(animation=False, write_still=True)

    # Render backgrounds
    reset_scene()
    configure_blender()

    c2w = np.array([
        [-1, 0, 0, 0],
        [0, 0, 1, 0],
        [0, 1, 0, 0],
        [0, 0, 0, 0]
    ])
    camera = create_camera(c2w, fov)
    bpy.context.scene.camera = camera
    if output_folder is None:
        output_folder = tempfile.mkdtemp()
    for i in tqdm(range(len(pls)), desc='Rendering Env Backgrounds'):
        z_angle = i / len(pls) * np.pi * 2.
        set_env_light(env_map, rotation_euler=[0, 0, z_angle])

        with stdout_redirected():
            render_env_bg(output_folder + f'/bg{i:02d}')

    return output_folder


def render_hint_images_wrapper(model_path, fov, pls, power, geo_smooth, output_folder, env_map, resolution, return_dict):
    output_folder = render_hint_images(model_path, fov, pls, power, geo_smooth, output_folder, env_map, resolution)
    return_dict['output_folder'] = output_folder


def render_hint_images_btn_func(model_path, fov, pls, power=500., geo_smooth=True, output_folder: Optional[str] = None,
                                env_map: Optional[str] = None, resolution=512):
    manager = multiprocessing.Manager()
    return_dict = manager.dict()
    p = Process(target=render_hint_images_wrapper, args=(model_path, fov, pls, power, geo_smooth, output_folder, env_map, resolution, return_dict))
    p.start()
    p.join()
    return return_dict['output_folder']