Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save Shreeyak/cfdf29a915d10354996927e63b39a483 to your computer and use it in GitHub Desktop.
Save Shreeyak/cfdf29a915d10354996927e63b39a483 to your computer and use it in GitHub Desktop.
For a Greppy Metaverse dataset: rearrange non-contiguous scenes in a dataset, generate camera normals, move to separate folders.
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 = masks_json['camera']['world_pose']['rotation']['inverted_quaternion']
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment