| 1023 | |
| 1024 | |
| 1025 | class Ray: |
| 1026 | def __init__( |
| 1027 | self, |
| 1028 | origins_w: torch.Tensor, # (b, *m_shape, 3) |
| 1029 | directions_w: torch.Tensor, # (b, *m_shape, 3) |
| 1030 | ): |
| 1031 | self.origins_w = origins_w |
| 1032 | self.directions_w = directions_w |
| 1033 | |
| 1034 | def to(self, device: torch.device) -> 'Ray': |
| 1035 | for attr_name in ['origins_w', 'directions_w']: |
| 1036 | arr = getattr(self, attr_name, None) |
| 1037 | if arr is not None: |
| 1038 | setattr(self, attr_name, arr.to(device=device)) |
| 1039 | return self |
| 1040 | |
| 1041 | def reshape(self, *shape: T.List[int]): |
| 1042 | self.origins_w = self.origins_w.reshape(*shape) |
| 1043 | self.directions_w = self.directions_w.reshape(*shape) |
| 1044 | return self |
| 1045 | |
| 1046 | def chunk(self, chunks: int, dim: int = 0) -> T.List['Ray']: |
| 1047 | """Return a""" |
| 1048 | origins_w_list = self.origins_w.chunk(chunks, dim) |
| 1049 | directions_w_list = self.directions_w.chunk(chunks, dim) |
| 1050 | rays = [] |
| 1051 | for origins_w, directions_w in zip(origins_w_list, directions_w_list): |
| 1052 | ray = Ray( |
| 1053 | origins_w=origins_w, |
| 1054 | directions_w=directions_w, |
| 1055 | ) |
| 1056 | rays.append(ray) |
| 1057 | return rays |
| 1058 | |
| 1059 | def random_perturb_direction( |
| 1060 | self, |
| 1061 | shift: T.Optional[float], |
| 1062 | angle: T.Optional[float], |
| 1063 | ): |
| 1064 | """ |
| 1065 | Perturb the rays by randomly shifting the ray origin by [-shift, shift], |
| 1066 | and randomly rotating the ray direction with [-angle, angle] in degree. |
| 1067 | |
| 1068 | Args: |
| 1069 | shift: |
| 1070 | [-shift, shift] |
| 1071 | angle: |
| 1072 | [-angle, angle] in degrees |
| 1073 | rng: |
| 1074 | the random generator |
| 1075 | """ |
| 1076 | |
| 1077 | if shift is not None and math.fabs(shift) > 1e-6: |
| 1078 | r_shifts = (torch.rand_like(self.origins_w) - 0.5) * 2. * shift # [-shift, shift] |
| 1079 | self.origins_w = self.origins_w + r_shifts # (b, *m_shape, 3) |
| 1080 | if angle is not None and math.fabs(angle) > 1e-3: |
| 1081 | r_angles = ( |
| 1082 | torch.rand_like(self.directions_w) - 0.5 |
no outgoing calls
no test coverage detected