| 66 | # --------------------------------------------------------------------------------------------------------------------- |
| 67 | |
| 68 | class TrackerShardsDataModule: |
| 69 | |
| 70 | def __init__( |
| 71 | self, |
| 72 | tar_base: str | list | ListConfig, |
| 73 | batch_size: int, |
| 74 | train=None, |
| 75 | validation=None, |
| 76 | num_workers=4, |
| 77 | val_batch_size: int | None = None, |
| 78 | val_num_workers: int | None = None, |
| 79 | prefetch_factor: int = 8, |
| 80 | num_tracks: int = 16, |
| 81 | num_steps: int = 16, |
| 82 | allow_invisible_track_ends: bool = True, |
| 83 | allow_out_of_frame_track_ends: bool = False, |
| 84 | filter_static_camera: bool = False, |
| 85 | static_camera_flow_mag_threshold: float=0.00035, |
| 86 | static_camera_fraction_threshold: float=0.2, |
| 87 | return_full_sequence: bool=False, |
| 88 | center_crop: bool=False, |
| 89 | crop_size: int | None = None, # provide an integer if we want to further center crop the image after resizing |
| 90 | track_key: str = "tracks_yx", |
| 91 | shuffle: bool = True, |
| 92 | certainty_threshold: float = 0.6, |
| 93 | visibility_threshold: float = 0.5, |
| 94 | t_scaling_factor: float=30.0, |
| 95 | ): |
| 96 | super().__init__() |
| 97 | self.tar_base = tar_base |
| 98 | self.batch_size = batch_size |
| 99 | self.num_workers = num_workers |
| 100 | self.train = train |
| 101 | self.validation = validation |
| 102 | self.val_batch_size = val_batch_size if val_batch_size is not None else batch_size |
| 103 | self.val_num_workers = val_num_workers if val_num_workers is not None else num_workers |
| 104 | self.prefetch_factor = prefetch_factor |
| 105 | |
| 106 | self.num_tracks = num_tracks |
| 107 | self.num_steps = num_steps |
| 108 | self.allow_invisible_track_ends = allow_invisible_track_ends |
| 109 | self.allow_out_of_frame_track_ends = allow_out_of_frame_track_ends |
| 110 | # self.flip_axes = flip_axes |
| 111 | self.filter_static_camera = filter_static_camera |
| 112 | |
| 113 | self.static_camera_flow_mag_threshold = static_camera_flow_mag_threshold |
| 114 | self.static_camera_fraction_threshold = static_camera_fraction_threshold |
| 115 | |
| 116 | self.return_full_sequence = return_full_sequence |
| 117 | self.center_crop = center_crop |
| 118 | self.crop_size = crop_size |
| 119 | |
| 120 | self.shuffle = shuffle |
| 121 | self.track_key = track_key |
| 122 | |
| 123 | self.certainty_threshold = certainty_threshold |
| 124 | self.visibility_threshold = visibility_threshold |
| 125 | |