init

.gitignore

@@ -0,0 +1,154 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintainted in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+
+output/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 NCJ
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,15 @@
 ---
 title: DiLightNet
-emoji: 🏆
+emoji: 💡
 colorFrom: indigo
 colorTo: red
 sdk: gradio
 sdk_version: 4.27.0
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# DiLightNet: Fine-grained Lighting Control for Diffusion-based Image Generation
+
+https://arxiv.org/abs/2402.11929

app.py

@@ -0,0 +1,131 @@
+import gradio as gr
+import imageio
+import numpy as np
+
+from demo.img_gen import img_gen
+from demo.mesh_recon import mesh_reconstruction
+from demo.relighting_gen import relighting_gen
+from demo.render_hints import render_hint_images_btn_func
+from demo.rm_bg import rm_bg
+
+
+with gr.Blocks(title="DiLightNet Demo") as demo:
+    gr.Markdown("# DiLightNet: Fine-grained Lighting Control for Image Diffusion")
+
+    with gr.Row():
+        # 1. Reference Image Input / Generation
+        with gr.Column(variant="panel"):
+            gr.Markdown("## Step 1. Input or Generate Reference Image")
+            input_image = gr.Image(height=512, width=512, label="Input Image", interactive=True)
+            with gr.Accordion("Generate Image", open=False):
+                with gr.Group():
+                    prompt = gr.Textbox(value="", label="Prompt", lines=3, placeholder="Input prompt here")
+                    with gr.Row():
+                        seed = gr.Number(value=42, label="Seed", interactive=True)
+                        steps = gr.Number(value=20, label="Steps", interactive=True)
+                        cfg = gr.Number(value=7.5, label="CFG", interactive=True)
+                        down_from_768 = gr.Checkbox(label="Downsample from 768", value=True)
+                with gr.Row():
+                    generate_btn = gr.Button(value="Generate")
+                    generate_btn.click(fn=img_gen, inputs=[prompt, seed, steps, cfg, down_from_768], outputs=[input_image])
+
+        # 2. Background Removal
+        with gr.Column(variant="panel"):
+            gr.Markdown("## Step 2. Remove Background")
+            with gr.Tab("Masked Image"):
+                masked_image = gr.Image(height=512, width=512, label="Masked Image", interactive=True)
+            with gr.Tab("Mask"):
+                mask = gr.Image(height=512, width=512, label="Mask", interactive=False)
+            use_sam = gr.Checkbox(label="Use SAM for Refinement", value=False)
+            rm_bg_btn = gr.Button(value="Remove Background")
+            rm_bg_btn.click(fn=rm_bg, inputs=[input_image, use_sam], outputs=[masked_image, mask])
+
+        # 3. Depth Estimation & Mesh Reconstruction
+        with gr.Column(variant="panel"):
+            gr.Markdown("## Step 3. Depth Estimation & Mesh Reconstruction")
+            mesh = gr.Model3D(label="Mesh Reconstruction", clear_color=(1.0, 1.0, 1.0, 1.0), interactive=True)
+            with gr.Column():
+                with gr.Accordion("Options", open=False):
+                    with gr.Group():
+                        remove_edges = gr.Checkbox(label="Remove Occlusion Edges", value=False)
+                        fov = gr.Number(value=55., label="FOV", interactive=True)
+                        mask_threshold = gr.Slider(value=25., label="Mask Threshold", minimum=0., maximum=255., step=1.)
+                depth_estimation_btn = gr.Button(value="Estimate Depth")
+                depth_estimation_btn.click(
+                    fn=mesh_reconstruction,
+                    inputs=[masked_image, mask, remove_edges, fov, mask_threshold],
+                    outputs=[mesh]
+                )
+
+    gr.Markdown("## Step 4. Render Hints")
+    with gr.Row():
+        with gr.Column():
+            hint_image = gr.Image(label="Hint Image")
+        with gr.Column():
+            pl_pos_x = gr.Slider(value=3., label="Point Light X", minimum=-5., maximum=5., step=0.01)
+            pl_pos_y = gr.Slider(value=1., label="Point Light Y", minimum=-5., maximum=5., step=0.01)
+            pl_pos_z = gr.Slider(value=3., label="Point Light Z", minimum=-5., maximum=5., step=0.01)
+            power = gr.Slider(value=1000., label="Point Light Power", minimum=0., maximum=2000., step=1.)
+            render_btn = gr.Button(value="Render Hints")
+            res_folder_path = gr.Textbox("", visible=False)
+
+            def render_wrapper(mesh, fov, pl_pos_x, pl_pos_y, pl_pos_z, power,
+                               progress=gr.Progress(track_tqdm=True)):
+                res_path = render_hint_images_btn_func(mesh, fov, [(pl_pos_x, pl_pos_y, pl_pos_z)], power)
+                hint_files = [res_path + '/hint00' + mat for mat in ["_diffuse.png", "_ggx0.34.png"]]
+                hints = []
+                for hint_file in hint_files:
+                    hint = imageio.v3.imread(hint_file)
+                    hints.append(hint)
+                hints = np.concatenate(hints, axis=1)
+                return hints, res_path
+
+            render_btn.click(
+                fn=render_wrapper,
+                inputs=[mesh, fov, pl_pos_x, pl_pos_y, pl_pos_z, power],
+                outputs=[hint_image, res_folder_path]
+            )
+
+    gr.Markdown("## Step 5. Relighting!")
+    with gr.Row():
+        res_image = gr.Image(label="Result Image")
+        with gr.Column():
+            with gr.Group():
+                relighting_prompt = gr.Textbox(value="", label="Relighting Text Prompt", lines=3,
+                                               placeholder="Input prompt here",
+                                               interactive=True)
+                reuse_btn = gr.Button(value="Reuse Image Generation Prompt")
+                reuse_btn.click(fn=lambda x: x, inputs=[prompt], outputs=[relighting_prompt])
+                with gr.Accordion("Options", open=False):
+                    with gr.Row():
+                        relighting_seed = gr.Number(value=3407, label="Seed", interactive=True)
+                        relighting_steps = gr.Number(value=20, label="Steps", interactive=True)
+                        relighting_cfg = gr.Number(value=3.0, label="CFG", interactive=True)
+            with gr.Row():
+                relighting_generate_btn = gr.Button(value="Generate")
+
+            def gen_relighting_image(masked_image, mask, res_folder_path, relighting_prompt, relighting_seed,
+                                     relighting_steps, relighting_cfg,
+                                     progress=gr.Progress(track_tqdm=True)):
+                relighting_gen(
+                    masked_ref_img=masked_image,
+                    mask=mask,
+                    cond_path=res_folder_path,
+                    frames=1,
+                    prompt=relighting_prompt,
+                    steps=int(relighting_steps),
+                    seed=int(relighting_seed),
+                    cfg=relighting_cfg
+                )
+                res = imageio.v3.imread(res_folder_path + '/relighting00.png')
+                return res
+
+
+            relighting_generate_btn.click(fn=gen_relighting_image,
+                                          inputs=[masked_image, mask, res_folder_path, relighting_prompt, relighting_seed,
+                                                  relighting_steps, relighting_cfg],
+                                          outputs=[res_image])
+
+
+if __name__ == '__main__':
+    demo.queue().launch(server_name="0.0.0.0", share=True)

