File size: 1,660 Bytes
d661b19 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from typing import List
from pydantic import validator
from my.config import BaseConf, SingleOrList, dispatch
from my.utils.seed import seed_everything
import numpy as np
from voxnerf.vox import VOXRF_REGISTRY
from voxnerf.pipelines import train
class VoxConfig(BaseConf):
model_type: str = "VoxRF"
bbox_len: float = 1.5
grid_size: SingleOrList(int) = [128, 128, 128]
step_ratio: float = 0.5
density_shift: float = -10.
ray_march_weight_thres: float = 0.0001
c: int = 3
blend_bg_texture: bool = False
bg_texture_hw: int = 64
@validator("grid_size")
def check_gsize(cls, grid_size):
if isinstance(grid_size, int):
return [grid_size, ] * 3
else:
assert len(grid_size) == 3
return grid_size
def make(self):
params = self.dict()
m_type = params.pop("model_type")
model_fn = VOXRF_REGISTRY.get(m_type)
radius = params.pop('bbox_len')
aabb = radius * np.array([
[-1, -1, -1],
[1, 1, 1]
])
model = model_fn(aabb=aabb, **params)
return model
class TrainerConfig(BaseConf):
model: VoxConfig = VoxConfig()
scene: str = "lego"
n_epoch: int = 2
bs: int = 4096
lr: float = 0.02
def run(self):
args = self.dict()
args.pop("model")
model = self.model.make()
train(model, **args)
if __name__ == "__main__":
seed_everything(0)
dispatch(TrainerConfig)
|