fix: Add color render

trellis/utils/render_utils.py

@@ -83,18 +83,28 @@ def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=N
         else:
             res = renderer.render(sample, extr, intr)
             if 'normal' not in rets: rets['normal'] = []
+            if 'color' not in rets: rets['color'] = []
+            rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
             rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
     return rets
 
-
-def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, **kwargs):
-    yaws = torch.linspace(0, 2 * 3.1415, num_frames)
-    pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
+def render_video(sample, resolution=512, ssaa=4, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, 
+                 inverse_direction=False, pitch=-1, device="cuda", **kwargs):
+    if inverse_direction:
+        yaws = torch.linspace(3.1415, -3.1415, num_frames)
+    else:
+        yaws = torch.linspace(0, 2 * 3.1415, num_frames)
+    if pitch != -1:
+        pitch = pitch * torch.ones(num_frames)
+    else:
+        pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
     yaws = yaws.tolist()
     pitch = pitch.tolist()
     extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov)
-    return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
-
+    
+    res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color, 'ssaa': ssaa}, device=device, **kwargs)
+    res.update({'extrinsics': extrinsics, 'intrinsics': intrinsics})
+    return res
 
 def render_multiview(sample, resolution=512, nviews=30):
     r = 2