convert_old_snerffast_checkpoint.py

import sys
import os
import configargparse
import argparse
import torch
import re
from pathlib import Path
from typing import Any

sys.path.append(os.path.abspath(sys.path[0] + '/../'))


parser = argparse.ArgumentParser()
parser.add_argument('input', type=str)
args = parser.parse_args()


class SphericalViewSynConfig(object):

    def __init__(self):
        self.name = 'default'
        self.COLOR = "rgb"

        # Net parameters
        self.NET_TYPE = 'msl'
        self.N_ENCODE_DIM = 10
        self.N_DIR_ENCODE = None
        self.NORMALIZE = False
        self.DEPTH_REF = False
        self.FC_PARAMS = {
            'nf': 256,
            'n_layers': 8,
            'skips': [],
            'activation': 'relu'
        }
        self.SAMPLE_PARAMS = {
            'spherical': True,
            'depth_range': (1, 50),
            'n_samples': 32,
            'perturb_sample': True,
            'lindisp': True,
            'inverse_r': True,
        }
        self.NERF_FINE_NET_PARAMS = {
            'enable': False,
            'nf': 256,
            'n_layers': 8,
            'additional_samples': 64
        }

    def from_id(self, id: str):
        id_splited = id.split('@')
        if len(id_splited) == 2:
            self.name = id_splited[0]
        segs = id_splited[-1].split('_')
        for i, seg in enumerate(segs):
            if seg.startswith('ffc'):  # Full-connected network parameters
                self.NERF_FINE_NET_PARAMS['nf'], self.NERF_FINE_NET_PARAMS['n_layers'] = (
                    int(str) for str in seg[3:].split('x'))
                self.NERF_FINE_NET_PARAMS['enable'] = True
                continue
            if seg.startswith('fs'):  # Number of samples
                try:
                    self.NERF_FINE_NET_PARAMS['additional_samples'] = int(seg[2:])
                    self.NERF_FINE_NET_PARAMS['enable'] = True
                    continue
                except ValueError:
                    pass
            if seg.startswith('fc'):  # Full-connected network parameters
                self.FC_PARAMS['nf'], self.FC_PARAMS['n_layers'] = (
                    int(str) for str in seg[2:].split('x'))
                continue
            if seg.startswith('skip'):  # Skip connection
                self.FC_PARAMS['skips'] = [int(str)
                                           for str in seg[4:].split(',')]
                continue
            if seg.startswith('*'):  # Activation
                self.FC_PARAMS['activation'] = seg[1:]
                continue
            if seg.startswith('ed'):  # Encode direction
                self.N_DIR_ENCODE = int(seg[2:])
                if self.N_DIR_ENCODE == 0:
                    self.N_DIR_ENCODE = None
                continue
            if seg.startswith('e'):  # Encode
                self.N_ENCODE_DIM = int(seg[1:])
                continue
            if seg.startswith('d'):  # Depth range
                try:
                    self.SAMPLE_PARAMS['depth_range'] = tuple(
                        float(str) for str in seg[1:].split('-'))
                    continue
                except ValueError:
                    pass
            if seg.startswith('s'):  # Number of samples
                try:
                    self.SAMPLE_PARAMS['n_samples'] = int(seg[1:])
                    continue
                except ValueError:
                    pass
            if seg.startswith('~'):  # Negative flags
                if seg.find('p') >= 0:
                    self.SAMPLE_PARAMS['perturb_sample'] = False
                if seg.find('l') >= 0:
                    self.SAMPLE_PARAMS['lindisp'] = False
                if seg.find('i') >= 0:
                    self.SAMPLE_PARAMS['inverse_r'] = False
                if seg.find('n') >= 0:
                    self.NORMALIZE = False
                if seg.find('d') >= 0:
                    self.DEPTH_REF = False
                continue
            if seg.startswith('+'):  # Positive flags
                if seg.find('p') >= 0:
                    self.SAMPLE_PARAMS['perturb_sample'] = True
                if seg.find('l') >= 0:
                    self.SAMPLE_PARAMS['lindisp'] = True
                if seg.find('i') >= 0:
                    self.SAMPLE_PARAMS['inverse_r'] = True
                if seg.find('n') >= 0:
                    self.NORMALIZE = True
                if seg.find('d') >= 0:
                    self.DEPTH_REF = True
                continue
            if i == 0:  # NetType
                self.NET_TYPE, color_str = seg.split('-')
                self.COLOR = color_str

    def print(self):
        print('==== Config %s ====' % self.name)
        print('Net type: ', self.NET_TYPE)
        print('Encode dim: ', self.N_ENCODE_DIM)
        print('Normalize: ', self.NORMALIZE)
        print('Train with depth: ', self.DEPTH_REF)
        print('Support direction: ', False if self.N_DIR_ENCODE is None
              else f'encode to {self.N_DIR_ENCODE}')
        print('Full-connected network parameters:', self.FC_PARAMS)
        print('Sample parameters', self.SAMPLE_PARAMS)
        if self.NERF_FINE_NET_PARAMS['enable']:
            print('NeRF fine network parameters', self.NERF_FINE_NET_PARAMS)
        print('==========================')


def convert_snerffast(input_network_state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    output_model_state: dict[str, torch.Tensor] = {}
    for key, value in input_network_state.items():
        key = re.sub("mlp_(\d+)\.core\.layer(\d+)\.net\.0\.(\w+)",
                     "core.subnets.\\1.field.net.layers.\\2.net.0.\\3", key)
        key = re.sub("mlp_(\d+)\.density_out\.net\.(\w+)",
                     "core.subnets.\\1.density_decoder.net.net.0.\\2", key)
        key = re.sub("mlp_(\d+)\.color_out\.net\.(\w+)",
                     "core.subnets.\\1.color_decoder.net.net.0.\\2", key)
        output_model_state[key] = value
    return output_model_state


ckpt_path = Path(args.input)
expdir = ckpt_path.parent
config_id = expdir.name
ckpt_epochs = int(re.match("model-epoch_(\d+)", ckpt_path.stem).group(1))

config = SphericalViewSynConfig()
config.from_id(config_id)
input_checkpoint: dict[str, Any] = torch.load(ckpt_path)

output_model_state = convert_snerffast(input_checkpoint["model"])
config.print()
print(f"Epochs: {ckpt_epochs}")

output_args = {
    "model": "FsNeRF",
    "model_args": {
        "color": "rgb",
        "n_samples": config.SAMPLE_PARAMS["n_samples"],
        "n_fields": int(config.NET_TYPE[9:]),
        "depth": config.FC_PARAMS["n_layers"],
        "width": config.FC_PARAMS["nf"],
        "skips": config.FC_PARAMS["skips"],
        "act": config.FC_PARAMS["activation"],
        "ln": False,
        "xfreqs": config.N_ENCODE_DIM,
        "raw_noise_std": 0.,
        "near": config.SAMPLE_PARAMS["depth_range"][0],
        "far": config.SAMPLE_PARAMS["depth_range"][1],
        "white_bg": False,
        "coord": "dx"
    },
    "trainer": None
}

output_path = expdir / f"checkpoint_{ckpt_epochs}.tar"
torch.save({
    "args": output_args,
    "states": {
        "model": output_model_state
    }
}, output_path)
print(f"Checkpoint is saved to {output_path}")