Spaces:
Runtime error
Runtime error
File size: 3,528 Bytes
fc16538 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
# TRI-VIDAR - Copyright 2022 Toyota Research Institute. All rights reserved.
import os
import pickle as pkl
import cv2
import numpy as np
import torchvision.transforms as transforms
from vidar.utils.decorators import multi_write
from vidar.utils.types import is_tensor, is_numpy
def create_folder(filename):
"""Create a new folder if it doesn't exist"""
if '/' in filename:
os.makedirs(os.path.dirname(filename), exist_ok=True)
def write_pickle(filename, data):
"""
Write a pickle file
Parameters
----------
filename : String
File where the pickle file will be saved
data : Value
Data to be saved
"""
create_folder(filename)
if not filename.endswith('.pkl'):
filename = filename + '.pkl'
pkl.dump(data, open(filename, 'wb'))
def write_npz(filename, data):
"""
Write a numpy compressed file
Parameters
----------
filename : String
File where the numpy file will be saved
data : Value
Data to be saved
"""
np.savez_compressed(filename, **data)
@multi_write
@multi_write
def write_image(filename, image):
"""
Write an image to file
Parameters
----------
filename : String
File where image will be saved
image : np.Array [H,W,3]
RGB image
"""
# Create folder if it doesn't exist
create_folder(filename)
# If image is a tensor
if is_tensor(image):
if len(image.shape) == 4:
image = image[0]
image = image.detach().cpu().numpy().transpose(1, 2, 0)
cv2.imwrite(filename, image[:, :, ::-1] * 255)
# If image is a numpy array
elif is_numpy(image):
cv2.imwrite(filename, image[:, :, ::-1] * 255)
# Otherwise, assume it's a PIL image
else:
image.save(filename)
@multi_write
def write_depth(filename, depth, intrinsics=None):
"""
Write a depth map to file, and optionally its corresponding intrinsics.
Parameters
----------
filename : String
File where depth map will be saved (.npz or .png)
depth : np.Array or torch.Tensor
Depth map [H,W]
intrinsics : np.Array
Optional camera intrinsics matrix [3,3]
"""
# If depth is a tensor
if is_tensor(depth):
depth = depth.detach().squeeze().cpu().numpy()
# If intrinsics is a tensor
if is_tensor(intrinsics):
intrinsics = intrinsics.detach().cpu().numpy()
# If we are saving as a .npz
if filename.endswith('.npz'):
np.savez_compressed(filename, depth=depth, intrinsics=intrinsics)
# If we are saving as a .png
elif filename.endswith('.png'):
depth = transforms.ToPILImage()((depth * 256).astype(np.int32))
depth.save(filename)
# Something is wrong
else:
raise NotImplementedError('Depth filename not valid.')
@multi_write
def write_optical_flow(filename, optflow):
"""
Write a depth map to file, and optionally its corresponding intrinsics.
Parameters
----------
filename : String
File where depth map will be saved (.npz or .png)
optflow : np.Array or torch.Tensor
Optical flow map [H,W]
"""
# If depth is a tensor
if is_tensor(optflow):
optflow = optflow.detach().squeeze().cpu().numpy()
# If we are saving as a .npz
if filename.endswith('.npz'):
np.savez_compressed(filename, optflow=optflow)
# Something is wrong
else:
raise NotImplementedError('Optical flow filename not valid.') |