For a Greppy Metaverse dataset: rearrange non-contiguous scenes in a dataset, generate camera normals, move to separate folders.

greppy_metaverse_dataset_renamer.py

import os, fnmatch, argparse
import numpy as np
import OpenEXR, Imath, json
import shutil, glob

# TODO update to handle stereo camera
#
# python3 greppy_metaverse_dataset_renamer.py --p /path/to/dataset

SUBFOLDER_MAP = {
    '-componentMasks.exr': 'components',
    '-depth.exr': 'depths',
    '-masks.json': 'mask-jsons',
    '-normals.exr': 'world-normals',
    '-cameraNormals.npy': 'camera-normals',
    '-rgb.jpg': 'rgbs',
    '-variantMasks.exr': 'variants',
}

NORMALS_X_CHANNEL_NAME = 'R'
NORMALS_Y_CHANNEL_NAME = 'G'
NORMALS_Z_CHANNEL_NAME = 'B'

# Get a list of all possible scenes
def scene_prefixes(dataset_path):
    dataset_prefixes = []
    for root, dirs, files in os.walk(dataset_path):
        # one mask json file per scene so we can get the prefixes from them
        for filename in fnmatch.filter(files, '*masks.json'):
            dataset_prefixes.append(filename[0:0-len('-masks.json')])
    dataset_prefixes.sort()
    return dataset_prefixes

def string_prefixes_to_sorted_ints(prefixes_list):
    unsorted = list(map(lambda x: int(x), prefixes_list))
    unsorted.sort()
    return unsorted

def swap_files_to_new_location(old_prefix_int, new_prefix_int, dataset_path):
    old_prefix_str = "{:09}".format(old_prefix_int)
    new_prefix_str = "{:09}".format(new_prefix_int)
    print("  Moving",old_prefix_str, "files to", new_prefix_str)
    for root, dirs, files in os.walk(dataset_path):
        for filename in fnmatch.filter(files, old_prefix_str+'-*'):
            os.rename(os.path.join(dataset_path,filename), os.path.join(dataset_path,filename.replace(old_prefix_str+"-", new_prefix_str+"-")))


def fix_non_contiguities(dataset_path):
    sorted_ints = string_prefixes_to_sorted_ints(scene_prefixes(dataset_path))
    i = 0
    while i < len(sorted_ints):
        if i == 0:
            i += 1
            continue
        if sorted_ints[i-1] + 1 == sorted_ints[i]: # contiguous
            i += 1
            continue
        # non-contiguous scenario: take last file, insert it at the correct number
        swap_files_to_new_location(sorted_ints[-1], sorted_ints[i-1] + 1, dataset_path)
        # insert the elemnent before this index, and try looking at this index again
        # don't increment i
        sorted_ints.insert(i, sorted_ints[i-1] + 1)
        del sorted_ints[-1]



########## Camera normal #######################################################



##
# q: quaternion
# v: 3-element array
# @see adapted from blender's math_rotation.c
#
# \note:
# Assumes a unit quaternion?
#
# in fact not, but you may want to use a unit quat, read on...
#
# Shortcut for 'q v q*' when \a v is actually a quaternion.
# This removes the need for converting a vector to a quaternion,
# calculating q's conjugate and converting back to a vector.
# It also happens to be faster (17+,24* vs * 24+,32*).
# If \a q is not a unit quaternion, then \a v will be both rotated by
# the same amount as if q was a unit quaternion, and scaled by the square of
# the length of q.
#
# For people used to python mathutils, its like:
# def mul_qt_v3(q, v): (q * Quaternion((0.0, v[0], v[1], v[2])) * q.conjugated())[1:]
#
# \note: multiplying by 3x3 matrix is ~25% faster.
##
def multiply_quaternion_vec3(q, v):
    t0 = -q[1] * v[0] - q[2] * v[1] - q[3] * v[2]
    t1 = q[0] * v[0] + q[2] * v[2] - q[3] * v[1]
    t2 = q[0] * v[1] + q[3] * v[0] - q[1] * v[2]
    i = [t1, t2, q[0] * v[2] + q[1] * v[1] - q[2] * v[0]]
    t1 = t0 * -q[1] + i[0] * q[0] - i[1] * q[3] + i[2] * q[2]
    t2 = t0 * -q[2] + i[1] * q[0] - i[2] * q[1] + i[0] * q[3]
    i[2] = t0 * -q[3] + i[2] * q[0] - i[0] * q[2] + i[1] * q[1]
    i[0] = t1
    i[1] = t2
    return i