demo/img_gen.py

@@ -0,0 +1,26 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+
+
+model_id = "stabilityai/stable-diffusion-2-1"
+
+device = torch.device('cpu')
+dtype = torch.float32
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    dtype = torch.float16
+
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to(device)
+
+
+def img_gen(prompt, seed, steps, cfg, down_from_768=False, progress=gr.Progress(track_tqdm=True)):
+    generator = torch.Generator(device=device).manual_seed(int(seed))
+    hw = 512 if not down_from_768 else 768
+    image = pipe(prompt, generator=generator, num_inference_steps=int(steps), guidance_scale=cfg, output_type='np', height=hw, width=hw).images[0]
+    if down_from_768:
+        image = F.interpolate(torch.from_numpy(image)[None].permute(0, 3, 1, 2), size=(512, 512), mode='bilinear', align_corners=False, antialias=True).permute(0, 2, 3, 1)[0].cpu().numpy()
+    return image

demo/mesh_recon.py

@@ -0,0 +1,129 @@
+import tempfile
+
+import numpy as np
+import torch
+import trimesh
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+# use torch hub
+model = torch.hub.load("isl-org/ZoeDepth", "ZoeD_NK", pretrained=True).to(device).eval()
+
+
+def get_intrinsics(H, W, fov=55.):
+    """
+    Intrinsics for a pinhole camera model.
+    Assume central principal point.
+    """
+    f = 0.5 * W / np.tan(0.5 * fov * np.pi / 180.0)
+    cx = 0.5 * W
+    cy = 0.5 * H
+    return np.array([[f, 0, cx],
+                     [0, f, cy],
+                     [0, 0, 1]])
+
+
+def depth_to_points(depth, R=None, t=None, fov=55.):
+    K = get_intrinsics(depth.shape[1], depth.shape[2], fov=fov)
+    Kinv = np.linalg.inv(K)
+    if R is None:
+        R = np.eye(3)
+    if t is None:
+        t = np.zeros(3)
+
+    # M converts from your coordinate to PyTorch3D's coordinate system
+    M = np.eye(3)
+    M[0, 0] = -1.0
+    M[1, 1] = -1.0
+
+    height, width = depth.shape[1:3]
+
+    x = np.arange(width)
+    y = np.arange(height)
+    coord = np.stack(np.meshgrid(x, y), -1)
+    coord = np.concatenate((coord, np.ones_like(coord)[:, :, [0]]), -1)  # z=1
+    coord = coord.astype(np.float32)
+    coord = coord[None]  # bs, h, w, 3
+
+    D = depth[:, :, :, None, None]
+    pts3D_1 = D * Kinv[None, None, None, ...] @ coord[:, :, :, :, None]
+    # pts3D_1 live in your coordinate system. Convert them to Py3D's
+    pts3D_1 = M[None, None, None, ...] @ pts3D_1
+    # from reference to targe tviewpoint
+    pts3D_2 = R[None, None, None, ...] @ pts3D_1 + t[None, None, None, :, None]
+    return pts3D_2[:, :, :, :3, 0][0]
+
+
+def create_triangles(h, w, mask=None):
+    """
+    Reference: https://github.com/google-research/google-research/blob/e96197de06613f1b027d20328e06d69829fa5a89/infinite_nature/render_utils.py#L68
+    Creates mesh triangle indices from a given pixel grid size.
+        This function is not and need not be differentiable as triangle indices are
+        fixed.
+    Args:
+    h: (int) denoting the height of the image.
+    w: (int) denoting the width of the image.
+    Returns:
+    triangles: 2D numpy array of indices (int) with shape (2(W-1)(H-1) x 3)
+    """
+    x, y = np.meshgrid(range(w - 1), range(h - 1))
+    tl = y * w + x
+    tr = y * w + x + 1
+    bl = (y + 1) * w + x
+    br = (y + 1) * w + x + 1
+    triangles = np.array([tl, bl, tr, br, tr, bl])
+    triangles = np.transpose(triangles, (1, 2, 0)).reshape(
+        ((w - 1) * (h - 1) * 2, 3))
+    if mask is not None:
+        mask = mask.reshape(-1)
+        triangles = triangles[mask[triangles].all(1)]
+    return triangles
+
+
+def depth_edges_mask(depth):
+    """Returns a mask of edges in the depth map.
+    Args:
+    depth: 2D numpy array of shape (H, W) with dtype float32.
+    Returns:
+    mask: 2D numpy array of shape (H, W) with dtype bool.
+    """
+    # Compute the x and y gradients of the depth map.
+    depth_dx, depth_dy = np.gradient(depth)
+    # Compute the gradient magnitude.
+    depth_grad = np.sqrt(depth_dx ** 2 + depth_dy ** 2)
+    # Compute the edge mask.
+    mask = depth_grad > 0.05
+    return mask
+
+
+def mesh_reconstruction(
+        masked_image: np.ndarray,
+        mask: np.ndarray,
+        remove_edges: bool = True,
+        fov: float = 55.,
+        mask_threshold: float = 25.,
+):
+    rgb = masked_image[..., :3].transpose(2, 0, 1) / 255.
+    sample = torch.from_numpy(rgb).to(device).unsqueeze(0).float()
+    with torch.no_grad():
+        depth = model.infer(sample)
+        depth = depth.squeeze().cpu().numpy()
+
+    pts3d = depth_to_points(depth[None], fov=fov)
+    pts3d = pts3d.reshape(-1, 3)
+    pts3d = pts3d.reshape(-1, 3)
+    verts = pts3d.reshape(-1, 3)
+    rgb = rgb.transpose(1, 2, 0)
+    mask = mask[..., 0] > mask_threshold
+    edge_mask = depth_edges_mask(depth)
+    if remove_edges:
+        mask = np.logical_and(mask, ~edge_mask)
+    triangles = create_triangles(rgb.shape[0], rgb.shape[1], mask=mask)
+    colors = rgb.reshape(-1, 3)
+    mesh = trimesh.Trimesh(vertices=verts, faces=triangles, vertex_colors=colors)
+
+    # Save as glb tmp file (obj will look inverted in ui)
+    mesh_file = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
+    mesh_file_path = mesh_file.name
+    mesh.export(mesh_file_path)
+    return mesh_file_path

