sync

data/blender_gen/gen_fovea.py

+import bpy
+import math
+import json
+import os
+import math
+import numpy as np
+from itertools import product
+
+
+scene = bpy.context.scene
+cam_obj = scene.camera
+cam = cam_obj.data
+scene.cycles.device = 'GPU'
+
+dataset_name = 'train'
+tbox = [0.6, 0.6, 0.6]
+rbox = [320, 40]
+
+dataset_desc = {
+    'view_file_pattern': '%s/view_%%04d.png' % dataset_name,
+    "gl_coord": True,
+    'view_res': {
+        'x': 512,
+        'y': 512
+    },
+    'cam_params': {
+        'fov': 40.0, 
+        'cx': 0.5,
+        'cy': 0.5,
+        'normalized': True
+    },
+    'range': {
+        'min': [-tbox[0] / 2, -tbox[1] / 2, -tbox[2] / 2, -rbox[0] / 2, -rbox[1] / 2],
+        'max': [tbox[0] / 2, tbox[1] / 2, tbox[2] / 2, rbox[0] / 2, rbox[1] / 2]
+    },
+    'samples': [5, 5, 5, 9, 2],
+    #'samples': [2000],
+    'view_centers': [],
+    'view_rots': []
+}
+data_desc_file = f'output/{dataset_name}.json'
+
+if not os.path.exists('output'):
+    os.mkdir('output')
+
+if os.path.exists(data_desc_file):
+    with open(data_desc_file, 'r') as fp:
+        dataset_desc.update(json.load(fp))
+with open(data_desc_file, 'w') as fp:
+    json.dump(dataset_desc, fp, indent=4)
+
+# Output resolution
+scene.render.resolution_x = dataset_desc['view_res']['x']
+scene.render.resolution_y = dataset_desc['view_res']['y']
+
+# Field of view
+cam.lens_unit = 'FOV'
+cam.angle = math.radians(dataset_desc['cam_params']['fov'])
+cam.dof.use_dof = False
+
+
+def add_sample(i, x, y, z, rx, ry, render_only=False):
+    cam_obj.location = [x, y, z]
+    cam_obj.rotation_euler = [math.radians(ry), math.radians(rx), 0]
+    scene.render.filepath = 'output/' + dataset_desc['view_file_pattern'] % i
+    bpy.ops.render.render(write_still=True)
+    if not render_only:
+        dataset_desc['view_centers'].append(list(cam_obj.location))
+        dataset_desc['view_rots'].append([rx, ry])
+        with open(data_desc_file, 'w') as fp:
+            json.dump(dataset_desc, fp, indent=4)
+
+for i in range(len(dataset_desc['view_centers'])):
+    if not os.path.exists('output/' + dataset_desc['view_file_pattern'] % i):
+        add_sample(i, *dataset_desc['view_centers'][i], *dataset_desc['view_rots'][i], render_only=True)
+
+start_view = len(dataset_desc['view_centers'])
+if len(dataset_desc['samples']) == 1:
+    range_min = np.array(dataset_desc['range']['min'])
+    range_max = np.array(dataset_desc['range']['max'])
+    samples = (range_max - range_min) * np.random.rand(dataset_desc['samples'][0], 5) + range_min
+    for i in range(start_view, dataset_desc['samples'][0]):
+        add_sample(i, *list(samples[i]))
+else:
+    ranges = [
+        np.linspace(dataset_desc['range']['min'][i],
+                    dataset_desc['range']['max'][i],
+                    dataset_desc['samples'][i])
+        for i in range(0, 3)
+    ] + [
+        np.linspace(dataset_desc['range']['min'][i],
+                    dataset_desc['range']['max'][i],
+                    dataset_desc['samples'][i])
+        for i in range(3, 5)
+    ]
+
+    i = 0
+    for x, y, z, rx, ry in product(*ranges):
+        if i >= start_view:
+            add_sample(i, x, y, z, rx, ry)
+        i += 1

