voxel model, [duration,fps,cuda] opts for video

Author: parameterized

.gitignore

@@ -2,4 +2,5 @@
 /videos
 
 __pycache__/
-.vscode/
+.vscode/
+.ipynb_checkpoints

Main.ipynb

@@ -8,7 +8,7 @@
 
 from dataset import get_rays
 
-def gen_video(dataset, model, path):
+def gen_video(dataset, model, path, duration=2, fps=12, cuda=True):
     camxs = dataset.poses[:, 0, -1]
     camzs = dataset.poses[:, 2, -1]
     cam_left_id = np.argmin(camxs + camzs)
@@ -21,6 +21,7 @@ def gen_video(dataset, model, path):
     ease_quad = lambda x: 2 * x**2 if x < 0.5 else 1 - (-2 * x + 2)**2 / 2
 
     def make_frame(t):
+        t = t * 2 / duration
         t = ease_quad(saw(t))
 
         a = np.array
@@ -34,12 +35,14 @@ def make_frame(t):
 
         H, W, focal = dataset.H, dataset.W, dataset.focal
         with torch.no_grad():
-            rays = get_rays(H, W, focal, pose_t).cuda()
+            rays = get_rays(H, W, focal, pose_t)
+            if cuda:
+                rays = rays.cuda()
             pred = model(rays)
             pred_r = transforms.Resize(256)(pred.clip(0, 1).reshape(H, W, 3).permute(2,0,1))
             img = pred_r.permute(1,2,0).multiply(255).int().cpu().numpy()
         return img
 
     os.makedirs(os.path.dirname(path), exist_ok=True)
-    clip = mpy.VideoClip(make_frame, duration=2)
-    clip.write_videofile(path, fps=12)
+    clip = mpy.VideoClip(make_frame, duration=duration)
+    clip.write_videofile(path, fps=fps)