prefusion 2.4.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- prefusion-2.4.1/PKG-INFO +57 -0
- prefusion-2.4.1/README.md +20 -0
- prefusion-2.4.1/prefusion/__init__.py +6 -0
- prefusion-2.4.1/prefusion/dataset/__init__.py +12 -0
- prefusion-2.4.1/prefusion/dataset/dataset.py +334 -0
- prefusion-2.4.1/prefusion/dataset/index_info.py +163 -0
- prefusion-2.4.1/prefusion/dataset/model_feeder.py +43 -0
- prefusion-2.4.1/prefusion/dataset/sampler.py +40 -0
- prefusion-2.4.1/prefusion/dataset/schema.py +119 -0
- prefusion-2.4.1/prefusion/dataset/sequence_sampler.py +726 -0
- prefusion-2.4.1/prefusion/dataset/streaming_collate.py +66 -0
- prefusion-2.4.1/prefusion/dataset/streaming_dataset.py +254 -0
- prefusion-2.4.1/prefusion/dataset/streaming_sampler.py +240 -0
- prefusion-2.4.1/prefusion/dataset/subepoch_manager.py +279 -0
- prefusion-2.4.1/prefusion/dataset/tensor_smith.py +3107 -0
- prefusion-2.4.1/prefusion/dataset/transform.py +2170 -0
- prefusion-2.4.1/prefusion/dataset/transformable_loader.py +1055 -0
- prefusion-2.4.1/prefusion/dataset/utils.py +413 -0
- prefusion-2.4.1/prefusion/hooks/__init__.py +5 -0
- prefusion-2.4.1/prefusion/hooks/airflow.py +153 -0
- prefusion-2.4.1/prefusion/hooks/checkpoint.py +165 -0
- prefusion-2.4.1/prefusion/hooks/iter_timer_hook.py +57 -0
- prefusion-2.4.1/prefusion/hooks/log_processor.py +189 -0
- prefusion-2.4.1/prefusion/hooks/utils.py +10 -0
- prefusion-2.4.1/prefusion/loss/__init__.py +2 -0
- prefusion-2.4.1/prefusion/loss/basic.py +73 -0
- prefusion-2.4.1/prefusion/loss/planar.py +247 -0
- prefusion-2.4.1/prefusion/metric/average_precision.py +668 -0
- prefusion-2.4.1/prefusion/metric/wrongly_implemented_average_precision.py +214 -0
- prefusion-2.4.1/prefusion/models/__init__.py +1 -0
- prefusion-2.4.1/prefusion/models/model.py +76 -0
- prefusion-2.4.1/prefusion/registry.py +130 -0
- prefusion-2.4.1/prefusion/runner/__init__.py +2 -0
- prefusion-2.4.1/prefusion/runner/loops.py +256 -0
- prefusion-2.4.1/prefusion/runner/runner.py +18 -0
- prefusion-2.4.1/prefusion/utils/__init__.py +1 -0
- prefusion-2.4.1/prefusion/utils/camera_ego_coord3d.py +225 -0
- prefusion-2.4.1/prefusion/utils/config_confirmation.py +115 -0
- prefusion-2.4.1/prefusion/utils/hpr_based_visibility.py +236 -0
- prefusion-2.4.1/prefusion/utils/visualization.py +145 -0
- prefusion-2.4.1/prefusion/utils/z_buffering_based_visibility.py +195 -0
- prefusion-2.4.1/prefusion.egg-info/PKG-INFO +57 -0
- prefusion-2.4.1/prefusion.egg-info/SOURCES.txt +47 -0
- prefusion-2.4.1/prefusion.egg-info/dependency_links.txt +1 -0
- prefusion-2.4.1/prefusion.egg-info/requires.txt +31 -0
- prefusion-2.4.1/prefusion.egg-info/top_level.txt +1 -0
- prefusion-2.4.1/pyproject.toml +64 -0
- prefusion-2.4.1/setup.cfg +4 -0
- prefusion-2.4.1/tests/test_import_boundary.py +25 -0
prefusion-2.4.1/PKG-INFO
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: prefusion
|
|
3
|
+
Version: 2.4.1
|
|
4
|
+
Summary: MMEngine-based core framework for Prefusion projects
|
|
5
|
+
Requires-Python: >=3.10
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
Requires-Dist: cachetools
|
|
8
|
+
Requires-Dist: copious>=0.1.23
|
|
9
|
+
Requires-Dist: easydict==1.13
|
|
10
|
+
Requires-Dist: loguru
|
|
11
|
+
Requires-Dist: matplotlib
|
|
12
|
+
Requires-Dist: mmengine<0.11,>=0.10.5
|
|
13
|
+
Requires-Dist: mmdet<3.3,>=3.2.0
|
|
14
|
+
Requires-Dist: mtv4d
|
|
15
|
+
Requires-Dist: numba==0.60.0
|
|
16
|
+
Requires-Dist: numpy==1.26.4
|
|
17
|
+
Requires-Dist: nuscenes-devkit
|
|
18
|
+
Requires-Dist: opencv-python-headless
|
|
19
|
+
Requires-Dist: open3d==0.18.0
|
|
20
|
+
Requires-Dist: pandas
|
|
21
|
+
Requires-Dist: polars==1.28.1
|
|
22
|
+
Requires-Dist: pydantic>=2
|
|
23
|
+
Requires-Dist: pypcd-imp==0.1.5
|
|
24
|
+
Requires-Dist: redis==5.2.0
|
|
25
|
+
Requires-Dist: requests
|
|
26
|
+
Requires-Dist: scipy==1.14.1
|
|
27
|
+
Requires-Dist: shapely
|
|
28
|
+
Requires-Dist: tabulate
|
|
29
|
+
Requires-Dist: tensorboard==2.17.1
|
|
30
|
+
Requires-Dist: torch==2.4.1
|
|
31
|
+
Requires-Dist: torchvision==0.19.1
|
|
32
|
+
Requires-Dist: tqdm
|
|
33
|
+
Requires-Dist: virtual-camera>=0.0.4.3
|
|
34
|
+
Provides-Extra: test
|
|
35
|
+
Requires-Dist: pytest==8.3.3; extra == "test"
|
|
36
|
+
Requires-Dist: pytest-cov; extra == "test"
|
|
37
|
+
|
|
38
|
+
# Prefusion
|
|
39
|
+
|
|
40
|
+
`prefusion` is the MMEngine-based core framework shared by Prefusion task projects.
|
|
41
|
+
|
|
42
|
+
Install the framework with:
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
pip install prefusion
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
Task-specific models and configurations live in the separately installable
|
|
49
|
+
[`prefusion-contrib`](../prefusion-contrib) project. The public Python API is
|
|
50
|
+
kept stable as `prefusion.*`.
|
|
51
|
+
|
|
52
|
+
## Development validation
|
|
53
|
+
|
|
54
|
+
```bash
|
|
55
|
+
python -m pip wheel --no-deps --wheel-dir dist .
|
|
56
|
+
python -m pytest tests
|
|
57
|
+
```
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
# Prefusion
|
|
2
|
+
|
|
3
|
+
`prefusion` is the MMEngine-based core framework shared by Prefusion task projects.
|
|
4
|
+
|
|
5
|
+
Install the framework with:
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install prefusion
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
Task-specific models and configurations live in the separately installable
|
|
12
|
+
[`prefusion-contrib`](../prefusion-contrib) project. The public Python API is
|
|
13
|
+
kept stable as `prefusion.*`.
|
|
14
|
+
|
|
15
|
+
## Development validation
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
python -m pip wheel --no-deps --wheel-dir dist .
|
|
19
|
+
python -m pytest tests
|
|
20
|
+
```
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
from .dataset import *
|
|
2
|
+
from .index_info import *
|
|
3
|
+
from .model_feeder import *
|
|
4
|
+
from .sampler import *
|
|
5
|
+
from .sequence_sampler import *
|
|
6
|
+
from .streaming_collate import *
|
|
7
|
+
from .streaming_dataset import *
|
|
8
|
+
from .streaming_sampler import *
|
|
9
|
+
from .subepoch_manager import *
|
|
10
|
+
from .tensor_smith import *
|
|
11
|
+
from .transform import *
|
|
12
|
+
from .transformable_loader import *
|
|
@@ -0,0 +1,334 @@
|
|
|
1
|
+
# Written by AlphaLFC. All rights reserved.
|
|
2
|
+
import copy
|
|
3
|
+
import functools
|
|
4
|
+
import logging
|
|
5
|
+
import os
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from typing import TYPE_CHECKING, Optional, Union
|
|
8
|
+
|
|
9
|
+
from cachetools import Cache, cached
|
|
10
|
+
from mmengine.logging import print_log
|
|
11
|
+
from torch.utils.data import Dataset
|
|
12
|
+
|
|
13
|
+
from prefusion.dataset.subepoch_manager import EndOfAllSubEpochs
|
|
14
|
+
from prefusion.dataset.transform import (
|
|
15
|
+
Bbox3D,
|
|
16
|
+
CameraDepthSet,
|
|
17
|
+
CameraImageSet,
|
|
18
|
+
CameraSegMaskSet,
|
|
19
|
+
EgoPoseSet,
|
|
20
|
+
LidarPoints,
|
|
21
|
+
OccSdf3D,
|
|
22
|
+
OccSdfBev,
|
|
23
|
+
ParkingSlot3D,
|
|
24
|
+
Polygon3D,
|
|
25
|
+
Polyline3D,
|
|
26
|
+
SegBev,
|
|
27
|
+
)
|
|
28
|
+
from prefusion.dataset.transformable_loader import (
|
|
29
|
+
Bbox3DLoader,
|
|
30
|
+
CameraDepthSetLoader,
|
|
31
|
+
CameraImageSetLoader,
|
|
32
|
+
CameraSegMaskSetLoader,
|
|
33
|
+
EgoPoseSetLoader,
|
|
34
|
+
LidarPointsLoader,
|
|
35
|
+
OccSdf3DLoader,
|
|
36
|
+
OccSdfBevLoader,
|
|
37
|
+
ParkingSlot3DLoader,
|
|
38
|
+
Polygon3DLoader,
|
|
39
|
+
Polyline3DLoader,
|
|
40
|
+
SegBevLoader,
|
|
41
|
+
TransformableLoader,
|
|
42
|
+
)
|
|
43
|
+
from prefusion.registry import DATASETS, MMENGINE_FUNCTIONS, TENSOR_SMITHS
|
|
44
|
+
|
|
45
|
+
from .utils import (
|
|
46
|
+
build_model_feeder,
|
|
47
|
+
build_sequence_sampler,
|
|
48
|
+
build_subepoch_manager,
|
|
49
|
+
build_tensor_smith,
|
|
50
|
+
build_transformable_loader,
|
|
51
|
+
build_transforms,
|
|
52
|
+
divide,
|
|
53
|
+
load_frame_info,
|
|
54
|
+
read_frame_pickle,
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
if TYPE_CHECKING:
|
|
58
|
+
from .index_info import IndexInfo
|
|
59
|
+
from .model_feeder import BaseModelFeeder
|
|
60
|
+
from .sequence_sampler import (
|
|
61
|
+
SequenceSampler,
|
|
62
|
+
)
|
|
63
|
+
from .subepoch_manager import SubEpochManager
|
|
64
|
+
from .tensor_smith import TensorSmith
|
|
65
|
+
from .transform import Transform, Transformable
|
|
66
|
+
|
|
67
|
+
__all__ = ["SequenceBatchDataset"]
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
@MMENGINE_FUNCTIONS.register_module()
|
|
71
|
+
def collate_dict(batch):
|
|
72
|
+
return batch[0]
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
SequenceBatch = list[list[dict]]
|
|
76
|
+
|
|
77
|
+
_print_log = functools.partial(print_log, logger="current")
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@DATASETS.register_module()
|
|
81
|
+
class SequenceBatchDataset(Dataset):
|
|
82
|
+
"""
|
|
83
|
+
A novel dataset class for batching sequences for multi-module data.
|
|
84
|
+
"""
|
|
85
|
+
|
|
86
|
+
# TODO: implement visualization?
|
|
87
|
+
DEFAULT_LOADERS = {
|
|
88
|
+
CameraImageSet.__name__: CameraImageSetLoader,
|
|
89
|
+
CameraDepthSet.__name__: CameraDepthSetLoader,
|
|
90
|
+
CameraSegMaskSet.__name__: CameraSegMaskSetLoader,
|
|
91
|
+
LidarPoints.__name__: LidarPointsLoader,
|
|
92
|
+
EgoPoseSet.__name__: EgoPoseSetLoader,
|
|
93
|
+
Bbox3D.__name__: Bbox3DLoader,
|
|
94
|
+
Polyline3D.__name__: Polyline3DLoader,
|
|
95
|
+
Polygon3D.__name__: Polygon3DLoader,
|
|
96
|
+
ParkingSlot3D.__name__: ParkingSlot3DLoader,
|
|
97
|
+
OccSdfBev.__name__: OccSdfBevLoader,
|
|
98
|
+
OccSdf3D.__name__: OccSdf3DLoader,
|
|
99
|
+
SegBev.__name__: SegBevLoader,
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
def __init__(
|
|
103
|
+
self,
|
|
104
|
+
name,
|
|
105
|
+
data_root: str | Path,
|
|
106
|
+
info_path: str | Path,
|
|
107
|
+
transformables: dict,
|
|
108
|
+
transforms: list[Union[dict, "Transform"]] = None,
|
|
109
|
+
model_feeder: Union["BaseModelFeeder", dict] = None,
|
|
110
|
+
subepoch_manager: Union["SubEpochManager", dict] = None,
|
|
111
|
+
sequence_sampler: Union["SequenceSampler", dict] = None,
|
|
112
|
+
batch_size: int = 1,
|
|
113
|
+
drop_last: bool = False,
|
|
114
|
+
sequence_backtime_prob: float = 0.0,
|
|
115
|
+
seed_dataset: int = None,
|
|
116
|
+
):
|
|
117
|
+
"""
|
|
118
|
+
Initializes the dataset instance.
|
|
119
|
+
|
|
120
|
+
Args:
|
|
121
|
+
- name (str): Name of the dataset.
|
|
122
|
+
- data_root (str): Root directory of the dataset.
|
|
123
|
+
- info_path (str): Path to the information file.
|
|
124
|
+
- transformables (dict): Dict of transformable definitions.
|
|
125
|
+
- transforms (list): Transform classes for preprocessing transformables. Build by TRANSFORMS.
|
|
126
|
+
- batch_size (int, optional): Batch size; defualt is 1.
|
|
127
|
+
- sequence_backtime_prob (float): Probability of sequencing backtime frames.
|
|
128
|
+
- subepoch_manager (SubEpochManager or dict): Only works for training. If set, only `num_sequence_batches_per_subepoch` will be trained in a epoch and the
|
|
129
|
+
rest of data will be spit out in the next epoch.
|
|
130
|
+
- seed_dataset (int): Random seed for dataset
|
|
131
|
+
|
|
132
|
+
Notes:
|
|
133
|
+
- `scene`: a scene contains a sequence of frames
|
|
134
|
+
- `sequence`: a subset of a scene, when testing, a scene can be one sequence
|
|
135
|
+
- `frame`: a time instance of a scene
|
|
136
|
+
- one index should get a batch of sequences, one sequence should be a sequential list
|
|
137
|
+
- the length of the dataset should be len(self.sequences) // self.batch_size
|
|
138
|
+
- a frame data is a dictionary containing the transformed data.
|
|
139
|
+
"""
|
|
140
|
+
super().__init__()
|
|
141
|
+
self.name = name
|
|
142
|
+
self.data_root = Path(data_root)
|
|
143
|
+
self.frame_info = load_frame_info(info_path)
|
|
144
|
+
self.transformables = transformables
|
|
145
|
+
self.transforms = build_transforms(transforms)
|
|
146
|
+
self.model_feeder = build_model_feeder(model_feeder)
|
|
147
|
+
self.sequence_sampler = build_sequence_sampler(sequence_sampler)
|
|
148
|
+
self.subepoch_manager = None
|
|
149
|
+
self.batch_size = batch_size
|
|
150
|
+
self.drop_last = drop_last
|
|
151
|
+
self.sequence_backtime_prob = sequence_backtime_prob
|
|
152
|
+
self.seed_dataset = seed_dataset
|
|
153
|
+
|
|
154
|
+
self._sample_sequences()
|
|
155
|
+
|
|
156
|
+
if type(self.sequence_sampler).__name__ in [
|
|
157
|
+
"ClassBalancedSequenceSampler",
|
|
158
|
+
"SceneLevelBalancedSequenceSampler",
|
|
159
|
+
"TrainIndexSequenceSampler",
|
|
160
|
+
]:
|
|
161
|
+
self.subepoch_manager = build_subepoch_manager(subepoch_manager, batch_size)
|
|
162
|
+
if self.subepoch_manager is not None:
|
|
163
|
+
_print_log(
|
|
164
|
+
"Since you have set SubEpochManager, "
|
|
165
|
+
"you should consider set a larger max_epochs for your training.",
|
|
166
|
+
level=logging.WARNING,
|
|
167
|
+
)
|
|
168
|
+
self.subepoch_manager.init(len(self.sequences))
|
|
169
|
+
|
|
170
|
+
def _sample_sequences(self):
|
|
171
|
+
self.sequences = self.sequence_sampler.sample(self.data_root, self.frame_info)
|
|
172
|
+
self.num_total_sequence_batches = divide(len(self.sequences), self.batch_size, drop_last=self.drop_last)
|
|
173
|
+
|
|
174
|
+
@property
|
|
175
|
+
def sequence_length(self):
|
|
176
|
+
return self.sequence_sampler.sequence_length
|
|
177
|
+
|
|
178
|
+
def __repr__(self):
|
|
179
|
+
return "".join(
|
|
180
|
+
[
|
|
181
|
+
f"An instance of {self.__class__}: (\n",
|
|
182
|
+
f" name={self.name}, \n",
|
|
183
|
+
f" num_sequences={len(self.sequences)}, \n"
|
|
184
|
+
f" batch_size={self.batch_size}, \n"
|
|
185
|
+
f" sequence_length={self.sequence_sampler.possible_sequence_lengths}, \n"
|
|
186
|
+
f" frame_interval={self.sequence_sampler.possible_frame_intervals}, \n"
|
|
187
|
+
f" sequence_backtime_prob={self.sequence_backtime_prob}\n)",
|
|
188
|
+
]
|
|
189
|
+
)
|
|
190
|
+
|
|
191
|
+
def load_all_transformables(self, index_info: "IndexInfo") -> dict:
|
|
192
|
+
transformables = {}
|
|
193
|
+
frame_data = read_frame_pickle(self.data_root / self.frame_info[index_info.scene_frame_id])
|
|
194
|
+
for name in self.transformables:
|
|
195
|
+
_t_cfg = self.transformables[name]
|
|
196
|
+
loader_cfg = _t_cfg["loader"] if "loader" in _t_cfg else None
|
|
197
|
+
if not loader_cfg:
|
|
198
|
+
loader_cfg = self._get_default_loader_cfg(_t_cfg["type"])
|
|
199
|
+
loader_cfg.update(data_root=self.data_root)
|
|
200
|
+
tensor_smith_cfg = _t_cfg["tensor_smith"] if "tensor_smith" in _t_cfg else None
|
|
201
|
+
loader = self._build_transformable_loader(loader_cfg)
|
|
202
|
+
tensor_smith = build_tensor_smith(tensor_smith_cfg) if "tensor_smith" in _t_cfg else None
|
|
203
|
+
rest_kwargs = {k: v for k, v in _t_cfg.items() if k not in ["type", "loader", "tensor_smith"]}
|
|
204
|
+
transformables[name] = self._build_transformable(
|
|
205
|
+
name,
|
|
206
|
+
frame_data,
|
|
207
|
+
index_info,
|
|
208
|
+
loader,
|
|
209
|
+
tensor_smith=tensor_smith,
|
|
210
|
+
**rest_kwargs,
|
|
211
|
+
)
|
|
212
|
+
|
|
213
|
+
return transformables
|
|
214
|
+
|
|
215
|
+
@cached(
|
|
216
|
+
cache=Cache(maxsize=float("inf")),
|
|
217
|
+
key=lambda self_, cfg: str(sorted((cfg or {}).items())),
|
|
218
|
+
)
|
|
219
|
+
def _build_transformable_loader(self, loader_cfg: dict) -> TransformableLoader:
|
|
220
|
+
loader = build_transformable_loader(loader_cfg)
|
|
221
|
+
return loader
|
|
222
|
+
|
|
223
|
+
def _get_default_loader_cfg(self, transformable_type: str) -> dict:
|
|
224
|
+
if transformable_type not in self.DEFAULT_LOADERS:
|
|
225
|
+
raise ValueError(
|
|
226
|
+
f"No default transformable loader for transformable type: {transformable_type}. Please provide one explicitly."
|
|
227
|
+
)
|
|
228
|
+
return {
|
|
229
|
+
"type": self.DEFAULT_LOADERS[transformable_type],
|
|
230
|
+
"data_root": self.data_root,
|
|
231
|
+
}
|
|
232
|
+
|
|
233
|
+
@staticmethod
|
|
234
|
+
def _build_tensor_smith(tensor_smith: dict = None):
|
|
235
|
+
tensor_smith = copy.deepcopy(tensor_smith)
|
|
236
|
+
if isinstance(tensor_smith, dict):
|
|
237
|
+
tensor_smith = TENSOR_SMITHS.build(tensor_smith)
|
|
238
|
+
return tensor_smith
|
|
239
|
+
|
|
240
|
+
def __len__(self):
|
|
241
|
+
if self.subepoch_manager is None:
|
|
242
|
+
return self.num_total_sequence_batches
|
|
243
|
+
return self.subepoch_manager.num_sequence_batches_per_subepoch
|
|
244
|
+
|
|
245
|
+
@staticmethod
|
|
246
|
+
def _batch_sequences(sequence_batch_ind, sequences, batch_size):
|
|
247
|
+
batched_sequences = []
|
|
248
|
+
for batch_idx in range(batch_size):
|
|
249
|
+
sequence_idx = sequence_batch_ind * batch_size + batch_idx
|
|
250
|
+
if sequence_idx >= len(sequences):
|
|
251
|
+
sequence_idx = max(0, 2 * (len(sequences) - 1) - sequence_idx)
|
|
252
|
+
batched_sequences.append(sequences[sequence_idx])
|
|
253
|
+
return batched_sequences
|
|
254
|
+
|
|
255
|
+
def post_epoch_processing(self):
|
|
256
|
+
if self.subepoch_manager is None: # training on the normal epoch settings
|
|
257
|
+
self._sample_sequences()
|
|
258
|
+
return
|
|
259
|
+
|
|
260
|
+
try: # training on sub-epochs
|
|
261
|
+
self.subepoch_manager.to_next_sub_epoch()
|
|
262
|
+
except EndOfAllSubEpochs:
|
|
263
|
+
# from mmengine.dist import get_dist_info
|
|
264
|
+
# from loguru import logger
|
|
265
|
+
# rank, _ = get_dist_info()
|
|
266
|
+
# logger.debug(f"[rank{rank}] id={id(self.subepoch_manager)}")
|
|
267
|
+
# logger.debug(f"[rank{rank}] drop_last_sequence_batch={self.subepoch_manager.drop_last_sequence_batch}")
|
|
268
|
+
# logger.debug(f"[rank{rank}] drop_last_subepoch={self.subepoch_manager.drop_last_subepoch}")
|
|
269
|
+
# logger.debug(f"[rank{rank}] num_total_sequences={self.subepoch_manager.num_total_sequences}")
|
|
270
|
+
# logger.debug(f"[rank{rank}] num_total_sequence_batches={self.subepoch_manager.num_total_sequence_batches}")
|
|
271
|
+
# logger.debug(f"[rank{rank}] num_sequence_batches_per_subepoch={self.subepoch_manager.num_sequence_batches_per_subepoch}")
|
|
272
|
+
# logger.debug(f"[rank{rank}] batch_size={self.subepoch_manager.batch_size}")
|
|
273
|
+
# logger.debug(f"[rank{rank}] num_subepochs={self.subepoch_manager.num_subepochs}")
|
|
274
|
+
# logger.debug(f"[rank{rank}] id(visited)={id(self.subepoch_manager.visited)}")
|
|
275
|
+
# logger.debug(f"[rank{rank}] visited={self.subepoch_manager.visited.todict().keys()}")
|
|
276
|
+
|
|
277
|
+
self._sample_sequences()
|
|
278
|
+
self.subepoch_manager.reset(len(self.sequences))
|
|
279
|
+
|
|
280
|
+
def __getitem__(self, idx) -> SequenceBatch:
|
|
281
|
+
if self.subepoch_manager is not None:
|
|
282
|
+
idx = self.subepoch_manager.translate_index(idx)
|
|
283
|
+
|
|
284
|
+
batched_sequences = self._batch_sequences(idx, self.sequences, self.batch_size)
|
|
285
|
+
|
|
286
|
+
batch = []
|
|
287
|
+
sequence_batch_seed = int.from_bytes(os.urandom(2), byteorder="big")
|
|
288
|
+
for sequence in batched_sequences:
|
|
289
|
+
sequence_of_inputs = []
|
|
290
|
+
sequence_seed = int.from_bytes(os.urandom(2), byteorder="big")
|
|
291
|
+
for i, index_info in enumerate(sequence):
|
|
292
|
+
frame_seed = int.from_bytes(os.urandom(2), byteorder="big")
|
|
293
|
+
transformables = self.load_all_transformables(index_info)
|
|
294
|
+
for transform in self.transforms:
|
|
295
|
+
transform(
|
|
296
|
+
*transformables.values(),
|
|
297
|
+
seeds={
|
|
298
|
+
"sequence": sequence_seed,
|
|
299
|
+
"sequence_batch": sequence_batch_seed,
|
|
300
|
+
"frame": frame_seed,
|
|
301
|
+
},
|
|
302
|
+
)
|
|
303
|
+
input_dict = {
|
|
304
|
+
"index_info": index_info,
|
|
305
|
+
"transformables": transformables,
|
|
306
|
+
}
|
|
307
|
+
sequence_of_inputs.append(input_dict)
|
|
308
|
+
batch.append(sequence_of_inputs)
|
|
309
|
+
|
|
310
|
+
sequence_batch = []
|
|
311
|
+
for frame_batch in zip(*batch):
|
|
312
|
+
sequence_batch.append(frame_batch)
|
|
313
|
+
|
|
314
|
+
model_food = self.model_feeder(sequence_batch)
|
|
315
|
+
|
|
316
|
+
return model_food
|
|
317
|
+
|
|
318
|
+
def _build_transformable(
|
|
319
|
+
self,
|
|
320
|
+
name: str,
|
|
321
|
+
frame_data: dict[str, dict],
|
|
322
|
+
index_info: "IndexInfo",
|
|
323
|
+
loader: "TransformableLoader",
|
|
324
|
+
tensor_smith: Optional["TensorSmith"] = None,
|
|
325
|
+
**kwargs,
|
|
326
|
+
) -> "Transformable":
|
|
327
|
+
return loader.load(
|
|
328
|
+
name,
|
|
329
|
+
self.frame_info,
|
|
330
|
+
frame_data,
|
|
331
|
+
index_info,
|
|
332
|
+
tensor_smith=tensor_smith,
|
|
333
|
+
**kwargs,
|
|
334
|
+
)
|
|
@@ -0,0 +1,163 @@
|
|
|
1
|
+
import copy
|
|
2
|
+
from typing import Optional, Union
|
|
3
|
+
|
|
4
|
+
__all__ = ["IndexInfo"]
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class IndexInfo:
|
|
8
|
+
def __init__(
|
|
9
|
+
self,
|
|
10
|
+
scene_id: str,
|
|
11
|
+
frame_id: str,
|
|
12
|
+
s_prev: Optional["IndexInfo"] = None,
|
|
13
|
+
s_next: Optional["IndexInfo"] = None,
|
|
14
|
+
q_prev: Optional["IndexInfo"] = None,
|
|
15
|
+
q_next: Optional["IndexInfo"] = None,
|
|
16
|
+
):
|
|
17
|
+
"""Object that stores scene_id and frame_id as well as the adjacent linking relationships.
|
|
18
|
+
|
|
19
|
+
Parameters
|
|
20
|
+
----------
|
|
21
|
+
scene_id : str
|
|
22
|
+
scene id
|
|
23
|
+
frame_id : str
|
|
24
|
+
frame id
|
|
25
|
+
s_prev : IndexInfo, optional
|
|
26
|
+
previous index_info within the same scene, by default None
|
|
27
|
+
s_next : IndexInfo, optional
|
|
28
|
+
next index_info within the same scene, by default None
|
|
29
|
+
q_prev : IndexInfo, optional
|
|
30
|
+
previous index_info within the same sequence, by default None
|
|
31
|
+
q_next : IndexInfo, optional
|
|
32
|
+
next index_info within the same sequence, by default None
|
|
33
|
+
"""
|
|
34
|
+
self.scene_id = scene_id
|
|
35
|
+
self.frame_id = frame_id
|
|
36
|
+
self.s_prev = s_prev
|
|
37
|
+
self.s_next = s_next
|
|
38
|
+
self.q_prev = q_prev
|
|
39
|
+
self.q_next = q_next
|
|
40
|
+
if s_prev:
|
|
41
|
+
s_prev.s_next = self
|
|
42
|
+
if s_next:
|
|
43
|
+
s_next.s_prev = self
|
|
44
|
+
if q_prev:
|
|
45
|
+
q_prev.q_next = self
|
|
46
|
+
if q_next:
|
|
47
|
+
q_next.q_prev = self
|
|
48
|
+
|
|
49
|
+
@property
|
|
50
|
+
def prev(self) -> Union["IndexInfo", None]:
|
|
51
|
+
return self.q_prev or self.s_prev
|
|
52
|
+
|
|
53
|
+
@property
|
|
54
|
+
def next(self) -> Union["IndexInfo", None]:
|
|
55
|
+
return self.q_next or self.s_next
|
|
56
|
+
|
|
57
|
+
def __repr__(self) -> str:
|
|
58
|
+
return f"{self.scene_id}/{self.frame_id} (prev: {self.prev.scene_frame_id if self.prev else None}, next: {self.next.scene_frame_id if self.next else None})"
|
|
59
|
+
|
|
60
|
+
def __eq__(self, other: object) -> bool:
|
|
61
|
+
if not isinstance(other, IndexInfo):
|
|
62
|
+
return False
|
|
63
|
+
if not self.two_index_info_are_equal(self.s_prev, other.s_prev):
|
|
64
|
+
return False
|
|
65
|
+
if not self.two_index_info_are_equal(self.s_next, other.s_next):
|
|
66
|
+
return False
|
|
67
|
+
if not self.two_index_info_are_equal(self.q_prev, other.q_prev):
|
|
68
|
+
return False
|
|
69
|
+
if not self.two_index_info_are_equal(self.q_next, other.q_next):
|
|
70
|
+
return False
|
|
71
|
+
if self.scene_frame_id != other.scene_frame_id:
|
|
72
|
+
return False
|
|
73
|
+
|
|
74
|
+
return True
|
|
75
|
+
|
|
76
|
+
@staticmethod
|
|
77
|
+
def two_index_info_are_equal(a_index_info, another_index_info):
|
|
78
|
+
if (a_index_info is None) ^ (another_index_info is None):
|
|
79
|
+
return False
|
|
80
|
+
if (
|
|
81
|
+
a_index_info is not None
|
|
82
|
+
and another_index_info is not None
|
|
83
|
+
and a_index_info.scene_frame_id != another_index_info.scene_frame_id
|
|
84
|
+
):
|
|
85
|
+
return False
|
|
86
|
+
return True
|
|
87
|
+
|
|
88
|
+
@property
|
|
89
|
+
def scene_frame_id(self) -> str:
|
|
90
|
+
return f"{self.scene_id}/{self.frame_id}"
|
|
91
|
+
|
|
92
|
+
def as_dict(self) -> dict:
|
|
93
|
+
return {
|
|
94
|
+
"scene_id": self.scene_id,
|
|
95
|
+
"frame_id": self.frame_id,
|
|
96
|
+
"s_prev": {"scene_id": self.s_prev.scene_id, "frame_id": self.s_prev.frame_id} if self.s_prev else None,
|
|
97
|
+
"s_next": {"scene_id": self.s_next.scene_id, "frame_id": self.s_next.frame_id} if self.s_next else None,
|
|
98
|
+
"q_prev": {"scene_id": self.q_prev.scene_id, "frame_id": self.q_prev.frame_id} if self.q_prev else None,
|
|
99
|
+
"q_next": {"scene_id": self.q_next.scene_id, "frame_id": self.q_next.frame_id} if self.q_next else None,
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
@classmethod
|
|
103
|
+
def from_str(
|
|
104
|
+
cls,
|
|
105
|
+
index_str: str,
|
|
106
|
+
s_prev: Optional["IndexInfo"] = None,
|
|
107
|
+
s_next: Optional["IndexInfo"] = None,
|
|
108
|
+
q_prev: Optional["IndexInfo"] = None,
|
|
109
|
+
q_next: Optional["IndexInfo"] = None,
|
|
110
|
+
sep: str = "/",
|
|
111
|
+
):
|
|
112
|
+
"""Create IndexInfo instance from index string.
|
|
113
|
+
|
|
114
|
+
Parameters
|
|
115
|
+
----------
|
|
116
|
+
index_str : str
|
|
117
|
+
expected format: `scene_id/frame_id`
|
|
118
|
+
s_prev : IndexInfo, optional
|
|
119
|
+
previous index_info within the same scene, by default None
|
|
120
|
+
s_next : IndexInfo, optional
|
|
121
|
+
next index_info within the same scene, by default None
|
|
122
|
+
q_prev : IndexInfo, optional
|
|
123
|
+
previous index_info within the same sequence, by default None
|
|
124
|
+
q_next : IndexInfo, optional
|
|
125
|
+
next index_info within the same sequence, by default None
|
|
126
|
+
sep : str, optional
|
|
127
|
+
separator, by default "/"
|
|
128
|
+
"""
|
|
129
|
+
scene_id, frame_id = index_str.split(sep)
|
|
130
|
+
return IndexInfo(scene_id, frame_id, s_prev=s_prev, s_next=s_next, q_prev=q_prev, q_next=q_next)
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def establish_linkings(index_info_list: list["IndexInfo"], sequence_linking=False) -> list["IndexInfo"]:
|
|
134
|
+
"""Establish doubly-linked list relationships between IndexInfo objects.
|
|
135
|
+
|
|
136
|
+
Callers (e.g., sequence samplers using generate_sequences) may pass the same
|
|
137
|
+
IndexInfo object multiple times in a single list—for instance, when padding
|
|
138
|
+
short scenes to a fixed sequence length repeats boundary frame indices. To
|
|
139
|
+
prevent corrupting existing links with self-loops, duplicates are cloned via
|
|
140
|
+
copy.copy(), preserving scene-level links (s_prev/s_next) while giving each
|
|
141
|
+
position in the sequence its own independent sequence-level links.
|
|
142
|
+
"""
|
|
143
|
+
attr_prev = "q_prev" if sequence_linking else "s_prev"
|
|
144
|
+
attr_next = "q_next" if sequence_linking else "s_next"
|
|
145
|
+
prev = None
|
|
146
|
+
seen_ids = set()
|
|
147
|
+
for i in range(len(index_info_list)):
|
|
148
|
+
cur = index_info_list[i]
|
|
149
|
+
if id(cur) in seen_ids:
|
|
150
|
+
cur = copy.copy(cur)
|
|
151
|
+
setattr(cur, attr_prev, None)
|
|
152
|
+
setattr(cur, attr_next, None)
|
|
153
|
+
index_info_list[i] = cur
|
|
154
|
+
seen_ids.add(id(cur))
|
|
155
|
+
setattr(cur, attr_prev, prev)
|
|
156
|
+
if prev is not None:
|
|
157
|
+
setattr(prev, attr_next, cur)
|
|
158
|
+
prev = cur
|
|
159
|
+
return index_info_list
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
def establish_sequence_linkings(index_info_list: list["IndexInfo"]) -> list["IndexInfo"]:
|
|
163
|
+
return establish_linkings(index_info_list, sequence_linking=True)
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
from prefusion.registry import MODEL_FEEDERS
|
|
2
|
+
|
|
3
|
+
__all__ = ["BaseModelFeeder"]
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@MODEL_FEEDERS.register_module()
|
|
7
|
+
class BaseModelFeeder:
|
|
8
|
+
"""BaseModelFeeder.
|
|
9
|
+
|
|
10
|
+
An alternative implementation of data_preprocessor.
|
|
11
|
+
|
|
12
|
+
Args
|
|
13
|
+
----
|
|
14
|
+
Any: Any parameter or keyword arguments.
|
|
15
|
+
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
def __init__(self, *args, **kwargs) -> None:
|
|
19
|
+
pass
|
|
20
|
+
|
|
21
|
+
def process(self, frame_batch: list) -> dict | list:
|
|
22
|
+
"""
|
|
23
|
+
Process frame_batch, make it ready for model inputs
|
|
24
|
+
|
|
25
|
+
Parameters
|
|
26
|
+
----------
|
|
27
|
+
frame_batch : list
|
|
28
|
+
list of input_dicts
|
|
29
|
+
|
|
30
|
+
Returns
|
|
31
|
+
-------
|
|
32
|
+
processed_frame_batch: dict | list
|
|
33
|
+
}
|
|
34
|
+
"""
|
|
35
|
+
processed_frame_batch = frame_batch
|
|
36
|
+
return processed_frame_batch
|
|
37
|
+
|
|
38
|
+
def __call__(self, sequence_batch: list):
|
|
39
|
+
processed_sequence_batch = []
|
|
40
|
+
for frame_batch in sequence_batch:
|
|
41
|
+
processed_sequence_batch.append(self.process(frame_batch))
|
|
42
|
+
|
|
43
|
+
return processed_sequence_batch
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
# Copyright (c) OpenMMLab. All rights reserved.
|
|
2
|
+
import torch
|
|
3
|
+
from mmdet.models.task_modules.samplers import BaseSampler, SamplingResult
|
|
4
|
+
|
|
5
|
+
from prefusion.registry import DATA_SAMPLERS
|
|
6
|
+
|
|
7
|
+
__all__ = ["PseudoSampler"]
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@DATA_SAMPLERS.register_module()
|
|
11
|
+
class PseudoSampler(BaseSampler):
|
|
12
|
+
"""A pseudo sampler that does not do sampling actually."""
|
|
13
|
+
|
|
14
|
+
def __init__(self, **kwargs):
|
|
15
|
+
pass
|
|
16
|
+
|
|
17
|
+
def _sample_pos(self, **kwargs):
|
|
18
|
+
"""Sample positive samples."""
|
|
19
|
+
raise NotImplementedError
|
|
20
|
+
|
|
21
|
+
def _sample_neg(self, **kwargs):
|
|
22
|
+
"""Sample negative samples."""
|
|
23
|
+
raise NotImplementedError
|
|
24
|
+
|
|
25
|
+
def sample(self, assign_result, bboxes, gt_bboxes, *args, **kwargs):
|
|
26
|
+
"""Directly returns the positive and negative indices of samples.
|
|
27
|
+
|
|
28
|
+
Args:
|
|
29
|
+
assign_result (:obj:`AssignResult`): Assigned results
|
|
30
|
+
bboxes (torch.Tensor): Bounding boxes
|
|
31
|
+
gt_bboxes (torch.Tensor): Ground truth boxes
|
|
32
|
+
|
|
33
|
+
Returns:
|
|
34
|
+
:obj:`SamplingResult`: sampler results
|
|
35
|
+
"""
|
|
36
|
+
pos_inds = torch.nonzero(assign_result.gt_inds > 0, as_tuple=False).squeeze(-1).unique()
|
|
37
|
+
neg_inds = torch.nonzero(assign_result.gt_inds == 0, as_tuple=False).squeeze(-1).unique()
|
|
38
|
+
gt_flags = bboxes.new_zeros(bboxes.shape[0], dtype=torch.uint8)
|
|
39
|
+
sampling_result = SamplingResult(pos_inds, neg_inds, bboxes, gt_bboxes, assign_result, gt_flags)
|
|
40
|
+
return sampling_result
|