data/blender_gen/gen_periph.py

+import bpy
+import math
+import json
+import os
+import math
+import numpy as np
+from itertools import product
+
+
+scene = bpy.context.scene
+cam_obj = scene.camera
+cam = cam_obj.data
+scene.cycles.device = 'GPU'
+
+dataset_name = 'train'
+tbox = [0.7, 0.7, 0.7]
+rbox = [300, 120]
+
+dataset_desc = {
+    'view_file_pattern': '%s/view_%%04d.png' % dataset_name,
+    "gl_coord": True,
+    'view_res': {
+        'x': 512,
+        'y': 512
+    },
+    'cam_params': {
+        'fov': 60.0,
+        'cx': 0.5,
+        'cy': 0.5,
+        'normalized': True
+    },
+    'range': {
+        'min': [-tbox[0] / 2, -tbox[1] / 2, -tbox[2] / 2, -rbox[0] / 2, -rbox[1] / 2],
+        'max': [tbox[0] / 2, tbox[1] / 2, tbox[2] / 2, rbox[0] / 2, rbox[1] / 2]
+    },
+    'samples': [5, 5, 5, 6, 3],
+    #'samples': [2000],
+    'view_centers': [],
+    'view_rots': []
+}
+data_desc_file = f'output/{dataset_name}.json'
+
+if not os.path.exists('output'):
+    os.mkdir('output')
+
+if os.path.exists(data_desc_file):
+    with open(data_desc_file, 'r') as fp:
+        dataset_desc.update(json.load(fp))
+with open(data_desc_file, 'w') as fp:
+    json.dump(dataset_desc, fp, indent=4)
+
+# Output resolution
+scene.render.resolution_x = dataset_desc['view_res']['x']
+scene.render.resolution_y = dataset_desc['view_res']['y']
+
+# Field of view
+cam.lens_unit = 'FOV'
+cam.angle = math.radians(dataset_desc['cam_params']['fov'])
+cam.dof.use_dof = False
+
+
+def add_sample(i, x, y, z, rx, ry, render_only=False):
+    cam_obj.location = [x, y, z]
+    cam_obj.rotation_euler = [math.radians(ry), math.radians(rx), 0]
+    scene.render.filepath = 'output/' + dataset_desc['view_file_pattern'] % i
+    bpy.ops.render.render(write_still=True)
+    if not render_only:
+        dataset_desc['view_centers'].append(list(cam_obj.location))
+        dataset_desc['view_rots'].append([rx, ry])
+        with open(data_desc_file, 'w') as fp:
+            json.dump(dataset_desc, fp, indent=4)
+
+for i in range(len(dataset_desc['view_centers'])):
+    if not os.path.exists('output/' + dataset_desc['view_file_pattern'] % i):
+        add_sample(i, *dataset_desc['view_centers'][i], *dataset_desc['view_rots'][i], render_only=True)
+
+start_view = len(dataset_desc['view_centers'])
+if len(dataset_desc['samples']) == 1:
+    range_min = np.array(dataset_desc['range']['min'])
+    range_max = np.array(dataset_desc['range']['max'])
+    samples = (range_max - range_min) * np.random.rand(dataset_desc['samples'][0], 5) + range_min
+    for i in range(start_view, dataset_desc['samples'][0]):
+        add_sample(i, *list(samples[i]))
+else:
+    ranges = [
+        np.linspace(dataset_desc['range']['min'][i],
+                    dataset_desc['range']['max'][i],
+                    dataset_desc['samples'][i])
+        for i in range(0, 3)
+    ] + [
+        np.linspace(dataset_desc['range']['min'][i],
+                    dataset_desc['range']['max'][i],
+                    dataset_desc['samples'][i])
+        for i in range(3, 5)
+    ]
+
+    i = 0
+    for x, y, z, rx, ry in product(*ranges):
+        if i >= start_view:
+            add_sample(i, x, y, z, rx, ry)
+        i += 1

