Add Shapenet dataset

[pira998]

import torch
import glob
import json
import numpy as np
import os
from tqdm import tqdm
from pathlib import Path

from .ray_utils import get_ray_directions
from .color_utils import read_image

from .base import BaseDataset
import imageio


class ShapeNetDataset(BaseDataset):
    def __init__(self, root_dir, split='train', downsample=1.0, **kwargs):
        super().__init__(root_dir, split, downsample)

        self.read_intrinsics()

        if kwargs.get('read_meta', True):
            self.read_meta(split)
    
    def read_intrinsics(self):
        meta_path = os.path.join(self.root_dir, 'transforms.json')
        with open(meta_path, 'r') as f:
            meta_data = json.load(f)
        frame = meta_data["frames"][0]
        img_name = f"{Path(frame['file_path']).stem}.png"
        img_path = self.root_dir + '/' +img_name
        img = imageio.imread(img_path)
        H, W = img.shape[:2]
        camera_angle_x = meta_data["camera_angle_x"]
        camera_angle_x = torch.FloatTensor([camera_angle_x])
        focal = 0.5 * W / torch.tan(0.5 * camera_angle_x)
        K = [[ focal, 0, self.downsample*W],
                [ 0, focal, self.downsample*H], 
                [ 0, 0, 1]]
        K = torch.FloatTensor(K)
        self.meta_data = meta_data
        self.K = K
        self.directions = get_ray_directions(H, W, self.K)

        self.img_wh = (W, H)
    
    def read_meta(self, split):
        self.rays = []
        self.poses = []
        if split == 'test' or split =='val':
            for frame_idx in tqdm(range(30, 50)):
                frame = self.meta_data["frames"][frame_idx]
                pose = np.asarray(frame['transform_matrix'], dtype=np.float32)
                self.poses += [pose.reshape(4,4)[:3]]
        else:
            for frame_idx in tqdm(range(0, 30)):
                frame = self.meta_data["frames"][frame_idx]
                img_name = f"{Path(frame['file_path']).stem}.png"
                img_path = self.root_dir + '/'+img_name
                img = read_image(img_path, self.img_wh)
                pose = np.asarray(frame['transform_matrix'], dtype=np.float32)
                self.poses += [pose.reshape(4,4)[:3]]
                self.rays += [img]

            
            self.rays = torch.FloatTensor(np.stack(self.rays))
        self.poses = torch.FloatTensor(self.poses)

I try to add shapenet dataset. I got Floating point exception when computing loss.backward()

/home/piraveen/minicondadai_py38/envs/ngp_pl/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py(143)closure() -> self._backward_fn(step_output.closure_loss) (Pdb) n benchmarking/d.sh: line 7: 132109 Floating point exception(core dumped) python train.py --root_dir $ROOT_DIR --exp_name lamp --num_epochs 20 --batch_size 8192 --lr 2e-2 --eval_lpips (ngp_pl)

[pira998]

1a5ebc8575a4e5edcc901650bbbbb0b5.zip This is the dataset it contains 50 images

[maturk]

Hi @pira998,

It might be that your camera poses are in the wrong format, and you are not marching rays into your scene, and thus your loss does not get computed. For training with synthetic data, be careful and check your bbox (scene size) and camera poses. Make sure your camera poses are in the ngp format and that they are looking into the bbox. Also, for synthetic data, I noticed a large improvement when setting exp_step_factor=0 for small scenes.

Hope this helps.

[joejep]

Hi @pira998,

It might be that your camera poses are in the wrong format, and you are not marching rays into your scene, and thus your loss does not get computed. For training with synthetic data, be careful and check your bbox (scene size) and camera poses. Make sure your camera poses are in the ngp format and that they are looking into the bbox. Also, for synthetic data, I noticed a large improvement when setting exp_step_factor=0 for small scenes.

Hope this helps.

Please i am having a bad reconstruction when i use custom dataset.