(
self,
tar_base: str,
batch_size: int,
num_workers: int = 8,
shuffle: int = 1000,
min_num_tracks: int = 128,
context_length: int = 128,
num_targets_per_context: int = 15,
allow_poke_query_overlap: bool = True,
skip: int = 25,
allow_invisible_track_ends: bool = True,
allow_out_of_frame_track_ends: bool = True,
# Static camera detection heuristic
static_camera_flow_mag_threshold: float = 0.01,
static_camera_fraction_threshold: float = 0.4,
# Other
verbose: bool = False,
)
| 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: |
nothing calls this directly
no outgoing calls
no test coverage detected