data/calc_range.py

+import json
+import sys
+import os
+import argparse
+import numpy as np
+import torch
+
+sys.path.append(os.path.abspath(sys.path[0] + '/../'))
+
+from utils import misc
+
+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]
+}
+with open(data_desc_path, 'w') as fp:
+    json.dump(dataset_desc, fp, indent=4)
\ No newline at end of file

data/gen_seq.py

+import json
+import sys
+import os
+import argparse
+import numpy as np
+
+sys.path.append(os.path.abspath(sys.path[0] + '/../'))
+
+from utils import seqs
+from utils import misc
+from utils.constants import *
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-r', '--rot-range', nargs='+', type=int)
+parser.add_argument('-t', '--trans-range', nargs='+', type=float)
+parser.add_argument('--fov', type=float)
+parser.add_argument('--res', type=str)
+parser.add_argument('--gl', action='store_true')
+parser.add_argument('-s', '--seq', type=str, required=True)
+parser.add_argument('-n', '--views', type=int, required=True)
+parser.add_argument('-o', '--out-desc', type=str)
+parser.add_argument('--ref', type=str)
+parser.add_argument('dataset', type=str)
+args = parser.parse_args()
+
+
+data_dir = args.dataset
+misc.create_dir(data_dir)
+out_desc_path = os.path.join(data_dir, (args.out_desc if args.out_desc else f"{args.seq}.json"))
+
+if args.ref:
+    with open(os.path.join(data_dir, args.ref), 'r') as fp:
+        ref_desc = json.load(fp)
+else:
+    if not args.trans_range or not args.rot_range or not args.fov or not args.res:
+        print('-r, -t, --fov, --res options are required if --ref is not specified')
+        exit(-1)
+    ref_desc = None
+
+if args.trans_range:
+    trans_range = np.array(list(args.trans_range) * 3 if len(args.trans_range) == 1
+                           else args.trans_range)
+else:
+    trans_range = np.array(ref_desc['range']['max'][0:3]) - \
+        np.array(ref_desc['range']['min'][0:3])
+if args.rot_range:
+    rot_range = np.array(list(args.rot_range) * 2 if len(args.rot_range) == 1
+                         else args.rot_range)
+else:
+    rot_range = np.array(ref_desc['range']['max'][3:5]) - \
+        np.array(ref_desc['range']['min'][3:5])
+filter_range = np.concatenate([trans_range, rot_range])
+
+if args.fov:
+    cam_params = {
+        'fov': args.fov,
+        'cx': 0.5,
+        'cy': 0.5,
+        'normalized': True
+    }
+else:
+    cam_params = ref_desc['cam_params']
+
+if args.res:
+    res = tuple(int(s) for s in args.res.split('x'))
+    res = {'x': res[0], 'y': res[1]}
+else:
+    res = ref_desc['view_res']
+
+if args.seq == 'helix':
+    centers, rots = seqs.helix(trans_range, 4, args.views)
+elif args.seq == 'scan_around':
+    centers, rots = seqs.scan_around(trans_range, 1, args.views)
+elif args.seq == 'look_around':
+    centers, rots = seqs.look_around(trans_range, args.views)
+
+rots *= 180 / PI
+gl = args.gl or ref_desc.get('gl_coord')
+if gl:
+    centers[:, 2] *= -1
+    rots[:, 0] *= -1
+
+dataset_desc = {
+    'gl_coord': gl,
+    'view_res': res,
+    'cam_params': cam_params,
+    'range': {
+        'min': (-0.5 * filter_range).tolist(),
+        'max': (0.5 * filter_range).tolist()
+    },
+    'samples': [args.views],
+    'view_centers': centers.tolist(),
+    'view_rots': rots.tolist()
+}
+
+with open(out_desc_path, 'w') as fp:
+    json.dump(dataset_desc, fp, indent=4)