demo/relighting_gen.py

@@ -0,0 +1,53 @@
+import imageio
+import numpy as np
+import torch
+from diffusers import UniPCMultistepScheduler, StableDiffusionControlNetPipeline
+from diffusers.utils import get_class_from_dynamic_module
+
+from tqdm import tqdm
+
+device = torch.device('cpu')
+dtype = torch.float32
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+    dtype = torch.float16
+
+NeuralTextureControlNetModel = get_class_from_dynamic_module(
+    "dilightnet/model_helpers",
+    "neuraltexture_controlnet.py",
+    "NeuralTextureControlNetModel"
+)
+controlnet = NeuralTextureControlNetModel.from_pretrained(
+    "dilightnet/DiLightNet",
+    torch_dtype=dtype,
+)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1", controlnet=controlnet, torch_dtype=dtype
+).to(device)
+pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.set_progress_bar_config(disable=True)
+
+
+def relighting_gen(masked_ref_img, mask, cond_path, frames, prompt, steps, seed, cfg):
+    mask = mask[..., :1] / 255.
+    for i in tqdm(range(frames)):
+        source_image = masked_ref_img[..., :3] / 255.
+        cond_diffuse = imageio.v3.imread(f'{cond_path}/hint{i:02d}_diffuse.png') / 255.
+        if cond_diffuse.shape[-1] == 4:
+            cond_diffuse = cond_diffuse[..., :3] * cond_diffuse[..., 3:]
+        cond_ggx034 = imageio.v3.imread(f'{cond_path}/hint{i:02d}_ggx0.34.png') / 255.
+        if cond_ggx034.shape[-1] == 4:
+            cond_ggx034 = cond_ggx034[..., :3] * cond_ggx034[..., 3:]
+        cond_ggx013 = imageio.v3.imread(f'{cond_path}/hint{i:02d}_ggx0.13.png') / 255.
+        if cond_ggx013.shape[-1] == 4:
+            cond_ggx013 = cond_ggx013[..., :3] * cond_ggx013[..., 3:]
+        cond_ggx005 = imageio.v3.imread(f'{cond_path}/hint{i:02d}_ggx0.05.png') / 255.
+        if cond_ggx005.shape[-1] == 4:
+            cond_ggx005 = cond_ggx005[..., :3] * cond_ggx005[..., 3:]
+        hint = np.concatenate([mask, source_image, cond_diffuse, cond_ggx005, cond_ggx013, cond_ggx034], axis=2).astype(np.float32)[None]
+        hint = torch.from_numpy(hint).to(dtype).permute(0, 3, 1, 2).to(device)
+        generator = torch.manual_seed(seed)
+        image = pipe(
+            prompt, num_inference_steps=steps, generator=generator, image=hint, num_images_per_prompt=1, guidance_scale=cfg, output_type='np',
+        ).images[0]  # [H, W, C]
+        imageio.imwrite(f'{cond_path}/relighting{i:02d}.png', (image * 255).clip(0, 255).astype(np.uint8))

