Update app.py

app.py

@@ -1,9 +1,11 @@
+# final one
 import torch
 import gradio as gr
 import os
 import numpy as np
 import trimesh
 import mcubes
+import imageio
 from torchvision.utils import save_image
 from PIL import Image
 from transformers import AutoModel, AutoConfig
@@ -11,13 +13,13 @@ from rembg import remove, new_session
 from functools import partial
 from kiui.op import recenter
 import kiui
-
+from gradio_litmodel3d import LitModel3D
 
 # we load the pre-trained model from HF
 class LRMGeneratorWrapper:
     def __init__(self):
-        self.config = AutoConfig.from_pretrained("jadechoghari/custom-llrm", trust_remote_code=True)
-        self.model = AutoModel.from_pretrained("jadechoghari/custom-llrm", trust_remote_code=True)
+        self.config = AutoConfig.from_pretrained("jadechoghari/vfusion3d", trust_remote_code=True)
+        self.model = AutoModel.from_pretrained("jadechoghari/vfusion3d", trust_remote_code=True)
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.model.to(self.device)
         self.model.eval()
@@ -27,7 +29,7 @@ class LRMGeneratorWrapper:
 
 model_wrapper = LRMGeneratorWrapper()
 
-
+# we preprocess the input image
 def preprocess_image(image, source_size):
     session = new_session("isnet-general-use")
     rembg_remove = partial(remove, session=session)
@@ -42,11 +44,9 @@ def preprocess_image(image, source_size):
     image = torch.clamp(image, 0, 1)
     return image
 
+# Copied from https://github.com/facebookresearch/vfusion3d/blob/main/lrm/cam_utils.py and
+# https://github.com/facebookresearch/vfusion3d/blob/main/lrm/inferrer.py
 def get_normalized_camera_intrinsics(intrinsics: torch.Tensor):
-    """
-    intrinsics: (N, 3, 2), [[fx, fy], [cx, cy], [width, height]]
-    Return batched fx, fy, cx, cy
-    """
     fx, fy = intrinsics[:, 0, 0], intrinsics[:, 0, 1]
     cx, cy = intrinsics[:, 1, 0], intrinsics[:, 1, 1]
     width, height = intrinsics[:, 2, 0], intrinsics[:, 2, 1]
@@ -54,52 +54,112 @@ def get_normalized_camera_intrinsics(intrinsics: torch.Tensor):
     cx, cy = cx / width, cy / height
     return fx, fy, cx, cy
 
-
 def build_camera_principle(RT: torch.Tensor, intrinsics: torch.Tensor):
+    fx, fy, cx, cy = get_normalized_camera_intrinsics(intrinsics)
+    return torch.cat([
+        RT.reshape(-1, 12),
+        fx.unsqueeze(-1), fy.unsqueeze(-1), cx.unsqueeze(-1), cy.unsqueeze(-1),
+    ], dim=-1)
+
+def _default_intrinsics():
+    fx = fy = 384
+    cx = cy = 256
+    w = h = 512
+    intrinsics = torch.tensor([
+        [fx, fy],
+        [cx, cy],
+        [w, h],
+    ], dtype=torch.float32)
+    return intrinsics
+
+def _default_source_camera(batch_size: int = 1):
+    canonical_camera_extrinsics = torch.tensor([[
+        [0, 0, 1, 1],
+        [1, 0, 0, 0],
+        [0, 1, 0, 0],
+    ]], dtype=torch.float32)
+    canonical_camera_intrinsics = _default_intrinsics().unsqueeze(0)
+    source_camera = build_camera_principle(canonical_camera_extrinsics, canonical_camera_intrinsics)
+    return source_camera.repeat(batch_size, 1)
+
+def _center_looking_at_camera_pose(camera_position: torch.Tensor, look_at: torch.Tensor = None, up_world: torch.Tensor = None):
+    """
+    camera_position: (M, 3)
+    look_at: (3)
+    up_world: (3)
+    return: (M, 3, 4)
+    """
+    # by default, looking at the origin and world up is pos-z
+    if look_at is None:
+        look_at = torch.tensor([0, 0, 0], dtype=torch.float32)
+    if up_world is None:
+        up_world = torch.tensor([0, 0, 1], dtype=torch.float32)
+    look_at = look_at.unsqueeze(0).repeat(camera_position.shape[0], 1)
+    up_world = up_world.unsqueeze(0).repeat(camera_position.shape[0], 1)
+
+    z_axis = camera_position - look_at
+    z_axis = z_axis / z_axis.norm(dim=-1, keepdim=True)
+    x_axis = torch.cross(up_world, z_axis)
+    x_axis = x_axis / x_axis.norm(dim=-1, keepdim=True)
+    y_axis = torch.cross(z_axis, x_axis)
+    y_axis = y_axis / y_axis.norm(dim=-1, keepdim=True)
+    extrinsics = torch.stack([x_axis, y_axis, z_axis, camera_position], dim=-1)
+    return extrinsics
+
+def compose_extrinsic_RT(RT: torch.Tensor):
+    """
+    Compose the standard form extrinsic matrix from RT.
+    Batched I/O.
+    """
+    return torch.cat([
+        RT,
+        torch.tensor([[[0, 0, 0, 1]]], dtype=torch.float32).repeat(RT.shape[0], 1, 1).to(RT.device)
+        ], dim=1)
+
+def _build_camera_standard(RT: torch.Tensor, intrinsics: torch.Tensor):
     """
     RT: (N, 3, 4)
     intrinsics: (N, 3, 2), [[fx, fy], [cx, cy], [width, height]]
     """
