coordinates = dense_coordinates(tensor.shape)
(shape: Union[list, torch.Size])
| 244 | |
| 245 | |
| 246 | def dense_coordinates(shape: Union[list, torch.Size]): |
| 247 | """ |
| 248 | coordinates = dense_coordinates(tensor.shape) |
| 249 | """ |
| 250 | r""" |
| 251 | Assume the input to have BxCxD1xD2x....xDN format. |
| 252 | |
| 253 | If the shape of the tensor do not change, use |
| 254 | """ |
| 255 | spatial_dim = len(shape) - 2 |
| 256 | assert ( |
| 257 | spatial_dim > 0 |
| 258 | ), "Invalid shape. Shape must be batch x channel x spatial dimensions." |
| 259 | |
| 260 | # Generate coordinates |
| 261 | size = [i for i in shape] |
| 262 | B = size[0] |
| 263 | coordinates = torch.from_numpy( |
| 264 | np.stack( |
| 265 | [ |
| 266 | s.reshape(-1) |
| 267 | for s in np.meshgrid( |
| 268 | np.linspace(0, B - 1, B), |
| 269 | *(np.linspace(0, s - 1, s) for s in size[2:]), |
| 270 | indexing="ij", |
| 271 | ) |
| 272 | ], |
| 273 | 1, |
| 274 | ) |
| 275 | ).int() |
| 276 | return coordinates |
| 277 | |
| 278 | |
| 279 | def to_sparse(x: torch.Tensor, format: str = None, coordinates=None, device=None): |
no outgoing calls