def world_to_camera_normals(inverted_camera_quaternation, exr_x, exr_y, exr_z):
    camera_normal = np.empty([exr_x.shape[0], exr_x.shape[1], 3], dtype=np.float32)
    for i in range(exr_x.shape[0]):
        for j in range(exr_x.shape[1]):
            pixel_camera_normal = multiply_quaternion_vec3(
                inverted_camera_quaternation,
                [exr_x[i][j], exr_y[i][j], exr_z[i][j]]
            )
            camera_normal[i][j][0] = pixel_camera_normal[0]
            camera_normal[i][j][1] = pixel_camera_normal[1]
            camera_normal[i][j][2] = pixel_camera_normal[2]
    return camera_normal

# checks exr_x, exr_y, exr_z for original zero values
def normal_to_rgb(normals_to_convert, exr_x, exr_y, exr_z):
    camera_normal_rgb = normals_to_convert + 1
    camera_normal_rgb *= 127.5
    camera_normal_rgb = camera_normal_rgb.astype(np.uint8)
    # Correct case when we had an empty normal (0,0,0), and turn it to black instead of gray
    for i in range(exr_x.shape[0]):
        for j in range(exr_x.shape[1]):
            if exr_x[i][j] == 0 and exr_y[i][j] == 0 and exr_z[i][j] == 0:
                camera_normal_rgb[i][j][0] = 0
                camera_normal_rgb[i][j][1] = 0
                camera_normal_rgb[i][j][2] = 0
    return camera_normal_rgb

# Return X, Y, Z normals as numpy arrays
def read_exr_normal_file(exr_path):
    exr_file = OpenEXR.InputFile(exr_path)
    cm_dw = exr_file.header()['dataWindow']
    exr_x = np.fromstring(
        exr_file.channel(NORMALS_X_CHANNEL_NAME, Imath.PixelType(Imath.PixelType.HALF)),
        dtype=np.float16
    )
    exr_x.shape = (cm_dw.max.y - cm_dw.min.y + 1, cm_dw.max.x - cm_dw.min.x + 1) # rows, cols
    exr_y = np.fromstring(
        exr_file.channel(NORMALS_Y_CHANNEL_NAME, Imath.PixelType(Imath.PixelType.HALF)),
        dtype=np.float16
    )
    exr_y.shape = (cm_dw.max.y - cm_dw.min.y + 1, cm_dw.max.x - cm_dw.min.x + 1) # rows, cols

    exr_z = np.fromstring(
        exr_file.channel(NORMALS_Z_CHANNEL_NAME, Imath.PixelType(Imath.PixelType.HALF)),
        dtype=np.float16
    )
    exr_z.shape = (cm_dw.max.y - cm_dw.min.y + 1, cm_dw.max.x - cm_dw.min.x + 1) # rows, cols
    return exr_x, exr_y, exr_z

def generate_camera_normals(dataset_path):
    dataset_prefixes = scene_prefixes(dataset_path)
    for prefix in dataset_prefixes:
        # if it already exists, don't regenerate
        if os.path.isfile(os.path.join(dataset_path, prefix+'-cameraNormals.npy')):
            continue
        normals_path = os.path.join(dataset_path, prefix+'-normals.exr')
        exr_x, exr_y, exr_z = read_exr_normal_file(normals_path)
        masks_json = json.load(open(os.path.join(dataset_path, prefix+'-masks.json')))
        inv_q = np.asarray(masks_json['camera']['world_pose']['rotation']['inverted_quaternion'], dtype=np.float32)
        camera_normal = world_to_camera_normals(inv_q, exr_x, exr_y, exr_z)
        transposed_camera_normal = np.transpose(camera_normal, (2, 0, 1))
        np.save(os.path.join(dataset_path, prefix+'-cameraNormals.npy'), transposed_camera_normal)
        print("  Generated", prefix+'-cameraNormals.npy')





########## Move prefixes to subfolders #########################################



def move_to_subfolders(dataset_path):
    dataset_prefixes = scene_prefixes(dataset_path)
    for prefix in dataset_prefixes:
        for file_postfix, subfolder in SUBFOLDER_MAP.items():
            if not os.path.isdir(os.path.join(dataset_path, subfolder)):
                os.mkdir(os.path.join(dataset_path, subfolder))
                print ("  Created", subfolder, "folder.")
            for scene_file in glob.glob(os.path.join(dataset_path, prefix+file_postfix)):
                shutil.move(scene_file, os.path.join(dataset_path, subfolder))
                print ("    Moved", scene_file, "to", subfolder, "folder.")


########## Main ################################################################



def main():
    parser = argparse.ArgumentParser(description='Rearrange non-contiguous scenes in a dataset, generate camera normals, move to separate folders.')
    parser.add_argument('--p', required=True, help='path to dataset', metavar='/path/to/dataset')
    args = parser.parse_args()
    print("Finding non-contiguous scenes in dataset.")
    fix_non_contiguities(args.p)
    print("Generating camera normals.")
    generate_camera_normals(args.p)
    print("Separating dataset into folders.")
    move_to_subfolders(args.p)


if __name__ == "__main__":
    main()