+    E = compose_extrinsic_RT(RT)
     fx, fy, cx, cy = get_normalized_camera_intrinsics(intrinsics)
+    I = torch.stack([
+        torch.stack([fx, torch.zeros_like(fx), cx], dim=-1),
+        torch.stack([torch.zeros_like(fy), fy, cy], dim=-1),
+        torch.tensor([[0, 0, 1]], dtype=torch.float32, device=RT.device).repeat(RT.shape[0], 1),
+    ], dim=1)
     return torch.cat([
-        RT.reshape(-1, 12),
-        fx.unsqueeze(-1), fy.unsqueeze(-1), cx.unsqueeze(-1), cy.unsqueeze(-1),
+        E.reshape(-1, 16),
+        I.reshape(-1, 9),
     ], dim=-1)
 
+def _default_render_cameras(batch_size: int = 1):
+    M = 160
+    radius = 1.5
+    elevation = 0
+    camera_positions = []
+    rand_theta = np.random.uniform(0, np.pi/180)
+    elevation = np.radians(elevation)
+    for i in range(M):
+        theta = 2 * np.pi * i / M + rand_theta
+        x = radius * np.cos(theta) * np.cos(elevation)
+        y = radius * np.sin(theta) * np.cos(elevation)
+        z = radius * np.sin(elevation)
+        camera_positions.append([x, y, z])
+    camera_positions = torch.tensor(camera_positions, dtype=torch.float32)
+    extrinsics = _center_looking_at_camera_pose(camera_positions)
 
-def _default_intrinsics():
-        fx = fy = 384
-        cx = cy = 256
-        w = h = 512
-        intrinsics = torch.tensor([
-            [fx, fy],
-            [cx, cy],
-            [w, h],
-        ], dtype=torch.float32)
-        return intrinsics
+    render_camera_intrinsics = _default_intrinsics().unsqueeze(0).repeat(extrinsics.shape[0], 1, 1)
+    render_cameras = _build_camera_standard(extrinsics, render_camera_intrinsics)
+    return render_cameras.unsqueeze(0).repeat(batch_size, 1, 1)
 
-def _default_source_camera(batch_size: int = 1):
-        dist_to_center = 1.5
-        canonical_camera_extrinsics = torch.tensor([[
-            [0, 0, 1, 1],
-            [1, 0, 0, 0],
-            [0, 1, 0, 0],
-        ]], dtype=torch.float32)
-        canonical_camera_intrinsics = _default_intrinsics().unsqueeze(0)
-        source_camera = build_camera_principle(canonical_camera_extrinsics, canonical_camera_intrinsics)
-        return source_camera.repeat(batch_size, 1)
-
-
-#Ref: https://github.com/jadechoghari/vfusion3d/blob/main/lrm/inferrer.py
-def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export_mesh=True):
+def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export_mesh=False, export_video=True, fps=30):
     image = preprocess_image(image, source_size).to(model_wrapper.device)
     source_camera = _default_source_camera(batch_size=1).to(model_wrapper.device)
-    # TODO: export video we need render_camera
-    # render_camera = _default_render_cameras(batch_size=1).to(model_wrapper.device)
 
     with torch.no_grad():
         planes = model_wrapper.forward(image, source_camera)
-        
+
         if export_mesh:
             grid_out = model_wrapper.model.synthesizer.forward_grid(planes=planes, grid_size=mesh_size)
             vtx, faces = mcubes.marching_cubes(grid_out['sigma'].float().squeeze(0).squeeze(-1).cpu().numpy(), 1.0)
