split.py

import json
import sys
import os
import argparse
from pathlib import Path

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

from data import get_dataset_desc_path

parser = argparse.ArgumentParser()
parser.add_argument('-o', '--output', type=str, nargs="+", required=True)
parser.add_argument("-v", "--views", type=int, nargs="+", required=True)
parser.add_argument('dataset', type=str)
args = parser.parse_args()

input = get_dataset_desc_path(args.dataset)
outputs = [
    get_dataset_desc_path(input.with_name(f"{input.stem}_{appendix}"))
    for appendix in args.output
]

with open(input, 'r') as fp:
    input_desc: dict = json.load(fp)
n_views = len(input_desc['view_centers'])

assert(len(args.views) == len(outputs))
sum_views = sum(args.views)
for i in range(len(args.views)):
    if args.views[i] == -1:
        args.views[i] = n_views - sum_views - 1
        sum_views = n_views
        break
assert(sum_views <= n_views)
for i in range(len(args.views)):
    assert(args.views[i] > 0)

offset = 0
for i in range(len(outputs)):
    n = args.views[i]
    end = offset + n
    output_desc = input_desc.copy()
    output_desc['samples'] = args.views[i]
    if 'views' in output_desc:
        output_desc['views'] = output_desc['views'][offset:end]
    else:
        output_desc['views'] = list(range(offset, end))
    output_desc['view_centers'] = output_desc['view_centers'][offset:end]
    if 'view_rots' in output_desc:
        output_desc['view_rots'] = output_desc['view_rots'][offset:end]
    with open(outputs[i], 'w') as fp:
        json.dump(output_desc, fp, indent=4)

    # Create symbol links of images
    out_dir = outputs[i].with_suffix('')
    out_dir.mkdir(exist_ok=True)
    for k in range(n):
        os.symlink(Path("..") / input.stem / (output_desc['view_file_pattern'] % output_desc['views'][k]),
                   out_dir / (input_desc['view_file_pattern'] % output_desc['views'][k]))
    offset += args.views[i]