Spaces:
Running
Running
File size: 1,885 Bytes
9665c2c |
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 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
import numpy as np
from scipy.spatial.transform import Rotation
def crop_map(raster, xy, size, seed=None):
h, w = raster.shape[-2:]
state = np.random.RandomState(seed)
top = state.randint(0, h - size + 1)
left = state.randint(0, w - size + 1)
raster = raster[..., top : top + size, left : left + size]
xy -= np.array([left, top])
return raster, xy
def random_rot90(raster, xy, heading, seed=None):
rot = np.random.RandomState(seed).randint(0, 4)
heading = (heading + rot * np.pi / 2) % (2 * np.pi)
h, w = raster.shape[-2:]
if rot == 0:
xy2 = xy
elif rot == 2:
xy2 = np.array([w, h]) - 1 - xy
elif rot == 1:
xy2 = np.array([xy[1], w - 1 - xy[0]])
elif rot == 3:
xy2 = np.array([h - 1 - xy[1], xy[0]])
else:
raise ValueError(rot)
raster = np.rot90(raster, rot, axes=(-2, -1))
return raster, xy2, heading
def random_flip(image, raster, xy, heading, seed=None):
state = np.random.RandomState(seed)
if state.rand() > 0.5: # no flip
return image, raster, xy, heading
image = image[:, ::-1]
h, w = raster.shape[-2:]
if state.rand() > 0.5: # flip x
raster = raster[..., :, ::-1]
xy = np.array([w - 1 - xy[0], xy[1]])
heading = np.pi - heading
else: # flip y
raster = raster[..., ::-1, :]
xy = np.array([xy[0], h - 1 - xy[1]])
heading = -heading
heading = heading % (2 * np.pi)
return image, raster, xy, heading
def decompose_rotmat(R_c2w):
R_cv2xyz = Rotation.from_euler("X", -90, degrees=True)
rot_w2c = R_cv2xyz * Rotation.from_matrix(R_c2w).inv()
roll, pitch, yaw = rot_w2c.as_euler("YXZ", degrees=True)
# R_plane2c = Rotation.from_euler("ZX", [roll, pitch], degrees=True).as_matrix()
return roll, pitch, yaw
|