@@ -108,34 +168,92 @@ def generate_mesh(image, source_size=512, render_size=384, mesh_size=512, export
             vtx_colors = model_wrapper.model.synthesizer.forward_points(planes, vtx_tensor)['rgb'].float().squeeze(0).cpu().numpy()
             vtx_colors = (vtx_colors * 255).astype(np.uint8)
             mesh = trimesh.Trimesh(vertices=vtx, faces=faces, vertex_colors=vtx_colors)
-            
+
             mesh_path = "awesome_mesh.obj"
             mesh.export(mesh_path, 'obj')
-            return mesh_path
 
-# we will convert image to mesh
+            return mesh_path, mesh_path
+
+        if export_video:
+            render_cameras = _default_render_cameras(batch_size=1).to(model_wrapper.device)
+            frames = []
+            chunk_size = 2
+            for i in range(0, render_cameras.shape[1], chunk_size):
+                frame_chunk = model_wrapper.model.synthesizer(
+                    planes,
+                    render_cameras[:, i:i + chunk_size],
+                    render_size,
+                    render_size,
+                    0,
+                    0
+                )
+                frames.append(frame_chunk['images_rgb'])
+
+            frames = torch.cat(frames, dim=1)
+            frames = (frames.permute(0, 2, 3, 1).cpu().numpy() * 255).astype(np.uint8)
+
+            video_path = "awesome_video.mp4"
+            imageio.mimwrite(video_path, frames, fps=fps)
+
+            return None, video_path
+
+    return None, None
+
 def step_1_generate_obj(image):
-    mesh_path = generate_mesh(image)
-    return mesh_path
+    mesh_path, _ = generate_mesh(image, export_mesh=True)
+    return mesh_path, mesh_path
+
+def step_2_generate_video(image):
+    _, video_path = generate_mesh(image, export_video=True)
+    return video_path
 
-# we will convert mesh to 3d-image
-def step_2_display_3d_model(mesh_file):
+def step_3_display_3d_model(mesh_file):
     return mesh_file
 
+# set up the example files from assets folder, we limit to 10
+example_folder = "assets"
+examples = [os.path.join(example_folder, f) for f in os.listdir(example_folder) if f.endswith(('.png', '.jpg', '.jpeg'))][:10]
+
 with gr.Blocks() as demo:
     with gr.Row():
+
         with gr.Column():
             img_input = gr.Image(type="pil", label="Input Image")
-            generate_button = gr.Button("Generate and Visualize 3D Model")
+            examples_component = gr.Examples(examples=examples, inputs=img_input, outputs=None, examples_per_page=3)
+            generate_mesh_button = gr.Button("Generate and Download Mesh")
+            generate_video_button = gr.Button("Generate and Download Video")
             obj_file_output = gr.File(label="Download .obj File")
-            
+            video_file_output = gr.File(label="Download Video")
+
         with gr.Column():
-            model_output = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model Visualization")
+            model_output = LitModel3D(
+                clear_color=[0.1, 0.1, 0.1, 0],  # can adjust background color for better contrast
+                label="3D Model Visualization",
+                scale=1.0,
+                tonemapping="aces",  # can use aces tonemapping for more realistic lighting
+                exposure=1.0,        # can adjust exposure to control brightness
+                contrast=1.1,        # can slightly increase contrast for better depth
+                camera_position=(0, 0, 2),  # will set initial camera position to center the model
+                zoom_speed=0.5,      # will adjust zoom speed for better control
+                pan_speed=0.5,       # will adjust pan speed for better control
+                interactive=True     # this allow users to interact with the model
+            )
+            
+        
+    # clear outputs
+    def clear_model_viewer():
+        """Reset the Model3D component before loading a new model."""
+        return gr.update(value=None)
     
     def generate_and_visualize(image):
         mesh_path = step_1_generate_obj(image)
         return mesh_path, mesh_path
-    
-    generate_button.click(generate_and_visualize, inputs=img_input, outputs=[obj_file_output, model_output])
 
-demo.launch()
+    # first we clear the existing 3D model
+    img_input.change(clear_model_viewer, inputs=None, outputs=model_output)
+
+    # then, generate the mesh and video
+    generate_mesh_button.click(step_1_generate_obj, inputs=img_input, outputs=[obj_file_output, model_output])
+    generate_video_button.click(step_2_generate_video, inputs=img_input, outputs=video_file_output)
+
+demo.launch(debug=True)