| 54 | |
| 55 | |
| 56 | class TrackerShardsDataModule: |
| 57 | def __init__( |
| 58 | self, |
| 59 | tar_base: str, |
| 60 | batch_size: int, |
| 61 | num_workers: int = 8, |
| 62 | shuffle: int = 1000, |
| 63 | min_num_tracks: int = 128, |
| 64 | context_length: int = 128, |
| 65 | num_targets_per_context: int = 15, |
| 66 | allow_poke_query_overlap: bool = True, |
| 67 | skip: int = 25, |
| 68 | allow_invisible_track_ends: bool = True, |
| 69 | allow_out_of_frame_track_ends: bool = True, |
| 70 | # Static camera detection heuristic |
| 71 | static_camera_flow_mag_threshold: float = 0.01, |
| 72 | static_camera_fraction_threshold: float = 0.4, |
| 73 | # Other |
| 74 | verbose: bool = False, |
| 75 | ): |
| 76 | super().__init__() |
| 77 | self.tar_base = tar_base |
| 78 | self.batch_size = batch_size |
| 79 | self.num_workers = num_workers |
| 80 | self.shuffle = shuffle |
| 81 | |
| 82 | self.min_num_tracks = min_num_tracks |
| 83 | self.context_length = context_length |
| 84 | self.num_targets_per_context = num_targets_per_context |
| 85 | self.allow_poke_query_overlap = allow_poke_query_overlap |
| 86 | self.skip = skip |
| 87 | self.allow_invisible_track_ends = allow_invisible_track_ends |
| 88 | self.allow_out_of_frame_track_ends = allow_out_of_frame_track_ends |
| 89 | |
| 90 | self.static_camera_flow_mag_threshold = static_camera_flow_mag_threshold |
| 91 | self.static_camera_fraction_threshold = static_camera_fraction_threshold |
| 92 | |
| 93 | self.verbose = verbose |
| 94 | |
| 95 | def extract_training_sample(self, sample: dict[str, torch.Tensor]) -> dict[str, torch.Tensor | int | bool]: |
| 96 | try: |
| 97 | visibility: Bool[torch.Tensor, "t n_t"] = sample["visibility"] |
| 98 | tracks: Float[torch.Tensor, "t n_t 2"] = sample["tracks"] |
| 99 | track_in_frame: Bool[torch.Tensor, "t n_t"] = ( |
| 100 | (tracks[..., 0] >= 0) & (tracks[..., 0] <= 1) & (tracks[..., 1] >= 0) & (tracks[..., 1] <= 1) |
| 101 | ) |
| 102 | visible_and_in_frame: Bool[torch.Tensor, "t n_t"] = visibility & track_in_frame |
| 103 | valid_start_frames = ( |
| 104 | visible_and_in_frame.int().sum(dim=1) >= self.min_num_tracks |
| 105 | ) # Ones that have at least `num_tracks` visible tracks |
| 106 | if not valid_start_frames.any(): |
| 107 | return {"valid": False} |
| 108 | # skips: Float[torch.Tensor, "t_start t_end"] = sample["times"][None, :] - sample["times"][:, None] |
| 109 | # valid_skips: Bool[torch.Tensor, "t_start t_end"] = ( |
| 110 | # (skips >= self.time_skip_min) & (skips <= self.time_skip_max) |
| 111 | # ) & valid_start_frames[:, None] |
| 112 | # if not valid_skips.any(): |
| 113 | # return {"valid": False} |
nothing calls this directly
no outgoing calls
no test coverage detected