Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
process_nerf.py 1.64 KB
Newer Older
Nianchen Deng's avatar
tog'21 baseline    
Nianchen Deng committed
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 
import sys
import os
import argparse
import json
import numpy as np

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

from utils import img

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

data_desc_path = args.dataset
data_desc_name = os.path.splitext(os.path.basename(data_desc_path))[0]
data_dir = os.path.dirname(data_desc_path) + '/'

with open(data_desc_path, 'r') as fp:
    dataset_desc = json.load(fp)

centers = np.array(dataset_desc['view_centers'])
t_max = np.max(centers, axis=0)
t_min = np.min(centers, axis=0)
dataset_desc['range'] = {
    'min': [t_min[0], t_min[1], t_min[2], 0, 0],
    'max': [t_max[0], t_max[1], t_max[2], 0, 0]
}
dataset_desc['gl_coord'] = True

for subdir in ['images', 'images_4', 'images_8']:
    dir = os.path.join(data_dir, subdir)
    files = os.listdir(dir)
    files.sort()
    image = img.load(os.path.join(dir, files[0]))
    res = image.shape[-2:]
    for i, file in enumerate(files):
        src = os.path.join(dir, file)
        tgt = os.path.join(dir, "view_%04d.jpg" % i)
        print(f"Rename {src} to {tgt}")
        os.rename(src, tgt)
    out_desc = dataset_desc.copy()
    out_desc['view_file_pattern'] = f"{subdir}/view_%04d.jpg"
    k = res[0] / dataset_desc['view_res']['y']
    out_desc['view_res'] = {
        'x': res[1],
        'y': res[0]
    }
    out_desc['cam_params'] = {
        'fx': dataset_desc['cam_params']['fx'] * k,
        'fy': dataset_desc['cam_params']['fy'] * k,
        'cx': res[1] / 2,
        'cy': res[0] / 2,
    }
    
    with open(os.path.join(data_dir, f"{subdir}.json"), 'w') as fp:
        json.dump(out_desc, fp, indent=4)