-
Notifications
You must be signed in to change notification settings - Fork 8
/
py3d_tools.py
1799 lines (1552 loc) · 70.6 KB
/
py3d_tools.py
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import sys
import math
import warnings
from typing import List, Optional, Sequence, Tuple, Union, Any
import numpy as np
import torch
import torch.nn.functional as F
import copy
import inspect
import torch.nn as nn
Device = Union[str, torch.device]
# Default values for rotation and translation matrices.
_R = torch.eye(3)[None] # (1, 3, 3)
_T = torch.zeros(1, 3) # (1, 3)
# Provide get_origin and get_args even in Python 3.7.
if sys.version_info >= (3, 8, 0):
from typing import get_args, get_origin
elif sys.version_info >= (3, 7, 0):
def get_origin(cls): # pragma: no cover
return getattr(cls, "__origin__", None)
def get_args(cls): # pragma: no cover
return getattr(cls, "__args__", None)
else:
raise ImportError("This module requires Python 3.7+")
################################################################
## ██████╗██╗ █████╗ ███████╗███████╗███████╗███████╗ ##
## ██╔════╝██║ ██╔══██╗██╔════╝██╔════╝██╔════╝██╔════╝ ##
## ██║ ██║ ███████║███████╗███████╗█████╗ ███████╗ ##
## ██║ ██║ ██╔══██║╚════██║╚════██║██╔══╝ ╚════██║ ##
## ╚██████╗███████╗██║ ██║███████║███████║███████╗███████║ ##
## ╚═════╝╚══════╝╚═╝ ╚═╝╚══════╝╚══════╝╚══════╝╚══════╝ ##
################################################################
class Transform3d:
"""
A Transform3d object encapsulates a batch of N 3D transformations, and knows
how to transform points and normal vectors. Suppose that t is a Transform3d;
then we can do the following:
.. code-block:: python
N = len(t)
points = torch.randn(N, P, 3)
normals = torch.randn(N, P, 3)
points_transformed = t.transform_points(points) # => (N, P, 3)
normals_transformed = t.transform_normals(normals) # => (N, P, 3)
BROADCASTING
Transform3d objects supports broadcasting. Suppose that t1 and tN are
Transform3d objects with len(t1) == 1 and len(tN) == N respectively. Then we
can broadcast transforms like this:
.. code-block:: python
t1.transform_points(torch.randn(P, 3)) # => (P, 3)
t1.transform_points(torch.randn(1, P, 3)) # => (1, P, 3)
t1.transform_points(torch.randn(M, P, 3)) # => (M, P, 3)
tN.transform_points(torch.randn(P, 3)) # => (N, P, 3)
tN.transform_points(torch.randn(1, P, 3)) # => (N, P, 3)
COMBINING TRANSFORMS
Transform3d objects can be combined in two ways: composing and stacking.
Composing is function composition. Given Transform3d objects t1, t2, t3,
the following all compute the same thing:
.. code-block:: python
y1 = t3.transform_points(t2.transform_points(t1.transform_points(x)))
y2 = t1.compose(t2).compose(t3).transform_points(x)
y3 = t1.compose(t2, t3).transform_points(x)
Composing transforms should broadcast.
.. code-block:: python
if len(t1) == 1 and len(t2) == N, then len(t1.compose(t2)) == N.
We can also stack a sequence of Transform3d objects, which represents
composition along the batch dimension; then the following should compute the
same thing.
.. code-block:: python
N, M = len(tN), len(tM)
xN = torch.randn(N, P, 3)
xM = torch.randn(M, P, 3)
y1 = torch.cat([tN.transform_points(xN), tM.transform_points(xM)], dim=0)
y2 = tN.stack(tM).transform_points(torch.cat([xN, xM], dim=0))
BUILDING TRANSFORMS
We provide convenience methods for easily building Transform3d objects
as compositions of basic transforms.
.. code-block:: python
# Scale by 0.5, then translate by (1, 2, 3)
t1 = Transform3d().scale(0.5).translate(1, 2, 3)
# Scale each axis by a different amount, then translate, then scale
t2 = Transform3d().scale(1, 3, 3).translate(2, 3, 1).scale(2.0)
t3 = t1.compose(t2)
tN = t1.stack(t3, t3)
BACKPROP THROUGH TRANSFORMS
When building transforms, we can also parameterize them by Torch tensors;
in this case we can backprop through the construction and application of
Transform objects, so they could be learned via gradient descent or
predicted by a neural network.
.. code-block:: python
s1_params = torch.randn(N, requires_grad=True)
t_params = torch.randn(N, 3, requires_grad=True)
s2_params = torch.randn(N, 3, requires_grad=True)
t = Transform3d().scale(s1_params).translate(t_params).scale(s2_params)
x = torch.randn(N, 3)
y = t.transform_points(x)
loss = compute_loss(y)
loss.backward()
with torch.no_grad():
s1_params -= lr * s1_params.grad
t_params -= lr * t_params.grad
s2_params -= lr * s2_params.grad
CONVENTIONS
We adopt a right-hand coordinate system, meaning that rotation about an axis
with a positive angle results in a counter clockwise rotation.
This class assumes that transformations are applied on inputs which
are row vectors. The internal representation of the Nx4x4 transformation
matrix is of the form:
.. code-block:: python
M = [
[Rxx, Ryx, Rzx, 0],
[Rxy, Ryy, Rzy, 0],
[Rxz, Ryz, Rzz, 0],
[Tx, Ty, Tz, 1],
]
To apply the transformation to points which are row vectors, the M matrix
can be pre multiplied by the points:
.. code-block:: python
points = [[0, 1, 2]] # (1 x 3) xyz coordinates of a point
transformed_points = points * M
"""
def __init__(
self,
dtype: torch.dtype = torch.float32,
device: Device = "cpu",
matrix: Optional[torch.Tensor] = None,
) -> None:
"""
Args:
dtype: The data type of the transformation matrix.
to be used if `matrix = None`.
device: The device for storing the implemented transformation.
If `matrix != None`, uses the device of input `matrix`.
matrix: A tensor of shape (4, 4) or of shape (minibatch, 4, 4)
representing the 4x4 3D transformation matrix.
If `None`, initializes with identity using
the specified `device` and `dtype`.
"""
if matrix is None:
self._matrix = torch.eye(4, dtype=dtype, device=device).view(1, 4, 4)
else:
if matrix.ndim not in (2, 3):
raise ValueError('"matrix" has to be a 2- or a 3-dimensional tensor.')
if matrix.shape[-2] != 4 or matrix.shape[-1] != 4:
raise ValueError(
'"matrix" has to be a tensor of shape (minibatch, 4, 4)'
)
# set dtype and device from matrix
dtype = matrix.dtype
device = matrix.device
self._matrix = matrix.view(-1, 4, 4)
self._transforms = [] # store transforms to compose
self._lu = None
self.device = make_device(device)
self.dtype = dtype
def __len__(self) -> int:
return self.get_matrix().shape[0]
def __getitem__(
self, index: Union[int, List[int], slice, torch.Tensor]
) -> "Transform3d":
"""
Args:
index: Specifying the index of the transform to retrieve.
Can be an int, slice, list of ints, boolean, long tensor.
Supports negative indices.
Returns:
Transform3d object with selected transforms. The tensors are not cloned.
"""
if isinstance(index, int):
index = [index]
return self.__class__(matrix=self.get_matrix()[index])
def compose(self, *others: "Transform3d") -> "Transform3d":
"""
Return a new Transform3d representing the composition of self with the
given other transforms, which will be stored as an internal list.
Args:
*others: Any number of Transform3d objects
Returns:
A new Transform3d with the stored transforms
"""
out = Transform3d(dtype=self.dtype, device=self.device)
out._matrix = self._matrix.clone()
for other in others:
if not isinstance(other, Transform3d):
msg = "Only possible to compose Transform3d objects; got %s"
raise ValueError(msg % type(other))
out._transforms = self._transforms + list(others)
return out
def get_matrix(self) -> torch.Tensor:
"""
Return a matrix which is the result of composing this transform
with others stored in self.transforms. Where necessary transforms
are broadcast against each other.
For example, if self.transforms contains transforms t1, t2, and t3, and
given a set of points x, the following should be true:
.. code-block:: python
y1 = t1.compose(t2, t3).transform(x)
y2 = t3.transform(t2.transform(t1.transform(x)))
y1.get_matrix() == y2.get_matrix()
Returns:
A transformation matrix representing the composed inputs.
"""
composed_matrix = self._matrix.clone()
if len(self._transforms) > 0:
for other in self._transforms:
other_matrix = other.get_matrix()
composed_matrix = _broadcast_bmm(composed_matrix, other_matrix)
return composed_matrix
def _get_matrix_inverse(self) -> torch.Tensor:
"""
Return the inverse of self._matrix.
"""
return torch.inverse(self._matrix)
def inverse(self, invert_composed: bool = False) -> "Transform3d":
"""
Returns a new Transform3d object that represents an inverse of the
current transformation.
Args:
invert_composed:
- True: First compose the list of stored transformations
and then apply inverse to the result. This is
potentially slower for classes of transformations
with inverses that can be computed efficiently
(e.g. rotations and translations).
- False: Invert the individual stored transformations
independently without composing them.
Returns:
A new Transform3d object containing the inverse of the original
transformation.
"""
tinv = Transform3d(dtype=self.dtype, device=self.device)
if invert_composed:
# first compose then invert
tinv._matrix = torch.inverse(self.get_matrix())
else:
# self._get_matrix_inverse() implements efficient inverse
# of self._matrix
i_matrix = self._get_matrix_inverse()
# 2 cases:
if len(self._transforms) > 0:
# a) Either we have a non-empty list of transforms:
# Here we take self._matrix and append its inverse at the
# end of the reverted _transforms list. After composing
# the transformations with get_matrix(), this correctly
# right-multiplies by the inverse of self._matrix
# at the end of the composition.
tinv._transforms = [t.inverse() for t in reversed(self._transforms)]
last = Transform3d(dtype=self.dtype, device=self.device)
last._matrix = i_matrix
tinv._transforms.append(last)
else:
# b) Or there are no stored transformations
# we just set inverted matrix
tinv._matrix = i_matrix
return tinv
def stack(self, *others: "Transform3d") -> "Transform3d":
"""
Return a new batched Transform3d representing the batch elements from
self and all the given other transforms all batched together.
Args:
*others: Any number of Transform3d objects
Returns:
A new Transform3d.
"""
transforms = [self] + list(others)
matrix = torch.cat([t.get_matrix() for t in transforms], dim=0)
out = Transform3d(dtype=self.dtype, device=self.device)
out._matrix = matrix
return out
def transform_points(self, points, eps: Optional[float] = None) -> torch.Tensor:
"""
Use this transform to transform a set of 3D points. Assumes row major
ordering of the input points.
Args:
points: Tensor of shape (P, 3) or (N, P, 3)
eps: If eps!=None, the argument is used to clamp the
last coordinate before performing the final division.
The clamping corresponds to:
last_coord := (last_coord.sign() + (last_coord==0)) *
torch.clamp(last_coord.abs(), eps),
i.e. the last coordinates that are exactly 0 will
be clamped to +eps.
Returns:
points_out: points of shape (N, P, 3) or (P, 3) depending
on the dimensions of the transform
"""
points_batch = points.clone()
if points_batch.dim() == 2:
points_batch = points_batch[None] # (P, 3) -> (1, P, 3)
if points_batch.dim() != 3:
msg = "Expected points to have dim = 2 or dim = 3: got shape %r"
raise ValueError(msg % repr(points.shape))
N, P, _3 = points_batch.shape
ones = torch.ones(N, P, 1, dtype=points.dtype, device=points.device)
points_batch = torch.cat([points_batch, ones], dim=2)
composed_matrix = self.get_matrix()
points_out = _broadcast_bmm(points_batch, composed_matrix)
denom = points_out[..., 3:] # denominator
if eps is not None:
denom_sign = denom.sign() + (denom == 0.0).type_as(denom)
denom = denom_sign * torch.clamp(denom.abs(), eps)
points_out = points_out[..., :3] / denom
# When transform is (1, 4, 4) and points is (P, 3) return
# points_out of shape (P, 3)
if points_out.shape[0] == 1 and points.dim() == 2:
points_out = points_out.reshape(points.shape)
return points_out
def transform_normals(self, normals) -> torch.Tensor:
"""
Use this transform to transform a set of normal vectors.
Args:
normals: Tensor of shape (P, 3) or (N, P, 3)
Returns:
normals_out: Tensor of shape (P, 3) or (N, P, 3) depending
on the dimensions of the transform
"""
if normals.dim() not in [2, 3]:
msg = "Expected normals to have dim = 2 or dim = 3: got shape %r"
raise ValueError(msg % (normals.shape,))
composed_matrix = self.get_matrix()
# TODO: inverse is bad! Solve a linear system instead
mat = composed_matrix[:, :3, :3]
normals_out = _broadcast_bmm(normals, mat.transpose(1, 2).inverse())
# This doesn't pass unit tests. TODO investigate further
# if self._lu is None:
# self._lu = self._matrix[:, :3, :3].transpose(1, 2).lu()
# normals_out = normals.lu_solve(*self._lu)
# When transform is (1, 4, 4) and normals is (P, 3) return
# normals_out of shape (P, 3)
if normals_out.shape[0] == 1 and normals.dim() == 2:
normals_out = normals_out.reshape(normals.shape)
return normals_out
def translate(self, *args, **kwargs) -> "Transform3d":
return self.compose(
Translate(device=self.device, dtype=self.dtype, *args, **kwargs)
)
def scale(self, *args, **kwargs) -> "Transform3d":
return self.compose(
Scale(device=self.device, dtype=self.dtype, *args, **kwargs)
)
def rotate(self, *args, **kwargs) -> "Transform3d":
return self.compose(
Rotate(device=self.device, dtype=self.dtype, *args, **kwargs)
)
def rotate_axis_angle(self, *args, **kwargs) -> "Transform3d":
return self.compose(
RotateAxisAngle(device=self.device, dtype=self.dtype, *args, **kwargs)
)
def clone(self) -> "Transform3d":
"""
Deep copy of Transforms object. All internal tensors are cloned
individually.
Returns:
new Transforms object.
"""
other = Transform3d(dtype=self.dtype, device=self.device)
if self._lu is not None:
other._lu = [elem.clone() for elem in self._lu]
other._matrix = self._matrix.clone()
other._transforms = [t.clone() for t in self._transforms]
return other
def to(
self,
device: Device,
copy: bool = False,
dtype: Optional[torch.dtype] = None,
) -> "Transform3d":
"""
Match functionality of torch.Tensor.to()
If copy = True or the self Tensor is on a different device, the
returned tensor is a copy of self with the desired torch.device.
If copy = False and the self Tensor already has the correct torch.device,
then self is returned.
Args:
device: Device (as str or torch.device) for the new tensor.
copy: Boolean indicator whether or not to clone self. Default False.
dtype: If not None, casts the internal tensor variables
to a given torch.dtype.
Returns:
Transform3d object.
"""
device_ = make_device(device)
dtype_ = self.dtype if dtype is None else dtype
skip_to = self.device == device_ and self.dtype == dtype_
if not copy and skip_to:
return self
other = self.clone()
if skip_to:
return other
other.device = device_
other.dtype = dtype_
other._matrix = other._matrix.to(device=device_, dtype=dtype_)
other._transforms = [
t.to(device_, copy=copy, dtype=dtype_) for t in other._transforms
]
return other
def cpu(self) -> "Transform3d":
return self.to("cpu")
def cuda(self) -> "Transform3d":
return self.to("cuda")
class Translate(Transform3d):
def __init__(
self,
x,
y=None,
z=None,
dtype: torch.dtype = torch.float32,
device: Optional[Device] = None,
) -> None:
"""
Create a new Transform3d representing 3D translations.
Option I: Translate(xyz, dtype=torch.float32, device='cpu')
xyz should be a tensor of shape (N, 3)
Option II: Translate(x, y, z, dtype=torch.float32, device='cpu')
Here x, y, and z will be broadcast against each other and
concatenated to form the translation. Each can be:
- A python scalar
- A torch scalar
- A 1D torch tensor
"""
xyz = _handle_input(x, y, z, dtype, device, "Translate")
super().__init__(device=xyz.device, dtype=dtype)
N = xyz.shape[0]
mat = torch.eye(4, dtype=dtype, device=self.device)
mat = mat.view(1, 4, 4).repeat(N, 1, 1)
mat[:, 3, :3] = xyz
self._matrix = mat
def _get_matrix_inverse(self) -> torch.Tensor:
"""
Return the inverse of self._matrix.
"""
inv_mask = self._matrix.new_ones([1, 4, 4])
inv_mask[0, 3, :3] = -1.0
i_matrix = self._matrix * inv_mask
return i_matrix
class Rotate(Transform3d):
def __init__(
self,
R: torch.Tensor,
dtype: torch.dtype = torch.float32,
device: Optional[Device] = None,
orthogonal_tol: float = 1e-5,
) -> None:
"""
Create a new Transform3d representing 3D rotation using a rotation
matrix as the input.
Args:
R: a tensor of shape (3, 3) or (N, 3, 3)
orthogonal_tol: tolerance for the test of the orthogonality of R
"""
device_ = get_device(R, device)
super().__init__(device=device_, dtype=dtype)
if R.dim() == 2:
R = R[None]
if R.shape[-2:] != (3, 3):
msg = "R must have shape (3, 3) or (N, 3, 3); got %s"
raise ValueError(msg % repr(R.shape))
R = R.to(device=device_, dtype=dtype)
_check_valid_rotation_matrix(R, tol=orthogonal_tol)
N = R.shape[0]
mat = torch.eye(4, dtype=dtype, device=device_)
mat = mat.view(1, 4, 4).repeat(N, 1, 1)
mat[:, :3, :3] = R
self._matrix = mat
def _get_matrix_inverse(self) -> torch.Tensor:
"""
Return the inverse of self._matrix.
"""
return self._matrix.permute(0, 2, 1).contiguous()
class TensorAccessor(nn.Module):
"""
A helper class to be used with the __getitem__ method. This can be used for
getting/setting the values for an attribute of a class at one particular
index. This is useful when the attributes of a class are batched tensors
and one element in the batch needs to be modified.
"""
def __init__(self, class_object, index: Union[int, slice]) -> None:
"""
Args:
class_object: this should be an instance of a class which has
attributes which are tensors representing a batch of
values.
index: int/slice, an index indicating the position in the batch.
In __setattr__ and __getattr__ only the value of class
attributes at this index will be accessed.
"""
self.__dict__["class_object"] = class_object
self.__dict__["index"] = index
def __setattr__(self, name: str, value: Any):
"""
Update the attribute given by `name` to the value given by `value`
at the index specified by `self.index`.
Args:
name: str, name of the attribute.
value: value to set the attribute to.
"""
v = getattr(self.class_object, name)
if not torch.is_tensor(v):
msg = "Can only set values on attributes which are tensors; got %r"
raise AttributeError(msg % type(v))
# Convert the attribute to a tensor if it is not a tensor.
if not torch.is_tensor(value):
value = torch.tensor(
value, device=v.device, dtype=v.dtype, requires_grad=v.requires_grad
)
# Check the shapes match the existing shape and the shape of the index.
if v.dim() > 1 and value.dim() > 1 and value.shape[1:] != v.shape[1:]:
msg = "Expected value to have shape %r; got %r"
raise ValueError(msg % (v.shape, value.shape))
if (
v.dim() == 0
and isinstance(self.index, slice)
and len(value) != len(self.index)
):
msg = "Expected value to have len %r; got %r"
raise ValueError(msg % (len(self.index), len(value)))
self.class_object.__dict__[name][self.index] = value
def __getattr__(self, name: str):
"""
Return the value of the attribute given by "name" on self.class_object
at the index specified in self.index.
Args:
name: string of the attribute name
"""
if hasattr(self.class_object, name):
return self.class_object.__dict__[name][self.index]
else:
msg = "Attribute %s not found on %r"
return AttributeError(msg % (name, self.class_object.__name__))
BROADCAST_TYPES = (float, int, list, tuple, torch.Tensor, np.ndarray)
class TensorProperties(nn.Module):
"""
A mix-in class for storing tensors as properties with helper methods.
"""
def __init__(
self,
dtype: torch.dtype = torch.float32,
device: Device = "cpu",
**kwargs,
) -> None:
"""
Args:
dtype: data type to set for the inputs
device: Device (as str or torch.device)
kwargs: any number of keyword arguments. Any arguments which are
of type (float/int/list/tuple/tensor/array) are broadcasted and
other keyword arguments are set as attributes.
"""
super().__init__()
self.device = make_device(device)
self._N = 0
if kwargs is not None:
# broadcast all inputs which are float/int/list/tuple/tensor/array
# set as attributes anything else e.g. strings, bools
args_to_broadcast = {}
for k, v in kwargs.items():
if v is None or isinstance(v, (str, bool)):
setattr(self, k, v)
elif isinstance(v, BROADCAST_TYPES):
args_to_broadcast[k] = v
else:
msg = "Arg %s with type %r is not broadcastable"
warnings.warn(msg % (k, type(v)))
names = args_to_broadcast.keys()
# convert from type dict.values to tuple
values = tuple(v for v in args_to_broadcast.values())
if len(values) > 0:
broadcasted_values = convert_to_tensors_and_broadcast(
*values, device=device
)
# Set broadcasted values as attributes on self.
for i, n in enumerate(names):
setattr(self, n, broadcasted_values[i])
if self._N == 0:
self._N = broadcasted_values[i].shape[0]
def __len__(self) -> int:
return self._N
def isempty(self) -> bool:
return self._N == 0
def __getitem__(self, index: Union[int, slice]) -> TensorAccessor:
"""
Args:
index: an int or slice used to index all the fields.
Returns:
if `index` is an index int/slice return a TensorAccessor class
with getattribute/setattribute methods which return/update the value
at the index in the original class.
"""
if isinstance(index, (int, slice)):
return TensorAccessor(class_object=self, index=index)
msg = "Expected index of type int or slice; got %r"
raise ValueError(msg % type(index))
# pyre-fixme[14]: `to` overrides method defined in `Module` inconsistently.
def to(self, device: Device = "cpu") -> "TensorProperties":
"""
In place operation to move class properties which are tensors to a
specified device. If self has a property "device", update this as well.
"""
device_ = make_device(device)
for k in dir(self):
v = getattr(self, k)
if k == "device":
setattr(self, k, device_)
if torch.is_tensor(v) and v.device != device_:
setattr(self, k, v.to(device_))
return self
def cpu(self) -> "TensorProperties":
return self.to("cpu")
# pyre-fixme[14]: `cuda` overrides method defined in `Module` inconsistently.
def cuda(self, device: Optional[int] = None) -> "TensorProperties":
return self.to(f"cuda:{device}" if device is not None else "cuda")
def clone(self, other) -> "TensorProperties":
"""
Update the tensor properties of other with the cloned properties of self.
"""
for k in dir(self):
v = getattr(self, k)
if inspect.ismethod(v) or k.startswith("__"):
continue
if torch.is_tensor(v):
v_clone = v.clone()
else:
v_clone = copy.deepcopy(v)
setattr(other, k, v_clone)
return other
def gather_props(self, batch_idx) -> "TensorProperties":
"""
This is an in place operation to reformat all tensor class attributes
based on a set of given indices using torch.gather. This is useful when
attributes which are batched tensors e.g. shape (N, 3) need to be
multiplied with another tensor which has a different first dimension
e.g. packed vertices of shape (V, 3).
Example
.. code-block:: python
self.specular_color = (N, 3) tensor of specular colors for each mesh
A lighting calculation may use
.. code-block:: python
verts_packed = meshes.verts_packed() # (V, 3)
To multiply these two tensors the batch dimension needs to be the same.
To achieve this we can do
.. code-block:: python
batch_idx = meshes.verts_packed_to_mesh_idx() # (V)
This gives index of the mesh for each vertex in verts_packed.
.. code-block:: python
self.gather_props(batch_idx)
self.specular_color = (V, 3) tensor with the specular color for
each packed vertex.
torch.gather requires the index tensor to have the same shape as the
input tensor so this method takes care of the reshaping of the index
tensor to use with class attributes with arbitrary dimensions.
Args:
batch_idx: shape (B, ...) where `...` represents an arbitrary
number of dimensions
Returns:
self with all properties reshaped. e.g. a property with shape (N, 3)
is transformed to shape (B, 3).
"""
# Iterate through the attributes of the class which are tensors.
for k in dir(self):
v = getattr(self, k)
if torch.is_tensor(v):
if v.shape[0] > 1:
# There are different values for each batch element
# so gather these using the batch_idx.
# First clone the input batch_idx tensor before
# modifying it.
_batch_idx = batch_idx.clone()
idx_dims = _batch_idx.shape
tensor_dims = v.shape
if len(idx_dims) > len(tensor_dims):
msg = "batch_idx cannot have more dimensions than %s. "
msg += "got shape %r and %s has shape %r"
raise ValueError(msg % (k, idx_dims, k, tensor_dims))
if idx_dims != tensor_dims:
# To use torch.gather the index tensor (_batch_idx) has
# to have the same shape as the input tensor.
new_dims = len(tensor_dims) - len(idx_dims)
new_shape = idx_dims + (1,) * new_dims
expand_dims = (-1,) + tensor_dims[1:]
_batch_idx = _batch_idx.view(*new_shape)
_batch_idx = _batch_idx.expand(*expand_dims)
v = v.gather(0, _batch_idx)
setattr(self, k, v)
return self
class CamerasBase(TensorProperties):
"""
`CamerasBase` implements a base class for all cameras.
For cameras, there are four different coordinate systems (or spaces)
- World coordinate system: This is the system the object lives - the world.
- Camera view coordinate system: This is the system that has its origin on the camera
and the and the Z-axis perpendicular to the image plane.
In PyTorch3D, we assume that +X points left, and +Y points up and
+Z points out from the image plane.
The transformation from world --> view happens after applying a rotation (R)
and translation (T)
- NDC coordinate system: This is the normalized coordinate system that confines
in a volume the rendered part of the object or scene. Also known as view volume.
For square images, given the PyTorch3D convention, (+1, +1, znear)
is the top left near corner, and (-1, -1, zfar) is the bottom right far
corner of the volume.
The transformation from view --> NDC happens after applying the camera
projection matrix (P) if defined in NDC space.
For non square images, we scale the points such that smallest side
has range [-1, 1] and the largest side has range [-u, u], with u > 1.
- Screen coordinate system: This is another representation of the view volume with
the XY coordinates defined in image space instead of a normalized space.
A better illustration of the coordinate systems can be found in
pytorch3d/docs/notes/cameras.md.
It defines methods that are common to all camera models:
- `get_camera_center` that returns the optical center of the camera in
world coordinates
- `get_world_to_view_transform` which returns a 3D transform from
world coordinates to the camera view coordinates (R, T)
- `get_full_projection_transform` which composes the projection
transform (P) with the world-to-view transform (R, T)
- `transform_points` which takes a set of input points in world coordinates and
projects to the space the camera is defined in (NDC or screen)
- `get_ndc_camera_transform` which defines the transform from screen/NDC to
PyTorch3D's NDC space
- `transform_points_ndc` which takes a set of points in world coordinates and
projects them to PyTorch3D's NDC space
- `transform_points_screen` which takes a set of points in world coordinates and
projects them to screen space
For each new camera, one should implement the `get_projection_transform`
routine that returns the mapping from camera view coordinates to camera
coordinates (NDC or screen).
Another useful function that is specific to each camera model is
`unproject_points` which sends points from camera coordinates (NDC or screen)
back to camera view or world coordinates depending on the `world_coordinates`
boolean argument of the function.
"""
# Used in __getitem__ to index the relevant fields
# When creating a new camera, this should be set in the __init__
_FIELDS: Tuple[str, ...] = ()
# Names of fields which are a constant property of the whole batch, rather
# than themselves a batch of data.
# When joining objects into a batch, they will have to agree.
_SHARED_FIELDS: Tuple[str, ...] = ()
def get_projection_transform(self):
"""
Calculate the projective transformation matrix.
Args:
**kwargs: parameters for the projection can be passed in as keyword
arguments to override the default values set in `__init__`.
Return:
a `Transform3d` object which represents a batch of projection
matrices of shape (N, 3, 3)
"""
raise NotImplementedError()
def unproject_points(self, xy_depth: torch.Tensor, **kwargs):
"""
Transform input points from camera coodinates (NDC or screen)
to the world / camera coordinates.
Each of the input points `xy_depth` of shape (..., 3) is
a concatenation of the x, y location and its depth.
For instance, for an input 2D tensor of shape `(num_points, 3)`
`xy_depth` takes the following form:
`xy_depth[i] = [x[i], y[i], depth[i]]`,
for a each point at an index `i`.
The following example demonstrates the relationship between
`transform_points` and `unproject_points`:
.. code-block:: python
cameras = # camera object derived from CamerasBase
xyz = # 3D points of shape (batch_size, num_points, 3)
# transform xyz to the camera view coordinates
xyz_cam = cameras.get_world_to_view_transform().transform_points(xyz)
# extract the depth of each point as the 3rd coord of xyz_cam
depth = xyz_cam[:, :, 2:]
# project the points xyz to the camera
xy = cameras.transform_points(xyz)[:, :, :2]
# append depth to xy
xy_depth = torch.cat((xy, depth), dim=2)
# unproject to the world coordinates
xyz_unproj_world = cameras.unproject_points(xy_depth, world_coordinates=True)
print(torch.allclose(xyz, xyz_unproj_world)) # True
# unproject to the camera coordinates
xyz_unproj = cameras.unproject_points(xy_depth, world_coordinates=False)
print(torch.allclose(xyz_cam, xyz_unproj)) # True
Args:
xy_depth: torch tensor of shape (..., 3).
world_coordinates: If `True`, unprojects the points back to world
coordinates using the camera extrinsics `R` and `T`.
`False` ignores `R` and `T` and unprojects to
the camera view coordinates.
from_ndc: If `False` (default), assumes xy part of input is in
NDC space if self.in_ndc(), otherwise in screen space. If
`True`, assumes xy is in NDC space even if the camera
is defined in screen space.
Returns
new_points: unprojected points with the same shape as `xy_depth`.
"""
raise NotImplementedError()
def get_camera_center(self, **kwargs) -> torch.Tensor:
"""
Return the 3D location of the camera optical center
in the world coordinates.
Args:
**kwargs: parameters for the camera extrinsics can be passed in
as keyword arguments to override the default values
set in __init__.
Setting T here will update the values set in init as this
value may be needed later on in the rendering pipeline e.g. for
lighting calculations.
Returns:
C: a batch of 3D locations of shape (N, 3) denoting
the locations of the center of each camera in the batch.
"""
w2v_trans = self.get_world_to_view_transform(**kwargs)
P = w2v_trans.inverse().get_matrix()
# the camera center is the translation component (the first 3 elements
# of the last row) of the inverted world-to-view
# transform (4x4 RT matrix)
C = P[:, 3, :3]
return C
def get_world_to_view_transform(self, **kwargs) -> Transform3d:
"""
Return the world-to-view transform.
Args:
**kwargs: parameters for the camera extrinsics can be passed in
as keyword arguments to override the default values
set in __init__.
Setting R and T here will update the values set in init as these
values may be needed later on in the rendering pipeline e.g. for
lighting calculations.
Returns:
A Transform3d object which represents a batch of transforms
of shape (N, 3, 3)
"""
R: torch.Tensor = kwargs.get("R", self.R)
T: torch.Tensor = kwargs.get("T", self.T)
self.R = R # pyre-ignore[16]
self.T = T # pyre-ignore[16]
world_to_view_transform = get_world_to_view_transform(R=R, T=T)
return world_to_view_transform
def get_full_projection_transform(self, **kwargs) -> Transform3d:
"""
Return the full world-to-camera transform composing the
world-to-view and view-to-camera transforms.
If camera is defined in NDC space, the projected points are in NDC space.
If camera is defined in screen space, the projected points are in screen space.
Args:
**kwargs: parameters for the projection transforms can be passed in
as keyword arguments to override the default values
set in __init__.
Setting R and T here will update the values set in init as these
values may be needed later on in the rendering pipeline e.g. for
lighting calculations.
Returns:
a Transform3d object which represents a batch of transforms
of shape (N, 3, 3)
"""
self.R: torch.Tensor = kwargs.get("R", self.R) # pyre-ignore[16]
self.T: torch.Tensor = kwargs.get("T", self.T) # pyre-ignore[16]