demo/render_hints.py

@@ -0,0 +1,150 @@
+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']

demo/rm_bg.py

@@ -0,0 +1,23 @@
+import numpy as np
+import rembg
+
+
+def rm_bg(img, use_sam=False):
+    output = rembg.remove(img)
+    mask = np.array(output)[:, :, 3]
+
+    # use sam for mask refinement
+    if use_sam:
+        session = rembg.new_session('sam', sam_model='sam_vit_h_4b8939')
+        bool_mask = mask > 0
+        y1, y2, x1, x2 = (
+            np.nonzero(bool_mask)[0].min(),
+            np.nonzero(bool_mask)[0].max(),
+            np.nonzero(bool_mask)[1].min(),
+            np.nonzero(bool_mask)[1].max()
+        )
+        output = rembg.remove(img, session=session, sam_prompt=[
+            {'type': 'rectangle', 'label': 1, 'data': [x1, y1, x2, y2]}
+        ])
+
+    return output, mask

requirements.txt

@@ -0,0 +1,15 @@
+numpy==1.26.4
+scipy==1.13.0
+diffusers==0.27.2
+transformers==4.39.3
+accelerate==0.29.3
+timm==0.6.12  # must use this version, required by MiDaS
+rembg==2.0.56
+trimesh==4.3.1
+opencv-contrib-python==4.9.0.80
+tqdm==4.66.2
+bpy==3.6.0
+bpy-helper==0.0.0
+gradio==4.27.0
+einops==0.7.0
+imageio[ffmpeg]==2.34.0