From b653eef91cd37bcc0b7d2069e5460fb37b07d1ef Mon Sep 17 00:00:00 2001
From: LSong <12553253+lsongx@users.noreply.github.com>
Date: Thu, 9 Feb 2023 16:43:21 -0800
Subject: [PATCH] add time (#1391)

* add time

* fix black

---------

Co-authored-by: LcDog <12553253+LcDog@users.noreply.github.com>
---
 nerfstudio/cameras/camera_paths.py           | 4 ++++
 nerfstudio/cameras/cameras.py                | 4 ++--
 nerfstudio/field_components/temporal_grid.py | 3 +--
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/nerfstudio/cameras/camera_paths.py b/nerfstudio/cameras/camera_paths.py
index 61c489edf..a43504694 100644
--- a/nerfstudio/cameras/camera_paths.py
+++ b/nerfstudio/cameras/camera_paths.py
@@ -100,12 +100,16 @@ def get_spiral_path(
         new_c2ws.append(c2wh[:3, :4])
     new_c2ws = torch.stack(new_c2ws, dim=0)
 
+    times = None
+    if camera.times is not None:
+        times = torch.linspace(0, 1, steps)[:, None]
     return Cameras(
         fx=camera.fx[0],
         fy=camera.fy[0],
         cx=camera.cx[0],
         cy=camera.cy[0],
         camera_to_worlds=new_c2ws,
+        times=times,
     )
 
 
diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py
index d7d228bae..80b7f6487 100644
--- a/nerfstudio/cameras/cameras.py
+++ b/nerfstudio/cameras/cameras.py
@@ -230,7 +230,7 @@ class Cameras(TensorDataclass):
             c_x_y: cx or cy for when h_w == None
         """
         if isinstance(h_w, int):
-            h_w = torch.Tensor([h_w]).to(torch.int64).to(self.device)
+            h_w = torch.as_tensor([h_w]).to(torch.int64).to(self.device)
         elif isinstance(h_w, torch.Tensor):
             assert not torch.is_floating_point(h_w), f"height and width tensor must be of type int, not: {h_w.dtype}"
             h_w = h_w.to(torch.int64).to(self.device)
@@ -238,7 +238,7 @@ class Cameras(TensorDataclass):
                 h_w = h_w.unsqueeze(-1)
         # assert torch.all(h_w == h_w.view(-1)[0]), "Batched cameras of different h, w will be allowed in the future."
         elif h_w is None:
-            h_w = torch.Tensor((c_x_y * 2).to(torch.int64).to(self.device))
+            h_w = torch.as_tensor((c_x_y * 2)).to(torch.int64).to(self.device)
         else:
             raise ValueError("Height must be an int, tensor, or None, received: " + str(type(h_w)))
         return h_w
diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py
index 6bccc110a..3bbe274b7 100644
--- a/nerfstudio/field_components/temporal_grid.py
+++ b/nerfstudio/field_components/temporal_grid.py
@@ -116,7 +116,6 @@ class TemporalGridEncodeFunc(Function):
     @staticmethod
     @custom_bwd
     def backward(ctx, grad):
-
         inputs, temporal_row_index, embeddings, offsets, dy_dx = ctx.saved_tensors
         B, D, grid_channel, C, L, S, H, gridtype = ctx.dims
         align_corners = ctx.align_corners
@@ -336,7 +335,7 @@ class TemporalGridEncoder(nn.Module):
         """
         outputs = TemporalGridEncodeFunc.apply(
             xyz,
-            self.get_temporal_index(time[:, 0]),
+            self.get_temporal_index(time[:, 0].float()),
             self.embeddings,
             self.offsets,
             self.per_level_scale,