add: init import

Seeed-Studio · Dec 22, 2023 · 8eb4ec0 · 8eb4ec0
1 parent 2a487c0
commit 8eb4ec0
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/sscma/datasets/transforms/__init__.py b/sscma/datasets/transforms/__init__.py
@@ -1,4 +1,5 @@
 from .formatting import PackSensorInputs
 from .loading import LoadSensorFromFile
+from .wrappers import MutiBranchPipe
 
-__all__ = ['PackSensorInputs', 'LoadSensorFromFile']
+__all__ = ['PackSensorInputs', 'LoadSensorFromFile', 'MutiBranchPipe']
diff --git a/sscma/datasets/transforms/wrappers.py b/sscma/datasets/transforms/wrappers.py
@@ -17,7 +17,7 @@ def transform(self, results: Dict) -> Optional[Union[Dict, Tuple[List, List]]]:
             multi_results[branch] = {'inputs': None, 'data_samples': None}
         for branch, pipeline in self.branch_pipelines.items():
             branch_results = pipeline(copy.deepcopy(results))
-            if branch == 'unsup_teacher':
+            if branch == self.piece_key:
                 results['img'] = branch_results['inputs'].permute(1, 2, 0).cpu().numpy()
             # If one branch pipeline returns None,
             # it will sample another data from dataset.

diff --git a/tools/export.py b/tools/export.py
@@ -15,7 +15,7 @@
 import sscma.evaluation  # noqa
 import sscma.models  # noqa
 import sscma.visualization  # noqa
-from sscma.utils.check import check_lib
+from sscma.utils.check import check_lib # noqa
 
 
 def parse_args():