data/gen_subset.py

+import json
+import sys
+import os
+import argparse
+import numpy as np
+
+sys.path.append(os.path.abspath(sys.path[0] + '/../'))
+
+from utils import misc
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-r', '--rot-range', nargs='+', type=int)
+parser.add_argument('-t', '--trans-range', nargs='+', type=float)
+parser.add_argument('-k', '--trainset-ratio', type=float, default=0.7)
+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)
+
+if args.trans_range:
+    trans_range = np.array(args.trans_range)
+else:
+    trans_range = np.array(dataset_desc['range']['max'][0:3]) - \
+        np.array(dataset_desc['range']['min'][0:3])
+if args.rot_range:
+    rot_range = np.array(args.rot_range)
+else:
+    rot_range = np.array(dataset_desc['range']['max'][3:5]) - \
+        np.array(dataset_desc['range']['min'][3:5])
+filter_range = np.concatenate([trans_range, rot_range])
+
+out_data_dir = data_dir + 'r%dx%d_t%.1fx%.1fx%.1f/' % (
+    int(rot_range[0]), int(rot_range[1]),
+    trans_range[0], trans_range[1], trans_range[2]
+)
+
+dataset_version = 0
+while True:
+    out_trainset_name = f'train_{dataset_version}'
+    out_testset_name = f'test_{dataset_version}'
+    if not os.path.exists(out_data_dir + out_trainset_name):
+        break
+    dataset_version += 1
+
+
+def in_range(val, range): return val >= -range / 2 and val <= range / 2
+
+
+views = []
+for i in range(len(dataset_desc['view_centers'])):
+    if in_range(dataset_desc['view_rots'][i][0], rot_range[0]) and \
+            in_range(dataset_desc['view_rots'][i][1], rot_range[1]) and \
+            in_range(dataset_desc['view_centers'][i][0], trans_range[0]) and \
+            in_range(dataset_desc['view_centers'][i][1], trans_range[1]) and \
+            in_range(dataset_desc['view_centers'][i][2], trans_range[2]):
+        views.append(i)
+
+if len(views) < 100:
+    print(f'Number of views in range is too small ({len(views)})')
+    exit()
+
+views = np.random.permutation(views)
+n_train_views = int(len(views) * args.trainset_ratio)
+train_views = np.sort(views[:n_train_views])
+test_views = np.sort(views[n_train_views:])
+
+print('Train set views: ', len(train_views))
+print('Test set views: ', len(test_views))
+
+def create_subset(views, out_desc_name):
+    views = views.tolist()
+    subset_desc = dataset_desc.copy()
+    subset_desc['view_file_pattern'] = \
+        f"{out_desc_name}/{dataset_desc['view_file_pattern'].split('/')[-1]}"
+    subset_desc['range'] = {
+        'min': list(-filter_range / 2),
+        'max': list(filter_range / 2)
+    }
+    subset_desc['samples'] = [int(len(views))]
+    subset_desc['views'] = views
+    subset_desc['view_centers'] = np.array(dataset_desc['view_centers'])[views].tolist()
+    subset_desc['view_rots'] = np.array(dataset_desc['view_rots'])[views].tolist()
+
+    with open(os.path.join(out_data_dir, f'{out_desc_name}.json'), 'w') as fp:
+        json.dump(subset_desc, fp, indent=4)
+    misc.create_dir(os.path.join(out_data_dir, out_desc_name))
+    for i in range(len(views)):
+        os.symlink(os.path.join('../../', dataset_desc['view_file_pattern'] % views[i]),
+                   os.path.join(out_data_dir, subset_desc['view_file_pattern'] % views[i]))
+
+
+misc.create_dir(out_data_dir)
+create_subset(train_views, out_trainset_name)
+create_subset(train_views, out_testset_name)

data/split_dataset.py

+import json
+import sys
+import os
+import argparse
+import numpy as np
+import torch
+
+sys.path.append(os.path.abspath(sys.path[0] + '/../'))
+
+from utils import misc
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-o', '--output', type=str, default='train1')
+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)
+
+indices = torch.arange(len(dataset_desc['view_centers'])).view(dataset_desc['samples'])
+
+idx = 0
+'''
+for i in range(3):
+    for j in range(2):
+        out_desc_name = f'part{idx:d}'
+        out_desc = dataset_desc.copy()
+        out_desc['view_file_pattern'] = f'{out_desc_name}/view_%04d.png'
+        n_x = out_desc['samples'][3] // 3
+        n_y = out_desc['samples'][4] // 2
+        views = indices[..., i * n_x:(i + 1) * n_x, j * n_y:(j + 1) * n_y].flatten().tolist()
+        out_desc['samples'] = [len(views)]
+        out_desc['views'] = views
+        out_desc['view_centers'] = np.array(dataset_desc['view_centers'])[views].tolist()
+        out_desc['view_rots'] = np.array(dataset_desc['view_rots'])[views].tolist()
+        with open(os.path.join(data_dir, f'{out_desc_name}.json'), 'w') as fp:
+            json.dump(out_desc, fp, indent=4)
+        misc.create_dir(os.path.join(data_dir, out_desc_name))
+        for k in range(len(views)):
+            os.symlink(os.path.join('..', dataset_desc['view_file_pattern'] % views[k]),
+                    os.path.join(data_dir, out_desc['view_file_pattern'] % views[k]))
+        idx += 1
+'''
+
+'''
+for xi in range(0, 4, 2):
+    for yi in range(0, 4, 2):
+        for zi in range(0, 4, 2):
+            out_desc_name = f'part{idx:d}'
+            out_desc = dataset_desc.copy()
+            out_desc['view_file_pattern'] = f'{out_desc_name}/view_%04d.png'
+            views = indices[xi:xi + 2, yi:yi + 2, zi:zi + 2].flatten().tolist()
+            out_desc['samples'] = [len(views)]
+            out_desc['views'] = views
+            out_desc['view_centers'] = np.array(dataset_desc['view_centers'])[views].tolist()
+            out_desc['view_rots'] = np.array(dataset_desc['view_rots'])[views].tolist()
+            with open(os.path.join(data_dir, f'{out_desc_name}.json'), 'w') as fp:
+                json.dump(out_desc, fp, indent=4)
+            misc.create_dir(os.path.join(data_dir, out_desc_name))
+            for k in range(len(views)):
+                os.symlink(os.path.join('..', dataset_desc['view_file_pattern'] % views[k]),
+                           os.path.join(data_dir, out_desc['view_file_pattern'] % views[k]))
+            idx += 1
+'''
+from itertools import product
+out_desc_name = args.output
+out_desc = dataset_desc.copy()
+out_desc['view_file_pattern'] = f"{out_desc_name}/{dataset_desc['view_file_pattern'].split('/')[-1]}"
+views = []
+for idx in product([0, 2, 4], [0, 2, 4], [0, 2, 4], list(range(9)), [1]):#, [0, 2, 3, 5], [1, 2, 3, 4]):
+    views += indices[idx].flatten().tolist()
+out_desc['samples'] = [len(views)]
+out_desc['views'] = views
+out_desc['view_centers'] = np.array(dataset_desc['view_centers'])[views].tolist()
+out_desc['view_rots'] = np.array(dataset_desc['view_rots'])[views].tolist()
+with open(os.path.join(data_dir, f'{out_desc_name}.json'), 'w') as fp:
+    json.dump(out_desc, fp, indent=4)
+misc.create_dir(os.path.join(data_dir, out_desc_name))
+for k in range(len(views)):
+    os.symlink(os.path.join('..', dataset_desc['view_file_pattern'] % views[k]),
+               os.path.join(data_dir, out_desc['view_file_pattern'] % views[k]))