| 320 | |
| 321 | |
| 322 | def point_tracking( |
| 323 | F: torch.Tensor, |
| 324 | F0: torch.Tensor, |
| 325 | handle_points: torch.Tensor, |
| 326 | handle_points0: torch.Tensor, |
| 327 | r2: int = 3, |
| 328 | device: torch.device = torch.device("cuda"), |
| 329 | ) -> torch.Tensor: |
| 330 | |
| 331 | n = handle_points.shape[0] # Number of handle points |
| 332 | new_handle_points = torch.zeros_like(handle_points) |
| 333 | |
| 334 | for i in range(n): |
| 335 | # Compute the patch around the handle point |
| 336 | patch = utils.create_square_mask( |
| 337 | F.shape[2], F.shape[3], center=handle_points[i].tolist(), radius=r2 |
| 338 | ).to(device) |
| 339 | |
| 340 | # Find indices where the patch is True |
| 341 | patch_coordinates = torch.nonzero(patch) # shape [num_points, 2] |
| 342 | |
| 343 | # Extract features in the patch |
| 344 | F_qi = F[:, :, patch_coordinates[:, 0], patch_coordinates[:, 1]] |
| 345 | # Extract feature of the initial handle point |
| 346 | f_i = F0[:, :, handle_points0[i][0].long(), handle_points0[i][1].long()] |
| 347 | |
| 348 | # Compute the L1 distance between the patch features and the initial handle point feature |
| 349 | distances = torch.norm(F_qi - f_i[:, :, None], p=1, dim=1) |
| 350 | |
| 351 | # Find the new handle point as the one with minimum distance |
| 352 | min_index = torch.argmin(distances) |
| 353 | new_handle_points[i] = patch_coordinates[min_index] |
| 354 | |
| 355 | return new_handle_points |