Update app.py

app.py

@@ -1,35 +1,77 @@
-import subprocess
-import os
-import nerf
 import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.utils.data as data
 
-# Check if the directory called Model exists
-if not os.path.exists("Model"):
-    # Create the directory
-    os.mkdir("Model")
+from nerf import NeRF
+from dataset import Dataset
 
-# Check if the model is downloaded
-if not os.path.exists("Model/nerf_model.ckpt"):
-    # Download the model using the subprocess module
-    subprocess.run(["wget", "https://github.com/bmild/nerf/releases/download/v0.5/nerf_model.ckpt"])
 
-# Load the NeRF model
-nerf_model = nerf.NeRF.load("Model/nerf_model.ckpt")
+def train(nerf, dataloader, optimizer, device):
+    nerf.train()
+    for i, data in enumerate(dataloader):
+        # Get the input and target images.
+        viewdirs, radiances = data
+        viewdirs = viewdirs.to(device)
+        radiances = radiances.to(device)
 
-# Define the camera parameters
-camera = nerf.Camera(
-    fov=60,
-    focal_length=50,
-    znear=1,
-    zfar=100,
-    principal_point=(0.5, 0.5),
-)
+        # Forward pass.
+        outputs = nerf(viewdirs)
 
-# Define the viewing direction
-viewing_direction = torch.tensor([0, 0, 1])
+        # Compute the loss.
+        loss = nn.functional.mse_loss(outputs, radiances)
 
-# Render the novel view
-novel_view = nerf_model.render(viewing_direction, camera)
+        # Backpropagate the loss.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
 
-# Save the novel view
-torch.save(novel_view, "novel_view.png")
+
+def test(nerf, dataloader, device):
+    nerf.eval()
+    psnrs = []
+    for i, data in enumerate(dataloader):
+        # Get the input and target images.
+        viewdirs, radiances = data
+        viewdirs = viewdirs.to(device)
+        radiances = radiances.to(device)
+
+        # Forward pass.
+        outputs = nerf(viewdirs)
+
+        # Compute the PSNR.
+        psnrs.append(
+            torch.mean(
+                torch.nn.functional.psnr(outputs, radiances, data["intrinsics"])
+            )
+        )
+
+    return np.mean(psnrs)
+
+
+def main():
+    # Create the dataset.
+    dataset = Dataset.from_json("data/nerf_synthetic_data.json")
+    dataloader = data.DataLoader(dataset, batch_size=1, shuffle=True)
+
+    # Create the NeRF model.
+    nerf = NeRF(32, 64, 8).to(device)
+
+    # Create the optimizer.
+    optimizer = optim.Adam(nerf.parameters(), lr=1e-3)
+
+    # Train the NeRF model.
+    for i in range(1000):
+        train(nerf, dataloader, optimizer, device)
+
+        # Print the loss and PSNR every 100 iterations.
+        if i % 100 == 0:
+            loss = test(nerf, dataloader, device)
+            print(f"Loss: {loss:.4f}")
+
+    # Save the NeRF model.
+    nerf.save("nerf.pth")
+
+
+if __name__ == "__main__":
+    main()