new

video_render.py

@@ -22,6 +22,15 @@ from pytorch3d.transforms import RotateAxisAngle
 
 from shader import MultiOutputShader
 
+def _rgb_to_srgb(f: torch.Tensor) -> torch.Tensor:
+    return torch.where(f <= 0.0031308, f * 12.92, torch.pow(torch.clamp(f, 0.0031308), 1.0/2.4)*1.055 - 0.055)
+
+def rgb_to_srgb(f: torch.Tensor) -> torch.Tensor:
+    assert f.shape[-1] == 3 or f.shape[-1] == 4
+    out = torch.cat((_rgb_to_srgb(f[..., 0:3]), f[..., 3:4]), dim=-1) if f.shape[-1] == 4 else _rgb_to_srgb(f)
+    assert out.shape[0] == f.shape[0] and out.shape[1] == f.shape[1]
+    return out
+
 def render_video_from_obj(input_obj_path, output_video_path, num_frames=60, image_size=512, fps=15, device="cuda"):
     if not os.path.exists(input_obj_path):
         raise FileNotFoundError(f"Input OBJ file not found: {input_obj_path}")
@@ -87,12 +96,14 @@ def render_video_from_obj(input_obj_path, output_video_path, num_frames=60, imag
                 cameras=cameras,
                 lights=lights,
                 materials=materials,
-                choices=["rgb", "mask", "normal"]
+                choices=["rgb", "mask", "normal", "albedo"]
             )
 
         renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
         render_result = renderer(mesh, cameras=cameras)
-        rgb_image = render_result["rgb"] * render_result["mask"] + (1 - render_result["mask"]) * torch.ones_like(render_result["rgb"]) * 255.0
+
+        render_result["albedo"] = rgb_to_srgb(render_result["albedo"]/255.0)*255.0
+        rgb_image = render_result["albedo"] * render_result["mask"] + (1 - render_result["mask"]) * torch.ones_like(render_result["albedo"]) * 255.0
         normal_map = render_result["normal"]
 
         rgb = rgb_image[0, ..., :3].cpu().numpy()
@@ -111,6 +122,6 @@ def render_video_from_obj(input_obj_path, output_video_path, num_frames=60, imag
     print(f"Video saved to {output_video_path}")
 
 if __name__ == '__main__':
-    input_obj_path = "/hpc2hdd/home/jlin695/code/github/Kiss3DGen/outputs/a_owl_wearing_a_hat/ISOMER/rgb_projected.obj"
+    input_obj_path = "./354e2aee-091d-4dc6-bdb1-e09be5791218_isomer_recon_mesh.obj"
     output_video_path = "output.mp4"
     render_video_from_obj(input_obj_path, output_video_path)