Update app.py

app.py

@@ -1,77 +1,20 @@
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torch.utils.data as data
-
-from nerf import NeRF
-from dataset import Dataset
-
-
-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)
-
-        # Forward pass.
-        outputs = nerf(viewdirs)
-
-        # Compute the loss.
-        loss = nn.functional.mse_loss(outputs, radiances)
-
-        # Backpropagate the loss.
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-
-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()
+import os, sys
+# os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
+os.environ['CUDA_VISIBLE_DEVICES'] = '1'
+import tensorflow as tf
+tf.compat.v1.enable_eager_execution()
+
+import numpy as np
+import imageio
+import json
+import random
+import time
+import pprint
+
+import matplotlib.pyplot as plt
+
+import run_nerf
+
+from load_llff import load_llff_data
+from load_deepvoxels import load_dv_data
+from load_blender import load_blender_data