diff --git a/.github/dependabot.yml b/.github/dependabot.yml deleted file mode 100644 index a17f6a0be..000000000 --- a/.github/dependabot.yml +++ /dev/null @@ -1,11 +0,0 @@ -# To get started with Dependabot version updates, you'll need to specify which -# package ecosystems to update and where the package manifests are located. -# Please see the documentation for all configuration options: -# https://help.github.com/github/administering-a-repository/configuration-options-for-dependency-updates - -version: 2 -updates: - - package-ecosystem: "pip" # See documentation for possible values - directory: "/requirements.txt" # Location of package manifests - schedule: - interval: "daily" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0f9c371c8..d05da6a16 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,7 @@ jobs: - name: set up python uses: actions/setup-python@v2 with: - python-version: "3.10" + python-version: "3.11" - name: install run: | @@ -55,11 +55,11 @@ jobs: python3 -m mypy tests test: - name: test py${{ matrix.python-version }} + name: test runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.10"] + python-version: ["3.11"] steps: - uses: actions/checkout@v2 @@ -89,7 +89,7 @@ jobs: run: python3 -m coverage run --source=vis4d -m pytest --pyargs tests - name: test coverage - run: python3 -m coverage report --fail-under=88 -m + run: python3 -m coverage report -m - name: build run: python3 -m build diff --git a/.github/workflows/deploy_doc.yml b/.github/workflows/deploy_doc.yml deleted file mode 100644 index 2682ab325..000000000 --- a/.github/workflows/deploy_doc.yml +++ /dev/null @@ -1,42 +0,0 @@ -# This workflow will deploy the documentation in docs/ to http://docs.vis.xyz/4d/ - -name: deploy docs -on: - push: - branches: - - main - paths: - - "docs/**" - -jobs: - deploy: - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@v2 - - name: Set up Python 3.10 - uses: actions/setup-python@v1 - with: - python-version: '3.10' - - name: Install dependencies - run: | - eval `ssh-agent -s` - ssh-add - <<< '${{ secrets.CUDA_OPS_REPO_KEY }}' - python3 -m pip install --upgrade pip - bash ./scripts/install_cpu_dep_full.sh - python3 -m pip install -e . - python3 -m pip install --ignore-installed -r docs/requirements.txt - python3 -m pip freeze - - name: website build - run: | - cd docs - make html - - name: deploy website to AWS - uses: jakejarvis/s3-sync-action@master - with: - args: --acl public-read --follow-symlinks --exclude ".DS_Store" - env: - AWS_S3_BUCKET: docs.vis.xyz/4d - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - SOURCE_DIR: "docs/build/html/" diff --git a/.gitignore b/.gitignore index 7b72005b9..fc099e2e1 100644 --- a/.gitignore +++ b/.gitignore @@ -61,14 +61,13 @@ doc/media # python build build/ dist/ -scalabel.egg-info +vis4d.egg-info # coverage .coverage* # package default workspace vis4d-workspace -vis4d.egg-info *.tmp diff --git a/.gitmodules b/.gitmodules index f5eb86837..f2c8f27b6 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,4 +1,4 @@ [submodule "vis4d-test-data"] - path = tests/vis4d-test-data - url = git@github.com:SysCV/vis4d-test-data.git + path = tests/vis4d-test-data + url = git@github.com:SysCV/vis4d-test-data.git branch = main diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml deleted file mode 100644 index 84bbcf27e..000000000 --- a/.pre-commit-config.yaml +++ /dev/null @@ -1,32 +0,0 @@ -fail_fast: true -repos: - - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.3.0 - hooks: - - id: end-of-file-fixer - - id: trailing-whitespace - - id: check-yaml - - id: check-merge-conflict - - id: requirements-txt-fixer - - id: debug-statements - - repo: https://github.com/psf/black - rev: 23.1.0 - hooks: - - id: black - - repo: https://github.com/pycqa/isort - rev: 5.10.1 - hooks: - - id: isort -# - repo: https://github.com/pre-commit/mirrors-pylint -# rev: v2.7.4 -# hooks: -# - id: pylint -# - repo: https://github.com/pycqa/pydocstyle -# rev: 6.1.1 -# hooks: -# - id: pydocstyle -# args: [--convention=google] -# - repo: https://github.com/pre-commit/mirrors-mypy -# rev: v0.991 -# hooks: -# - id: mypy diff --git a/.pylintrc b/.pylintrc index 35f4549fa..2142ce8a4 100644 --- a/.pylintrc +++ b/.pylintrc @@ -411,6 +411,7 @@ max-spelling-suggestions=2 # Maximum number of arguments for function / method max-args=20 +max-positional-arguments=20 # Maximum number of locals for function / method body max-locals=100 diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 9138a7fe5..d22002cf2 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -8,6 +8,9 @@ build: python: install: - requirements: docs/requirements.txt + - requirements: requirements/install.txt + - requirements: requirements/torch-lib.txt + - requirements: requirements/viewer.txt sphinx: configuration: docs/source/conf.py diff --git a/docs/requirements.txt b/docs/requirements.txt index b2a0ed4e7..afe9fbce9 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -5,4 +5,5 @@ sphinx_autodoc_defaultargs sphinx_autodoc_typehints sphinx_copybutton sphinx_design +sphinx_rtd_theme sphinxcontrib-aafig diff --git a/docs/source/conf.py b/docs/source/conf.py index 08f9b4c9a..f1dfa4be8 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -24,7 +24,7 @@ project = "Vis4D" copyright = "2022, ETH Zurich" -author = "Tobias Fischer" +author = "Vis4D Team" # -- General configuration --------------------------------------------------- @@ -94,7 +94,7 @@ "Vis4D Documentation", author, "Vis4D", - "Dynamic Scene Understanding in Pytorch.", + "Dynamic Scene Understanding in PyTorch.", "Miscellaneous", ) ] @@ -102,13 +102,14 @@ # -- auto doc settings ------------------------------------------------------- autosummary_generate = True autodoc_member_order = "groupwise" -autoclass_content = "both" +# autoclass_content = "both" add_module_names = False # Remove namespaces from class/method signatures +autodoc_member_order = "bysource" autodoc_default_options = { "members": True, "methods": True, "special-members": "__call__", - "exclude-members": "_abc_impl,__init__", + "exclude-members": "_abc_impl,__repr__", } # -- Napoleon settings ------------------------------------------------------- @@ -117,7 +118,7 @@ # project. napoleon_google_docstring = True napoleon_numpy_docstring = False -napoleon_include_init_with_doc = False +napoleon_include_init_with_doc = True napoleon_include_private_with_doc = False napoleon_include_special_with_doc = True napoleon_use_admonition_for_examples = False diff --git a/docs/source/datasets.rst b/docs/source/datasets.rst new file mode 100644 index 000000000..3acad4ec9 --- /dev/null +++ b/docs/source/datasets.rst @@ -0,0 +1,3 @@ +******** +Datasets +******** \ No newline at end of file diff --git a/docs/source/dev_guide/cli.rst b/docs/source/dev_guide/cli.rst index fe62471bd..538e3ede5 100644 --- a/docs/source/dev_guide/cli.rst +++ b/docs/source/dev_guide/cli.rst @@ -1,10 +1,11 @@ ### CLI ### + We provide a command line interface for training and evaluating your models. Assuming you have installed the package using pip, you can use the command `vis4d` to access the CLI. -Alternatively, you can run the CLI using `python -m vis4d.engine.cli` or `python -m vis4d.pl.cli` if you want to use the PyTorch Lightning version. +Alternatively, you can run the CLI using `python -m vis4d.engine.run` or `python -m vis4d.pl.run` if you want to use the PyTorch Lightning version. The CLI relies on a configuration file to specify each experiment. We use `ml_collections `_ as underlying framework to define the configuration files. You can read up on our configuration files in the `Config System `_ section. @@ -12,6 +13,7 @@ You can read up on our configuration files in the `Config System `_ section. @@ -20,6 +22,7 @@ We support both, our own training engine as well as `PyTorch Lightning diff --git a/docs/source/user_guide/3D_visualization.ipynb b/docs/source/user_guide/3D_visualization.ipynb index 1d45f8138..b93550516 100644 --- a/docs/source/user_guide/3D_visualization.ipynb +++ b/docs/source/user_guide/3D_visualization.ipynb @@ -20,7 +20,7 @@ "os.environ[\"WEBRTC_IP\"] = \"127.0.0.1\"\n", "\n", "import pickle\n", - "from vis4d.vis.functional import show_points\n", + "from vis4d.vis.pointcloud.functional import show_points\n", "import numpy as np" ] }, diff --git a/docs/source/user_guide/faster_rcnn_example.py b/docs/source/user_guide/faster_rcnn_example.py index b639930d2..7703770c1 100644 --- a/docs/source/user_guide/faster_rcnn_example.py +++ b/docs/source/user_guide/faster_rcnn_example.py @@ -6,7 +6,7 @@ import lightning.pytorch as pl import numpy as np -from torch.optim import SGD +from torch.optim.sgd import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR from vis4d.config import class_config @@ -137,7 +137,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir, refresh_rate=1) + callbacks = get_default_callbacks_cfg(refresh_rate=1) # Evaluator callbacks.append( diff --git a/docs/source/user_guide/getting_started.ipynb b/docs/source/user_guide/getting_started.ipynb index db99ef46e..088244787 100644 --- a/docs/source/user_guide/getting_started.ipynb +++ b/docs/source/user_guide/getting_started.ipynb @@ -251,7 +251,7 @@ "from vis4d.model.detect.faster_rcnn import FasterRCNN\n", "\n", "from vis4d.data.const import CommonKeys as K\n", - "from vis4d.vis.functional.image import imshow_bboxes\n", + "from vis4d.vis.image.functional import imshow_bboxes\n", "\n", "from vis4d.config import instantiate_classes\n", "from vis4d.zoo.base.datasets.coco import get_coco_detection_cfg" diff --git a/docs/source/user_guide/visualization.ipynb b/docs/source/user_guide/visualization.ipynb index 30e92e35a..6c1c285c0 100644 --- a/docs/source/user_guide/visualization.ipynb +++ b/docs/source/user_guide/visualization.ipynb @@ -15,22 +15,12 @@ "cell_type": "code", "execution_count": 1, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Jupyter environment detected. Enabling Open3D WebVisualizer.\n", - "[Open3D INFO] WebRTC GUI backend enabled.\n", - "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n" - ] - } - ], + "outputs": [], "source": [ "from __future__ import annotations\n", "\n", "from vis4d.common.typing import NDArrayF64, NDArrayI64\n", - "from vis4d.vis.functional import imshow_bboxes, imshow_masks, imshow_topk_bboxes, imshow, draw_bboxes, draw_masks, imshow_track_matches\n", + "from vis4d.vis.image.functional import imshow_bboxes, imshow_masks, imshow_topk_bboxes, imshow, draw_bboxes, draw_masks, imshow_track_matches\n", "\n", "import pickle\n", "import numpy as np" @@ -453,7 +443,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.8" + "version": "3.11.9" }, "vscode": { "interpreter": { diff --git a/pyproject.toml b/pyproject.toml index 9a06910a8..4aa3652cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ line_length = 79 [tool.pyright] include = ["vis4d"] +typeCheckingMode = "off" [tool.coverage] [tool.coverage.report] @@ -112,12 +113,12 @@ plugins = ["numpy.typing.mypy_plugin"] [project] name = "vis4d" -version = "0.1.2" -authors = [{name = "VIS @ ETH", email = "i@yf.io"}] +version = "0.1.3" +authors = [{name = "Vis4D Team"}] description = "Vis4D Python package for Visual 4D scene understanding" readme = "README.md" license = {text = "Apache 2.0"} -requires-python = ">=3.10" +requires-python = ">=3.11" classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", diff --git a/requirements/install.txt b/requirements/install.txt index c1f6cfe52..09e8063ae 100644 --- a/requirements/install.txt +++ b/requirements/install.txt @@ -6,8 +6,8 @@ devtools h5py jsonargparse[signatures] lightning -ml_collections==0.1.1 # Config interface. Need exact version since we overwrite internal functions -numpy>=1.21.0 +ml_collections==1.1.0 # Config interface. Need exact version since we overwrite internal functions +numpy>=1.21.0,<2.0.0 opencv-python pandas pillow @@ -24,6 +24,8 @@ torchvision>=0.15.1 tqdm utm wheel +scipy +scipy-stubs # packages for datasets bdd100k diff --git a/requirements/torch-lib.txt b/requirements/torch-lib.txt index 423c07689..a3e6a1f4b 100644 --- a/requirements/torch-lib.txt +++ b/requirements/torch-lib.txt @@ -1 +1 @@ -git+ssh://git@github.com/SysCV/vis4d_cuda_ops.git +git+https://github.com/SysCV/vis4d_cuda_ops.git diff --git a/requirements/viewer.txt b/requirements/viewer.txt index 6c0ba1c16..2cf81fd62 100644 --- a/requirements/viewer.txt +++ b/requirements/viewer.txt @@ -1 +1,2 @@ open3d +matplotlib>3.9 diff --git a/tests/common/array_test.py b/tests/common/array_test.py index d0db7cd13..953bfe5c7 100644 --- a/tests/common/array_test.py +++ b/tests/common/array_test.py @@ -6,7 +6,7 @@ import numpy as np import torch -from vis4d.common.array import array_to_numpy, arrays_to_numpy +from vis4d.common.array import array_to_numpy class TestConvertToArray(unittest.TestCase): @@ -52,11 +52,3 @@ def test_dim_shaping(self) -> None: # And the right if we can not remove anything from the left anymore self.assertEqual(array_to_numpy(data.copy(), 2).shape, (2, 3)) - - def test_array_to_numpys_multiple(self) -> None: - """Test that multiple arrays are converted correctly.""" - out = arrays_to_numpy( - np.random.rand(2, 3), np.random.rand(1, 1, 2, 3), n_dims=3 - ) - for arr in out: - self.assertEqual(arr.shape, (1, 2, 3)) diff --git a/tests/config/show_connection_test.py b/tests/config/show_connection_test.py index 1011020f7..079f91065 100644 --- a/tests/config/show_connection_test.py +++ b/tests/config/show_connection_test.py @@ -24,10 +24,10 @@ def test_show_frcnn(self) -> None: model = instantiate_classes(config.model) # Change the data root of evaluator callback to the test data - config.callbacks[3].init_args.evaluator.init_args.data_root = ( + config.callbacks[2].init_args.evaluator.init_args.data_root = ( "tests/vis4d-test-data/coco_test" ) - config.callbacks[3].init_args.evaluator.init_args.split = "train" + config.callbacks[2].init_args.evaluator.init_args.split = "train" callbacks = [instantiate_classes(cb) for cb in config.callbacks] diff --git a/tests/config/sweep_test.py b/tests/config/sweep_test.py deleted file mode 100644 index 681488adc..000000000 --- a/tests/config/sweep_test.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Vis4d Sweep Config Tests.""" - -import unittest - -import numpy as np -from ml_collections import ConfigDict - -from vis4d.config.replicator import replicate_config -from vis4d.config.sweep import grid_search -from vis4d.config.typing import ExperimentConfig - - -class TestSweep(unittest.TestCase): - """Test for config sweeps.""" - - def test_one_param_sweep(self) -> None: - """Test if one param config sweep works.""" - expected = np.linspace(0.001, 0.01, 3) - - exp_names_expected = [f"test_lr_{lr:.3f}_" for lr in expected] - - sweep_config = grid_search("lr", list(expected)) - sweep_config.suffix = "lr_{lr:.3f}_" - config = ExperimentConfig() - config.lr = 0 - config.experiment_name = "test" - config.value_mode() - config_iter = replicate_config( - config, - method=sweep_config.method, - sampling_args=sweep_config.sampling_args, - fstring=sweep_config.get("suffix", ""), - ) - - lr_actual: list[float] = [] - exp_names: list[str] = [] - for c in config_iter: - lr_actual.append(c.lr) - exp_names.append(c.experiment_name) - lr_actual = np.array(lr_actual) # type: ignore - - self.assertTrue(np.allclose(expected, lr_actual)) - for pred, gt in zip(exp_names, exp_names_expected): - self.assertEqual(pred, gt) - - def test_one_nested_param_sweep(self) -> None: - """Test if one param config sweep works when it is nested.""" - expected = np.linspace(0.001, 0.01, 3) - - exp_names_expected = [f"test_lr_{lr:.3f}_" for lr in expected] - - sweep_config = grid_search("params.lr", list(expected)) - sweep_config.suffix = "lr_{params.lr:.3f}_" - config = ExperimentConfig() - config.params = ConfigDict() - config.params.lr = 0 - config.experiment_name = "test" - config.value_mode() - config_iter = replicate_config( - config, - method=sweep_config.method, - sampling_args=sweep_config.sampling_args, - fstring=sweep_config.get("suffix", ""), - ) - - lr_actual: list[float] = [] - exp_names: list[str] = [] - for c in config_iter: - lr_actual.append(c.params.lr) - exp_names.append(c.experiment_name) - lr_actual = np.array(lr_actual) # type: ignore - - self.assertTrue(np.allclose(expected, lr_actual)) - for pred, gt in zip(exp_names, exp_names_expected): - self.assertEqual(pred, gt) - - def test_two_param_sweeps(self) -> None: - """Test to sweep over two parameters (lr, bs).""" - learning_rates = np.linspace(0.001, 0.01, 3) - batch_sizes = [4, 8, 16] - - exp_names_expected = [] - lr_expected = [] - bs_expected = [] - for lr in learning_rates: - for bs in batch_sizes: - exp_names_expected.append(f"test_lr_{lr:.3f}_bs_{bs}_") - lr_expected.append(lr) - bs_expected.append(bs) - - sweep_config = grid_search( - ["lr", "bs"], [list(learning_rates), batch_sizes] - ) - sweep_config.suffix = "lr_{lr:.3f}_bs_{bs}_" - config = ExperimentConfig() - config.lr = 0 - config.experiment_name = "test" - config.value_mode() - config_iter = replicate_config( - config, - method=sweep_config.method, - sampling_args=sweep_config.sampling_args, - fstring=sweep_config.get("suffix", ""), - ) - - lr_actual: list[float] = [] - bs_actual: list[int] = [] - exp_names: list[str] = [] - for c in config_iter: - lr_actual.append(c.lr) - bs_actual.append(c.bs) - exp_names.append(c.experiment_name) - - self.assertTrue(np.allclose(lr_expected, lr_actual)) - self.assertTrue(np.allclose(bs_expected, bs_actual)) - for pred, gt in zip(exp_names, exp_names_expected): - self.assertEqual(pred, gt) diff --git a/tests/data/datasets/coco_test.py b/tests/data/datasets/coco_test.py index ef09035cd..21ee6f97d 100644 --- a/tests/data/datasets/coco_test.py +++ b/tests/data/datasets/coco_test.py @@ -46,17 +46,6 @@ def test_sample(self) -> None: """Test if sample loaded correctly.""" item = self.coco[0] item = ToTensor().apply_to_data([item])[0] # pylint: disable=no-member - self.assertEqual( - tuple(item.keys()), - ( - "sample_names", - "images", - "input_hw", - "boxes2d", - "boxes2d_classes", - "instance_masks", - ), - ) self.assertEqual(item[K.sample_names], 37777) self.assertEqual(item[K.input_hw], [230, 352]) @@ -142,14 +131,6 @@ def test_sample(self) -> None: """Test if sample loaded correctly.""" item = self.coco[0] item = ToTensor().apply_to_data([item])[0] # pylint: disable=no-member - assert tuple(item.keys()) == ( - "sample_names", - "images", - "input_hw", - "boxes2d_classes", - "instance_masks", - "seg_masks", - ) self.assertEqual(item[K.sample_names], 37777) self.assertEqual(item[K.input_hw], [230, 352]) diff --git a/tests/data/io/to_hdf5_test.py b/tests/data/io/to_hdf5_test.py index 5e469746f..caf619233 100644 --- a/tests/data/io/to_hdf5_test.py +++ b/tests/data/io/to_hdf5_test.py @@ -4,7 +4,7 @@ import unittest from tests.util import get_test_data -from vis4d.data.io.hdf5 import convert_dataset +from vis4d.data.io.to_hdf5 import convert_dataset class TestHDF5(unittest.TestCase): diff --git a/tests/data/transforms/mosaic_test.py b/tests/data/transforms/mosaic_test.py index 1e8608364..baacf9bdb 100644 --- a/tests/data/transforms/mosaic_test.py +++ b/tests/data/transforms/mosaic_test.py @@ -4,7 +4,6 @@ import unittest import numpy as np -import torch from PIL import Image from tests.util import get_test_file @@ -48,7 +47,7 @@ def test_mosaic_images(self) -> None: assert len(data["transforms"]["mosaic"]["im_scales"]) == 4 assert np.allclose( data[K.images], - torch.load(get_test_file("mosaic_images.npy")), + np.load(get_test_file("mosaic_images.npy")), atol=1e-4, ) @@ -75,12 +74,19 @@ def test_mosaic_boxes2d(self) -> None: data = params.apply_to_data([copy.deepcopy(data) for _ in range(4)]) data = MosaicImages().apply_to_data(data) data = transform.apply_to_data(data)[0] - box_data = [ + + assert np.allclose( data[K.boxes2d], + np.load(get_test_file("mosaic_boxes2d.npy")), + atol=1e-4, + ) + assert np.allclose( data[K.boxes2d_classes], + np.load(get_test_file("mosaic_boxes2d_classes.npy")), + atol=1e-4, + ) + assert np.allclose( data[K.boxes2d_track_ids], - ] - for pred, gt in zip( - box_data, torch.load(get_test_file("mosaic_boxes2d.npy")) - ): - assert np.allclose(pred, gt, atol=1e-4) + np.load(get_test_file("mosaic_boxes2d_track_ids.npy")), + atol=1e-4, + ) diff --git a/tests/data/transforms/photometric_test.py b/tests/data/transforms/photometric_test.py index f5c5b8da7..b4e6268f4 100644 --- a/tests/data/transforms/photometric_test.py +++ b/tests/data/transforms/photometric_test.py @@ -4,7 +4,6 @@ import unittest import numpy as np -import torch from PIL import Image from tests.util import get_test_file @@ -36,7 +35,7 @@ def test_random_gamma(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0], - torch.load(get_test_file("random_gamma_gt.npy")), + np.load(get_test_file("random_gamma_gt.npy")), atol=1e-4, ) @@ -47,7 +46,7 @@ def test_random_gamma(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0][..., [2, 1, 0]], - torch.load(get_test_file("random_gamma_gt.npy")), + np.load(get_test_file("random_gamma_gt.npy")), atol=1e-4, ) @@ -59,7 +58,7 @@ def test_random_brightness(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0], - torch.load(get_test_file("random_brightness_gt.npy")), + np.load(get_test_file("random_brightness_gt.npy")), atol=1e-4, ) @@ -71,7 +70,7 @@ def test_random_contrast(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0], - torch.load(get_test_file("random_contrast_gt.npy")), + np.load(get_test_file("random_contrast_gt.npy")), atol=1e-4, ) @@ -83,7 +82,7 @@ def test_random_saturation(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0], - torch.load(get_test_file("random_saturation_gt.npy")), + np.load(get_test_file("random_saturation_gt.npy")), atol=1e-4, ) @@ -96,7 +95,7 @@ def test_random_hue(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0], - torch.load(get_test_file("random_hue_gt.npy")), + np.load(get_test_file("random_hue_gt.npy")), atol=1e-4, ) @@ -107,7 +106,7 @@ def test_random_hue(self): self.assertEqual(data[K.images].shape, (1, 230, 352, 3)) assert np.allclose( data[K.images][0][..., [2, 1, 0]], - torch.load(get_test_file("random_hue_gt.npy")), + np.load(get_test_file("random_hue_gt.npy")), atol=1e-4, ) diff --git a/tests/data/transforms/testcases/mosaic_boxes2d.npy b/tests/data/transforms/testcases/mosaic_boxes2d.npy index b05eafec5..76d364a6f 100644 Binary files a/tests/data/transforms/testcases/mosaic_boxes2d.npy and b/tests/data/transforms/testcases/mosaic_boxes2d.npy differ diff --git a/tests/data/transforms/testcases/mosaic_boxes2d_classes.npy b/tests/data/transforms/testcases/mosaic_boxes2d_classes.npy new file mode 100644 index 000000000..64f7b0ba3 Binary files /dev/null and b/tests/data/transforms/testcases/mosaic_boxes2d_classes.npy differ diff --git a/tests/data/transforms/testcases/mosaic_boxes2d_track_ids.npy b/tests/data/transforms/testcases/mosaic_boxes2d_track_ids.npy new file mode 100644 index 000000000..922378b04 Binary files /dev/null and b/tests/data/transforms/testcases/mosaic_boxes2d_track_ids.npy differ diff --git a/tests/data/transforms/testcases/mosaic_images.npy b/tests/data/transforms/testcases/mosaic_images.npy index d8a0f88ee..2b618624a 100644 Binary files a/tests/data/transforms/testcases/mosaic_images.npy and b/tests/data/transforms/testcases/mosaic_images.npy differ diff --git a/tests/data/transforms/testcases/random_brightness_gt.npy b/tests/data/transforms/testcases/random_brightness_gt.npy index f1f7fc6c2..ca36ca883 100644 Binary files a/tests/data/transforms/testcases/random_brightness_gt.npy and b/tests/data/transforms/testcases/random_brightness_gt.npy differ diff --git a/tests/data/transforms/testcases/random_contrast_gt.npy b/tests/data/transforms/testcases/random_contrast_gt.npy index 9b9a927e5..a668cd3f7 100644 Binary files a/tests/data/transforms/testcases/random_contrast_gt.npy and b/tests/data/transforms/testcases/random_contrast_gt.npy differ diff --git a/tests/data/transforms/testcases/random_gamma_gt.npy b/tests/data/transforms/testcases/random_gamma_gt.npy index d52062a46..11d418f1c 100644 Binary files a/tests/data/transforms/testcases/random_gamma_gt.npy and b/tests/data/transforms/testcases/random_gamma_gt.npy differ diff --git a/tests/data/transforms/testcases/random_hue_gt.npy b/tests/data/transforms/testcases/random_hue_gt.npy index e30d500f0..6254c8028 100644 Binary files a/tests/data/transforms/testcases/random_hue_gt.npy and b/tests/data/transforms/testcases/random_hue_gt.npy differ diff --git a/tests/data/transforms/testcases/random_saturation_gt.npy b/tests/data/transforms/testcases/random_saturation_gt.npy index 24f486e53..1f4f6054c 100644 Binary files a/tests/data/transforms/testcases/random_saturation_gt.npy and b/tests/data/transforms/testcases/random_saturation_gt.npy differ diff --git a/tests/engine/callbacks/checkpoint_test.py b/tests/engine/callbacks/checkpoint_test.py deleted file mode 100644 index 560a433f1..000000000 --- a/tests/engine/callbacks/checkpoint_test.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Test cases for checkpoint callback.""" - -import shutil -import tempfile -import unittest - -from torch.optim import SGD - -from tests.util import MOCKLOSS, MockModel -from vis4d.config import class_config -from vis4d.engine.callbacks import CheckpointCallback, TrainerState - -from ..optim.optimizer_test import get_optimizer - - -class TestCheckpointCallback(unittest.TestCase): - """Test cases for callback functions.""" - - def setUp(self) -> None: - """Creates a tmp directory and setup callback.""" - self.test_dir = tempfile.mkdtemp() - - self.callback = CheckpointCallback(save_prefix=self.test_dir) - - self.callback.setup() - - optimizers, lr_scheulders = get_optimizer( - MockModel(0), class_config(SGD, lr=0.01) - ) - - self.trainer_state = TrainerState( - current_epoch=0, - num_epochs=0, - global_step=0, - train_dataloader=None, - num_train_batches=None, - test_dataloader=None, - num_test_batches=None, - optimizers=optimizers, - lr_schedulers=lr_scheulders, - ) - - def tearDown(self) -> None: - """Removes the tmp directory after the test.""" - shutil.rmtree(self.test_dir) - - def test_on_train_epoch_end(self) -> None: - """Test on_train_epoch_end function.""" - self.callback.on_train_epoch_end( - self.trainer_state, MockModel(0), MOCKLOSS - ) diff --git a/tests/engine/callbacks/ema_test.py b/tests/engine/callbacks/ema_test.py index 2342b3f2b..5553f1b7d 100644 --- a/tests/engine/callbacks/ema_test.py +++ b/tests/engine/callbacks/ema_test.py @@ -2,10 +2,11 @@ import unittest +import lightning.pytorch as pl import torch -from tests.util import MOCKLOSS, MockModel -from vis4d.engine.callbacks import EMACallback, TrainerState +from tests.util import MockModel +from vis4d.engine.callbacks import EMACallback from vis4d.model.adapter import ModelEMAAdapter @@ -14,29 +15,21 @@ class TestEMACallback(unittest.TestCase): def setUp(self) -> None: """Setup callback.""" + self.trainer = pl.Trainer() + self.training_module = pl.LightningModule() + self.callback = EMACallback() - self.callback.setup() + self.callback.setup(self.trainer, self.training_module, stage="fit") self.model = ModelEMAAdapter(MockModel(0)) self.model.ema_model.linear.weight.fill_(0) self.model.ema_model.linear.bias.fill_(0) - self.trainer_state = TrainerState( - current_epoch=0, - num_epochs=0, - global_step=0, - train_dataloader=None, - num_train_batches=None, - test_dataloader=None, - num_test_batches=None, - ) - def test_ema_callback(self) -> None: """Test EMA callback function.""" self.callback.on_train_batch_end( - self.trainer_state, + self.trainer, self.model, - MOCKLOSS, outputs={}, batch={}, batch_idx=0, diff --git a/tests/engine/callbacks/evaluator_test.py b/tests/engine/callbacks/evaluator_test.py index 40dfbf457..0674766d5 100644 --- a/tests/engine/callbacks/evaluator_test.py +++ b/tests/engine/callbacks/evaluator_test.py @@ -4,11 +4,12 @@ import tempfile import unittest +import lightning.pytorch as pl import torch -from tests.util import MockModel, get_test_data +from tests.util import get_test_data from vis4d.data.const import CommonKeys as K -from vis4d.engine.callbacks import EvaluatorCallback, TrainerState +from vis4d.engine.callbacks import EvaluatorCallback from vis4d.engine.connectors import CallbackConnector from vis4d.eval.coco import COCODetectEvaluator from vis4d.zoo.base.datasets.coco import CONN_COCO_MASK_EVAL @@ -21,27 +22,20 @@ def setUp(self) -> None: """Creates a tmp directory and setup callback.""" self.test_dir = tempfile.mkdtemp() + self.trainer = pl.Trainer() + self.training_module = pl.LightningModule() + self.callback = EvaluatorCallback( evaluator=COCODetectEvaluator( data_root=get_test_data("coco_test"), split="train" ), save_predictions=True, metrics_to_eval=[COCODetectEvaluator.METRIC_DET], - save_prefix=self.test_dir, + output_dir=self.test_dir, test_connector=CallbackConnector(CONN_COCO_MASK_EVAL), ) - self.callback.setup() - - self.trainer_state = TrainerState( - current_epoch=0, - num_epochs=0, - global_step=0, - train_dataloader=None, - num_train_batches=None, - test_dataloader=None, - num_test_batches=None, - ) + self.callback.setup(self.trainer, self.training_module, stage="test") def tearDown(self) -> None: """Removes the tmp directory after the test.""" @@ -50,8 +44,8 @@ def tearDown(self) -> None: def test_evaluator_callback(self) -> None: """Test evaluator callback function.""" self.callback.on_test_batch_end( - self.trainer_state, - MockModel(0), + self.trainer, + self.training_module, outputs={ "boxes": { "boxes": [torch.zeros((2, 4))], @@ -64,7 +58,4 @@ def test_evaluator_callback(self) -> None: batch_idx=0, ) - log_dict = self.callback.on_test_epoch_end( - self.trainer_state, MockModel(0) - ) - self.assertEqual(log_dict["Det/AP"], 0.0) + self.callback.on_test_epoch_end(self.trainer, self.training_module) diff --git a/tests/engine/callbacks/logging_test.py b/tests/engine/callbacks/logging_test.py deleted file mode 100644 index b60676b47..000000000 --- a/tests/engine/callbacks/logging_test.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Test cases for logging callback.""" - -import unittest - -from tests.util import MOCKLOSS, MockModel -from vis4d.engine.callbacks import LoggingCallback, TrainerState - - -class TestLoggingCallback(unittest.TestCase): - """Test cases for callback functions.""" - - def setUp(self) -> None: - """Setup callback.""" - self.callback = LoggingCallback(refresh_rate=1) - - self.trainer_state = TrainerState( - current_epoch=0, - num_epochs=0, - global_step=0, - num_steps=0, - train_dataloader=None, - num_train_batches=1, - test_dataloader=None, - num_test_batches=[1], - ) - - def test_on_train_epoch_start(self) -> None: - """Test on_train_epoch_start function.""" - self.callback.on_train_epoch_start( - self.trainer_state, MockModel(0), MOCKLOSS - ) - - def test_on_train_batch_end(self) -> None: - """Test on_train_batch_end function.""" - self.trainer_state["metrics"] = {"loss1": 0, "loss2": 1} - - self.callback.on_train_batch_end( - self.trainer_state, - MockModel(0), - MOCKLOSS, - outputs={}, - batch={}, - batch_idx=0, - ) - - self.trainer_state.pop("metrics") - - def test_on_test_epoch_start(self) -> None: - """Test on_test_epoch_start function.""" - self.callback.on_test_epoch_start(self.trainer_state, MockModel(0)) - - def test_on_test_batch_end(self) -> None: - """Test on_test_batch_end function.""" - self.callback.on_test_batch_end( - self.trainer_state, - MockModel(0), - outputs={}, - batch={}, - batch_idx=0, - ) diff --git a/tests/engine/callbacks/visualizer_test.py b/tests/engine/callbacks/visualizer_test.py index 2a4f590e7..6a5b3fc77 100644 --- a/tests/engine/callbacks/visualizer_test.py +++ b/tests/engine/callbacks/visualizer_test.py @@ -4,11 +4,11 @@ import tempfile import unittest +import lightning.pytorch as pl import torch -from tests.util import MOCKLOSS, MockModel from vis4d.data.const import CommonKeys as K -from vis4d.engine.callbacks import TrainerState, VisualizerCallback +from vis4d.engine.callbacks import VisualizerCallback from vis4d.engine.connectors import CallbackConnector from vis4d.vis.image import BoundingBoxVisualizer from vis4d.zoo.base.data_connectors import CONN_BBOX_2D_VIS @@ -21,24 +21,17 @@ def setUp(self) -> None: """Creates a tmp directory and setup callback.""" self.test_dir = tempfile.mkdtemp() + self.trainer = pl.Trainer() + self.training_module = pl.LightningModule() + self.callback = VisualizerCallback( visualizer=BoundingBoxVisualizer(), - save_prefix=self.test_dir, + output_dir=self.test_dir, train_connector=CallbackConnector(CONN_BBOX_2D_VIS), test_connector=CallbackConnector(CONN_BBOX_2D_VIS), ) - self.callback.setup() - - self.trainer_state = TrainerState( - current_epoch=0, - num_epochs=0, - global_step=0, - train_dataloader=None, - num_train_batches=None, - test_dataloader=None, - num_test_batches=None, - ) + self.callback.setup(self.trainer, self.training_module, stage="fit") def tearDown(self) -> None: """Removes the tmp directory after the test.""" @@ -47,9 +40,8 @@ def tearDown(self) -> None: def test_on_train_batch_end(self) -> None: """Test on_train_batch_end function.""" self.callback.on_train_batch_end( - self.trainer_state, - MockModel(0), - MOCKLOSS, + self.trainer, + self.training_module, outputs={ "boxes": [torch.zeros((0, 4))], "scores": [torch.zeros((0,))], @@ -65,8 +57,8 @@ def test_on_train_batch_end(self) -> None: def test_on_test_batch_end(self) -> None: """Test the visualizer callback.""" self.callback.on_test_batch_end( - self.trainer_state, - MockModel(0), + self.trainer, + self.training_module, outputs={ "boxes": [torch.zeros((0, 4))], "scores": [torch.zeros((0,))], diff --git a/tests/pl/data_module_test.py b/tests/engine/data_module_test.py similarity index 96% rename from tests/pl/data_module_test.py rename to tests/engine/data_module_test.py index c4ce4a023..381ed48f7 100644 --- a/tests/pl/data_module_test.py +++ b/tests/engine/data_module_test.py @@ -7,7 +7,7 @@ from torch.utils.data.dataloader import DataLoader from tests.util import get_test_data -from vis4d.pl.data_module import DataModule +from vis4d.engine.data_module import DataModule from vis4d.zoo.base.datasets.coco import get_coco_detection_cfg diff --git a/tests/engine/optim/scheduler_test.py b/tests/engine/optim/scheduler_test.py index 226e8eb43..e94cf7a50 100644 --- a/tests/engine/optim/scheduler_test.py +++ b/tests/engine/optim/scheduler_test.py @@ -5,8 +5,8 @@ import torch import torch.nn.functional as F -from torch import optim from torch.optim.lr_scheduler import LRScheduler +from torch.optim.sgd import SGD from torch.testing import assert_close from vis4d.engine.optim.scheduler import ConstantLR, PolyLR, QuadraticLRWarmup @@ -68,7 +68,7 @@ def setUp(self) -> None: model = ToyModel() self.lr = 0.05 self.l2_mult = 10 - self.optimizer = optim.SGD( + self.optimizer = SGD( [ {"params": model.conv1.parameters()}, { diff --git a/tests/engine/trainer_test.py b/tests/engine/trainer_test.py index 5d56dd497..94434f81a 100644 --- a/tests/engine/trainer_test.py +++ b/tests/engine/trainer_test.py @@ -1,4 +1,4 @@ -"""Engine trainer tests.""" +"""Pytorch lightning utilities for unit tests.""" from __future__ import annotations @@ -6,10 +6,11 @@ import tempfile import unittest -import torch +from ml_collections import ConfigDict +from torch.optim.sgd import SGD from torch.utils.data import DataLoader, Dataset -from tests.util import MockModel, get_test_data +from tests.util import get_test_data from vis4d.config import class_config from vis4d.data.const import CommonKeys as K from vis4d.data.datasets import COCO @@ -27,7 +28,7 @@ to_tensor, ) from vis4d.data.typing import DictData -from vis4d.engine.callbacks import LoggingCallback +from vis4d.engine.callbacks import LoggingCallback, LRSchedulerCallback from vis4d.engine.connectors import ( DataConnector, LossConnector, @@ -35,11 +36,11 @@ pred_key, ) from vis4d.engine.loss_module import LossModule -from vis4d.engine.trainer import Trainer +from vis4d.engine.trainer import PLTrainer +from vis4d.engine.training_module import TrainingModule from vis4d.model.seg.semantic_fpn import SemanticFPN from vis4d.op.loss import SegCrossEntropyLoss - -from .optim.optimizer_test import get_optimizer +from vis4d.zoo.base import get_optimizer_cfg def seg_pipeline(data: list[DictData]) -> DictData: @@ -84,73 +85,70 @@ def get_test_dataloader( ) -class EngineTrainerTest(unittest.TestCase): - """Engine trainer test class.""" +def get_training_module(model_cfg: ConfigDict): + """Build mockup training module. + + Args: + model_cfg (ConfigDict): Pytorch model + """ + train_data_connector = DataConnector(key_mapping={K.images: K.images}) + test_data_connector = DataConnector(key_mapping={K.images: K.images}) + loss_module = LossModule( + { + "loss": SegCrossEntropyLoss(), + "connector": LossConnector( + key_mapping={ + "output": pred_key("outputs"), + "target": data_key(K.seg_masks), + } + ), + } + ) + + optimizer_cfg = get_optimizer_cfg(class_config(SGD, lr=0.01)) + return TrainingModule( + model_cfg=model_cfg, + optimizers_cfg=[optimizer_cfg], + loss_module=loss_module, + train_data_connector=train_data_connector, + test_data_connector=test_data_connector, + seed=1, + ) + + +class PLTrainerTest(unittest.TestCase): + """Pytorch lightning trainer test class.""" def setUp(self) -> None: - """Set up test.""" + """Setup.""" self.test_dir = tempfile.mkdtemp() - dataset = COCO( - get_test_data("coco_test"), - keys_to_load=[ - K.images, - K.original_images, - K.boxes2d_classes, - K.instance_masks, - ], - split="train", - ) - train_dataloader = get_train_dataloader(dataset, 2) - test_dataloader = get_test_dataloader(dataset) - train_data_connector = DataConnector(key_mapping={"images": K.images}) - test_data_connector = DataConnector( - key_mapping={"images": K.images, "original_hw": K.original_hw} - ) + callbacks = [LRSchedulerCallback(), LoggingCallback()] - self.model = SemanticFPN(num_classes=80) - - self.trainer = Trainer( - device=torch.device("cpu"), - output_dir=self.test_dir, - num_epochs=2, - train_data_connector=train_data_connector, - test_data_connector=test_data_connector, - callbacks=[LoggingCallback(refresh_rate=1)], - train_dataloader=train_dataloader, - test_dataloader=test_dataloader, + self.trainer = PLTrainer( + work_dir=self.test_dir, + exp_name="test", + version="test", + callbacks=callbacks, + max_steps=2, + devices=0, + num_sanity_val_steps=0, ) - def tearDown(self) -> None: - """Tear down test.""" - shutil.rmtree(self.test_dir) - - def test_fit(self) -> None: - """Test trainer training.""" - optimizers, lr_scheulders = get_optimizer( - MockModel(0), class_config(torch.optim.SGD, lr=0.01) - ) - loss_module = LossModule( - { - "loss": SegCrossEntropyLoss(), - "connector": LossConnector( - key_mapping={ - "output": pred_key("outputs"), - "target": data_key(K.seg_masks), - } - ), - } + model_cfg = class_config( + SemanticFPN, + num_classes=80, ) - self.trainer.fit(self.model, optimizers, lr_scheulders, loss_module) + self.training_module = get_training_module(model_cfg=model_cfg) - # TODO: add callback to check loss - - def test_test(self) -> None: - """Test trainer testing.""" - state = torch.random.get_rng_state() - torch.random.set_rng_state(torch.manual_seed(0).get_state()) + def tearDown(self) -> None: + """Tear down.""" + shutil.rmtree(self.test_dir) - self.trainer.test(self.model) + def test_train(self) -> None: + """Test training.""" + dataset = COCO(get_test_data("coco_test"), split="train") + train_dataloader = get_train_dataloader(dataset, 2) - torch.random.set_rng_state(state) + self.trainer.fit(self.training_module, train_dataloader) diff --git a/tests/engine/util_test.py b/tests/engine/util_test.py deleted file mode 100644 index f5453fa65..000000000 --- a/tests/engine/util_test.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Test engine util.""" - -from collections import namedtuple -from dataclasses import dataclass - -from vis4d.engine.util import apply_to_collection - - -@dataclass -class Test: - """Test dataclass.""" - - aaa: int - bbb: int - - -def test_apply_to_collection(): - """Test apply_to_collection.""" - data = {"a": 1, "b": 2, "c": 3} - data = apply_to_collection(data, int, lambda x: x * 2) - assert data == {"a": 2, "b": 4, "c": 6} - - data = {"a": 1, "b": 2, "c": 3} - data = apply_to_collection(data, (int, str), lambda x: x * 2) - assert data == {"a": 2, "b": 4, "c": 6} - - data = {"a": 1, "b": 2, "c": 3} - data = apply_to_collection(data, int, lambda x: x * 2, wrong_dtype=str) - assert data == {"a": 2, "b": 4, "c": 6} - - data = {"a": 1, "b": 2, "c": 3} - data = apply_to_collection( - data, int, lambda x: x * 2, wrong_dtype=str, include_none=False - ) - assert data == {"a": 2, "b": 4, "c": 6} - - data = {"a": 1, "b": 2, "c": 3} - data = apply_to_collection( - data, int, lambda x: x * 2, wrong_dtype=(str, int), include_none=False - ) - assert data == {"a": 1, "b": 2, "c": 3} - - data = {"a": 1, "b": 2, "c": 3} - data = apply_to_collection( - data, int, lambda x: x * 2, wrong_dtype=(str, int), include_none=True - ) - assert data == {"a": 1, "b": 2, "c": 3} - - # test with data as namedtuple or dataclass - data_cls = Test(1, 2) - data_cls = apply_to_collection(data_cls, int, lambda x: x * 2) - assert data_cls == Test(2, 4) - - data_cls = Test(1, 2) - data_cls = apply_to_collection(data_cls, (int, str), lambda x: x * 2) - assert data_cls == Test(2, 4) - - data_cls = Test(1, 2) - data_cls = apply_to_collection( - data_cls, int, lambda x: x * 2, wrong_dtype=str - ) - assert data_cls == Test(2, 4) - - data_cls = Test(1, 2) - data_cls = apply_to_collection( - data_cls, - int, - lambda x: x * 2, - wrong_dtype=(str, int), - include_none=False, - ) - assert data_cls == Test(1, 2) - - data_cls = Test(1, 2) - data_cls = apply_to_collection( - data_cls, - int, - lambda x: x * 2, - wrong_dtype=(str, int), - include_none=True, - ) - assert data_cls == Test(1, 2) - - data_tup = namedtuple("test", "aaa bbb")(1, 2) - data_tup = apply_to_collection(data_tup, int, lambda x: x * 2) - assert data_tup == namedtuple("test", "aaa bbb")(2, 4) - - data_tup = namedtuple("test", "aaa bbb")(1, 2) - data_tup = apply_to_collection(data_tup, (int, str), lambda x: x * 2) - assert data_tup == namedtuple("test", "aaa bbb")(2, 4) - - data_tup = namedtuple("test", "aaa bbb")(1, 2) - data_tup = apply_to_collection( - data_tup, int, lambda x: x * 2, wrong_dtype=str - ) - assert data_tup == namedtuple("test", "aaa bbb")(2, 4) - - data_tup = namedtuple("test", "aaa bbb")(1, 2) - data_tup = apply_to_collection( - data_tup, - int, - lambda x: x * 2, - wrong_dtype=(str, int), - include_none=False, - ) - assert data_tup == namedtuple("test", "aaa bbb")(1, 2) - - data_tup = namedtuple("test", "aaa bbb")(1, 2) - data_tup = apply_to_collection( - data_tup, - int, - lambda x: x * 2, - wrong_dtype=(str, int), - include_none=True, - ) - assert data_tup == namedtuple("test", "aaa bbb")(1, 2) diff --git a/tests/eval/scalabel/detect_test.py b/tests/eval/scalabel/detect_test.py deleted file mode 100644 index ba380acef..000000000 --- a/tests/eval/scalabel/detect_test.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Testcases for BDD100K detection evaluator.""" - -from __future__ import annotations - -import os.path as osp -import unittest - -from tests.util import generate_boxes, generate_instance_masks, get_test_data -from vis4d.data.datasets.shift import SHIFT -from vis4d.engine.connectors import ( - data_key, - get_inputs_for_pred_and_data, - pred_key, -) -from vis4d.eval.scalabel import ScalabelDetectEvaluator - -from ..utils import get_dataloader - - -class TestBDD100KDetectEvaluator(unittest.TestCase): - """BDD100K detection evaluator testcase class.""" - - CONN_SHIFT_EVAL = { - "frame_ids": data_key("frame_ids"), - "sample_names": data_key("sample_names"), - "sequence_names": data_key("sequence_names"), - "pred_boxes": pred_key("boxes"), - "pred_classes": pred_key("class_ids"), - "pred_scores": pred_key("scores"), - "pred_masks": pred_key("masks"), - } - - def test_shift_eval(self) -> None: - """Testcase for SHIFT evaluation.""" - batch_size = 1 - - annotations = osp.join( - get_test_data("shift_test"), - "discrete/images/val/front/det_2d.json", - ) - scalabel_eval = ScalabelDetectEvaluator(annotation_path=annotations) - assert str(scalabel_eval) == "Scalabel Detection Evaluator" - assert scalabel_eval.metrics == ["Det", "InsSeg"] - - # test gt - dataset = SHIFT(data_root=get_test_data("shift_test"), split="val") - test_loader = get_dataloader(dataset, batch_size) - - boxes, scores, classes, track_ids = generate_boxes( - 800, 1280, 4, batch_size, True - ) - masks, _, _ = generate_instance_masks(800, 1280, 4, batch_size) - output = { - "boxes": boxes, - "scores": scores, - "class_ids": classes, - "track_ids": track_ids, - "masks": masks, - } - - for batch in test_loader: - scalabel_eval.process_batch( - **get_inputs_for_pred_and_data( - self.CONN_SHIFT_EVAL, output, batch - ) - ) - - _, log_str = scalabel_eval.evaluate("Det") - assert isinstance(log_str, str) - assert log_str.count("\n") == 12 - - _, log_str = scalabel_eval.evaluate("InsSeg") - assert isinstance(log_str, str) - assert log_str.count("\n") == 12 diff --git a/tests/eval/shift/detect_test.py b/tests/eval/shift/detect_test.py deleted file mode 100644 index 37c84bc13..000000000 --- a/tests/eval/shift/detect_test.py +++ /dev/null @@ -1,57 +0,0 @@ -"""SHIFT eval test cases.""" - -from __future__ import annotations - -import unittest - -import torch - -from tests.eval.utils import get_dataloader -from tests.util import get_test_data -from vis4d.data.datasets.shift import SHIFT -from vis4d.eval.shift import SHIFTDetectEvaluator - - -class TestSegEvaluator(unittest.TestCase): - """Tests for SegEvaluator.""" - - data_root = get_test_data("shift_test") - evaluator = SHIFTDetectEvaluator( - annotation_path=( - f"{data_root}/discrete/images/val/front/det_insseg_2d.json" - ) - ) - dataset = SHIFT( - data_root=data_root, - split="val", - keys_to_load=[ - "images", - "boxes2d", - "boxes2d_classes", - "instance_masks", - ], - ) - test_loader = get_dataloader(dataset, 1, sensors=["front"]) - - def test_shift_perfect_prediction(self) -> None: - """Tests when predictions are correct.""" - for batch in self.test_loader: - self.evaluator.process_batch( - frame_ids=batch["frame_ids"], - sample_names=batch["sample_names"], - sequence_names=batch["sequence_names"], - pred_boxes=batch["front"]["boxes2d"], - pred_classes=batch["front"]["boxes2d_classes"], - pred_scores=[ - torch.ones_like(batch["front"]["boxes2d_classes"][0]) - ], - pred_masks=batch["front"]["instance_masks"], - ) - - metrics, _ = self.evaluator.evaluate("Det") - self.assertAlmostEqual(metrics["AP"], 100.0, places=2) - self.assertAlmostEqual(metrics["AP/pedestrian"], 100.0, places=2) - - metrics, _ = self.evaluator.evaluate("InsSeg") - self.assertAlmostEqual(metrics["AP"], 100.0, places=2) - self.assertAlmostEqual(metrics["AP/pedestrian"], 100, places=2) diff --git a/tests/model/detect/mask_rcnn_test.py b/tests/model/detect/mask_rcnn_test.py index 4b5ff954a..b830bbcf7 100644 --- a/tests/model/detect/mask_rcnn_test.py +++ b/tests/model/detect/mask_rcnn_test.py @@ -3,7 +3,7 @@ import unittest import torch -from torch import optim +from torch.optim.sgd import SGD from tests.util import get_test_data, get_test_file from vis4d.common.ckpt import load_model_checkpoint @@ -120,7 +120,7 @@ def test_train(self): ] ) - optimizer = optim.SGD(mask_rcnn.parameters(), lr=0.001, momentum=0.9) + optimizer = SGD(mask_rcnn.parameters(), lr=0.001, momentum=0.9) dataset = COCO(get_test_data("coco_test"), split="train") train_loader = get_train_dataloader(dataset, 2, (256, 256)) diff --git a/tests/model/detect/retinanet_test.py b/tests/model/detect/retinanet_test.py index 58f34645c..dc878e7e0 100644 --- a/tests/model/detect/retinanet_test.py +++ b/tests/model/detect/retinanet_test.py @@ -3,7 +3,7 @@ import unittest import torch -from torch import optim +from torch.optim.sgd import SGD from tests.util import get_test_data, get_test_file from vis4d.common.ckpt import load_model_checkpoint @@ -75,7 +75,7 @@ def test_train(self) -> None: retina_net.retinanet_head.box_sampler, ) - optimizer = optim.SGD(retina_net.parameters(), lr=0.001, momentum=0.9) + optimizer = SGD(retina_net.parameters(), lr=0.001, momentum=0.9) dataset = COCO(get_test_data("coco_test"), split="train") train_loader = get_train_dataloader(dataset, 2, (256, 256)) diff --git a/tests/model/detect/yolox_test.py b/tests/model/detect/yolox_test.py index 871cb729d..90bcb08c4 100644 --- a/tests/model/detect/yolox_test.py +++ b/tests/model/detect/yolox_test.py @@ -65,7 +65,7 @@ def test_inference(self) -> None: dets = yolox(inputs, images_hw, original_hw=images_hw) assert isinstance(dets, DetOut) - testcase_gt = torch.load(get_test_file("yolox.pt")) + testcase_gt = torch.load(get_test_file("yolox.pt"), weights_only=False) def _assert_eq( prediction: list[torch.Tensor], gts: list[torch.Tensor] diff --git a/tests/model/detect3d/bevformer_test.py b/tests/model/detect3d/bevformer_test.py index 7549f4805..6cdfcbb0d 100644 --- a/tests/model/detect3d/bevformer_test.py +++ b/tests/model/detect3d/bevformer_test.py @@ -134,7 +134,9 @@ def test_tiny_inference(self): if cur_iter == 1: break - testcase_gt_list = torch.load(get_test_file("bevformer_tiny.pt")) + testcase_gt_list = torch.load( + get_test_file("bevformer_tiny.pt"), weights_only=False + ) for dets, testcase_gt in zip(dets_list, testcase_gt_list): for pred, expected in zip(dets, testcase_gt): diff --git a/tests/model/seg/fcn_resnet_test.py b/tests/model/seg/fcn_resnet_test.py index dc1a2f985..4a9f90fcf 100644 --- a/tests/model/seg/fcn_resnet_test.py +++ b/tests/model/seg/fcn_resnet_test.py @@ -5,7 +5,7 @@ import unittest import torch -from torch import optim +from torch.optim.sgd import SGD from tests.util import get_test_data, get_test_file from vis4d.common.ckpt import load_model_checkpoint @@ -46,7 +46,7 @@ def test_train(self) -> None: """Test FCNResNet training.""" model = FCNResNet(base_model="resnet50", resize=(64, 64)) loss_fn = MultiLevelSegLoss(feature_idx=(4, 5), weights=[0.5, 1]) - optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9) + optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9) dataset = COCO( get_test_data("coco_test"), split="train", use_pascal_voc_cats=True ) diff --git a/tests/model/seg/semantic_fpn_test.py b/tests/model/seg/semantic_fpn_test.py index 53a1c4de7..19932f195 100644 --- a/tests/model/seg/semantic_fpn_test.py +++ b/tests/model/seg/semantic_fpn_test.py @@ -5,7 +5,7 @@ import unittest import torch -from torch import optim +from torch.optim.sgd import SGD from tests.util import get_test_data, get_test_file from vis4d.data.const import CommonKeys as K @@ -59,7 +59,7 @@ def test_train(self) -> None: """Test SemanticFPN training.""" model = SemanticFPN(num_classes=21) loss_fn = SegCrossEntropyLoss() - optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9) + optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9) train_loader = get_train_dataloader(self.dataset, 2) model.train() diff --git a/tests/model/track/qdtrack_test.py b/tests/model/track/qdtrack_test.py index 937424e3e..81e59b216 100644 --- a/tests/model/track/qdtrack_test.py +++ b/tests/model/track/qdtrack_test.py @@ -77,7 +77,9 @@ def test_inference_fasterrcnn(self): tracks = qdtrack(images, inputs_hw, original_hw, frame_ids) assert isinstance(tracks, TrackOut) print("Testcase file:", get_test_file("qdtrack.pt")) - testcase_gt = torch.load(get_test_file("qdtrack.pt")) + testcase_gt = torch.load( + get_test_file("qdtrack.pt"), weights_only=False + ) for pred_entry, expected_entry in zip(tracks, testcase_gt): for pred, expected in zip(pred_entry, expected_entry): print("PREDICTION:", pred.shape, pred) @@ -132,7 +134,9 @@ def test_inference_yolox(self): tracks = qdtrack(images, inputs_hw, original_hw, frame_ids) assert isinstance(tracks, TrackOut) print("Testcase file:", get_test_file("qdtrack-yolox.pt")) - testcase_gt = torch.load(get_test_file("qdtrack-yolox.pt")) + testcase_gt = torch.load( + get_test_file("qdtrack-yolox.pt"), weights_only=False + ) for pred_entry, expected_entry in zip(tracks, testcase_gt): for pred, expected in zip(pred_entry, expected_entry): print("PREDICTION:", pred.shape, pred) diff --git a/tests/model/track3d/cc_3dt_test.py b/tests/model/track3d/cc_3dt_test.py index 32d922b3d..b3570bd09 100644 --- a/tests/model/track3d/cc_3dt_test.py +++ b/tests/model/track3d/cc_3dt_test.py @@ -98,7 +98,9 @@ def test_r50_fpn_inference(self): if cur_iter == 1: break - testcase_gt_list = torch.load(f"{data_root}/cc_3dt.pt") + testcase_gt_list = torch.load( + f"{data_root}/cc_3dt.pt", weights_only=False + ) for tracks, testcase_gt in zip(tracks_list, testcase_gt_list): for pred, expected in zip(tracks, testcase_gt): diff --git a/tests/op/base/dla_test.py b/tests/op/base/dla_test.py index 2fc4b475f..35488a296 100644 --- a/tests/op/base/dla_test.py +++ b/tests/op/base/dla_test.py @@ -13,15 +13,11 @@ class TestDLA(unittest.TestCase): def test_dla46_c(self) -> None: """Testcase for DLA46-C.""" - dla46_c = DLA( - name="dla46_c", - # weights= - # "http://dl.yf.io/dla/models/imagenet/dla46_c-2bfd52c3.pth", - ) + dla46_c = DLA(name="dla46_c") out = dla46_c(self.inputs) self.assertEqual(len(out), 6) - channels = [16, 32, 64, 64, 128, 256] - for i in range(6): + channels = [3, 3, 64, 64, 128, 256] + for i in range(2, 6): feat = out[i] self.assertEqual(feat.shape[0], 2) self.assertEqual(feat.shape[1], channels[i]) @@ -33,36 +29,10 @@ def test_dla46x_c(self) -> None: dla46x_c = DLA(name="dla46x_c") out = dla46x_c(self.inputs) self.assertEqual(len(out), 6) - channels = [16, 32, 64, 64, 128, 256] - for i in range(6): + channels = [3, 3, 64, 64, 128, 256] + for i in range(2, 6): feat = out[i] self.assertEqual(feat.shape[0], 2) self.assertEqual(feat.shape[1], channels[i]) self.assertEqual(feat.shape[2], 32 / (2**i)) self.assertEqual(feat.shape[3], 32 / (2**i)) - - def test_dla_custom(self) -> None: - """Testcase for custom DLA.""" - dla_custom = DLA( - levels=(1, 1, 1, 2, 2, 1), - channels=(16, 32, 64, 128, 256, 512), - block="BasicBlock", - residual_root=True, - ) - out = dla_custom(self.inputs) - self.assertEqual(tuple(out[2].shape[2:]), (8, 8)) - dla_custom = DLA( - levels=(1, 1, 1, 2, 2, 1), - channels=(16, 32, 64, 128, 256, 512), - block="BasicBlock", - residual_root=False, - ) - out = dla_custom(self.inputs) - self.assertEqual(tuple(out[2].shape[2:]), (8, 8)) - out = dla_custom(self.inputs) - for i in range(6): - feat = out[i] - self.assertEqual(feat.shape[0], 2) - self.assertEqual(feat.shape[1], 16 * (2**i)) - self.assertEqual(feat.shape[2], 32 / (2**i)) - self.assertEqual(feat.shape[3], 32 / (2**i)) diff --git a/tests/op/detect/rpn_test.py b/tests/op/detect/rpn_test.py index f3a9bfe46..db629a6bc 100644 --- a/tests/op/detect/rpn_test.py +++ b/tests/op/detect/rpn_test.py @@ -26,7 +26,7 @@ def test_rpn_head(): ) # rpn forward rpn_out = rpn_head(test_features) - rpn_gt = torch.load(get_test_file("rpn_gt.pth")) + rpn_gt = torch.load(get_test_file("rpn_gt.pth"), weights_only=False) assert len(rpn_out.cls) == len(rpn_out.box) == num_feats for i in range(num_feats): wh_ = wh // 2**i diff --git a/tests/pl/__init__.py b/tests/pl/__init__.py deleted file mode 100644 index c8e39e3d2..000000000 --- a/tests/pl/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""PL Tests.""" diff --git a/tests/pl/trainer_test.py b/tests/pl/trainer_test.py deleted file mode 100644 index 2608983ff..000000000 --- a/tests/pl/trainer_test.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Pytorch lightning utilities for unit tests.""" - -from __future__ import annotations - -import shutil -import tempfile -import unittest - -from ml_collections import ConfigDict -from torch import optim -from torch.utils.data import DataLoader, Dataset - -from tests.util import get_test_data -from vis4d.config import class_config -from vis4d.data.const import CommonKeys as K -from vis4d.data.datasets import COCO -from vis4d.data.loader import DataPipe, build_train_dataloader -from vis4d.data.transforms import ( - compose, - mask, - normalize, - pad, - resize, - to_tensor, -) -from vis4d.data.typing import DictData -from vis4d.engine.callbacks import LoggingCallback -from vis4d.engine.connectors import ( - DataConnector, - LossConnector, - data_key, - pred_key, -) -from vis4d.engine.loss_module import LossModule -from vis4d.model.seg.semantic_fpn import SemanticFPN -from vis4d.op.loss import SegCrossEntropyLoss -from vis4d.pl.callbacks import CallbackWrapper, LRSchedulerCallback -from vis4d.pl.trainer import PLTrainer -from vis4d.pl.training_module import TrainingModule -from vis4d.zoo.base import get_optimizer_cfg - - -def seg_pipeline(data: list[DictData]) -> DictData: - """Default data pipeline.""" - return compose([pad.PadImages(value=255), to_tensor.ToTensor()])(data) - - -def get_train_dataloader(datasets: Dataset, batch_size: int) -> DataLoader: - """Get data loader for training.""" - preprocess_fn = compose( - [ - resize.GenResizeParameters((64, 64)), - resize.ResizeImages(), - resize.ResizeInstanceMasks(), - normalize.NormalizeImages(), - mask.ConvertInstanceMaskToSegMask(), - ] - ) - datapipe = DataPipe(datasets, preprocess_fn) - return build_train_dataloader( - datapipe, - batchprocess_fn=seg_pipeline, - samples_per_gpu=batch_size, - workers_per_gpu=1, - ) - - -def get_training_module(model_cfg: ConfigDict): - """Build mockup training module. - - Args: - model_cfg (ConfigDict): Pytorch model - """ - train_data_connector = DataConnector(key_mapping={K.images: K.images}) - test_data_connector = DataConnector(key_mapping={K.images: K.images}) - loss_module = LossModule( - { - "loss": SegCrossEntropyLoss(), - "connector": LossConnector( - key_mapping={ - "output": pred_key("outputs"), - "target": data_key(K.seg_masks), - } - ), - } - ) - - optimizer_cfg = get_optimizer_cfg(class_config(optim.SGD, lr=0.01)) - return TrainingModule( - model_cfg=model_cfg, - optimizers_cfg=[optimizer_cfg], - loss_module=loss_module, - train_data_connector=train_data_connector, - test_data_connector=test_data_connector, - seed=1, - ) - - -class PLTrainerTest(unittest.TestCase): - """Pytorch lightning trainer test class.""" - - def setUp(self) -> None: - """Setup.""" - self.test_dir = tempfile.mkdtemp() - - callbacks = [LRSchedulerCallback(), CallbackWrapper(LoggingCallback())] - - self.trainer = PLTrainer( - work_dir=self.test_dir, - exp_name="test", - version="test", - callbacks=callbacks, - max_steps=2, - devices=0, - num_sanity_val_steps=0, - ) - - model_cfg = class_config( - SemanticFPN, - num_classes=80, - ) - - self.training_module = get_training_module(model_cfg=model_cfg) - - def tearDown(self) -> None: - """Tear down.""" - shutil.rmtree(self.test_dir) - - def test_train(self) -> None: - """Test training.""" - dataset = COCO(get_test_data("coco_test"), split="train") - train_dataloader = get_train_dataloader(dataset, 2) - - self.trainer.fit(self.training_module, train_dataloader) diff --git a/tests/vis/image/bbox3d_visualizer_test.py b/tests/vis/image/bbox3d_visualizer_test.py index 72b6082e0..d5f105a13 100644 --- a/tests/vis/image/bbox3d_visualizer_test.py +++ b/tests/vis/image/bbox3d_visualizer_test.py @@ -30,7 +30,9 @@ def setUp(self) -> None: self.test_dir = tempfile.mkdtemp() data_root = get_test_data("nuscenes_test", absolute_path=False) - self.testcase_gt = torch.load(f"{data_root}/cc_3dt.pt") + self.testcase_gt = torch.load( + f"{data_root}/cc_3dt.pt", weights_only=False + ) dataset = NuScenes( data_root=data_root, diff --git a/tests/vis/image/bev_visualizer_test.py b/tests/vis/image/bev_visualizer_test.py index af837468b..1abfcc36f 100644 --- a/tests/vis/image/bev_visualizer_test.py +++ b/tests/vis/image/bev_visualizer_test.py @@ -27,7 +27,9 @@ def setUp(self) -> None: self.test_dir = tempfile.mkdtemp() data_root = get_test_data("nuscenes_test", absolute_path=False) - self.testcase_gt = torch.load(f"{data_root}/cc_3dt.pt") + self.testcase_gt = torch.load( + f"{data_root}/cc_3dt.pt", weights_only=False + ) dataset = NuScenes( data_root=data_root, diff --git a/tests/vis/image/bounding_box_visualizer_test.py b/tests/vis/image/bounding_box_visualizer_test.py index 7626409e3..69cc5208e 100644 --- a/tests/vis/image/bounding_box_visualizer_test.py +++ b/tests/vis/image/bounding_box_visualizer_test.py @@ -26,14 +26,15 @@ def setUp(self) -> None: testcase_gt = pickle.load(f) self.images: list[NDArrayF64] = testcase_gt["imgs"] - self.image_names: list[str] = ["0000", "0001"] self.boxes: list[NDArrayF64] = testcase_gt["boxes"] self.classes: list[NDArrayI64] = testcase_gt["classes"] self.scores: list[NDArrayF64] = testcase_gt["scores"] self.tracks = [np.arange(len(b)) for b in self.boxes] + cat_mapping = {v: k for k, v in COCO_COLOR_MAPPING.items()} + self.vis = BoundingBoxVisualizer( - n_colors=20, class_id_mapping=COCO_COLOR_MAPPING, vis_freq=1 + n_colors=20, cat_mapping=cat_mapping, vis_freq=1 ) def tearDown(self) -> None: @@ -55,7 +56,7 @@ def test_single_bbox_vis(self) -> None: """Tests visualization of single boudning boxes.""" self.vis.process_single_image( image=self.images[0], - image_name=self.image_names[0], + image_name="bbox_with_cts_target", boxes=self.boxes[0], scores=self.scores[0], class_ids=self.classes[0], @@ -64,7 +65,7 @@ def test_single_bbox_vis(self) -> None: self.vis.save_to_disk(cur_iter=1, output_folder=self.test_dir) self.assert_img_equal( - os.path.join(self.test_dir, "0000.png"), + os.path.join(self.test_dir, "bbox_with_cts_target.png"), get_test_file("bbox_with_cts_target.png"), ) self.vis.reset() @@ -73,7 +74,7 @@ def test_single_bbox_vis_no_tracks(self) -> None: """Tests visualization of single bounding boxes without track ids.""" self.vis.process_single_image( image=self.images[0], - image_name=self.image_names[0], + image_name="bbox_with_cs_target", boxes=self.boxes[0], scores=self.scores[0], class_ids=self.classes[0], @@ -82,7 +83,7 @@ def test_single_bbox_vis_no_tracks(self) -> None: self.vis.save_to_disk(cur_iter=1, output_folder=self.test_dir) self.assert_img_equal( - os.path.join(self.test_dir, "0000.png"), + os.path.join(self.test_dir, "bbox_with_cs_target.png"), get_test_file("bbox_with_cs_target.png"), ) self.vis.reset() @@ -91,7 +92,7 @@ def test_single_bbox_vis_only_class(self) -> None: """Tests visualization of single bounding boxes with only classes.""" self.vis.process_single_image( image=self.images[0], - image_name=self.image_names[0], + image_name="bbox_with_c_target", boxes=self.boxes[0], scores=None, class_ids=self.classes[0], @@ -100,7 +101,7 @@ def test_single_bbox_vis_only_class(self) -> None: self.vis.save_to_disk(cur_iter=1, output_folder=self.test_dir) self.assert_img_equal( - os.path.join(self.test_dir, "0000.png"), + os.path.join(self.test_dir, "bbox_with_c_target.png"), get_test_file("bbox_with_c_target.png"), ) self.vis.reset() @@ -110,7 +111,7 @@ def test_batched_vis(self) -> None: self.vis.process( cur_iter=1, images=self.images, - image_names=self.image_names, + image_names=[f"bbox_batched_{i}" for i in range(len(self.images))], boxes=self.boxes, scores=self.scores, class_ids=self.classes, @@ -120,7 +121,7 @@ def test_batched_vis(self) -> None: self.vis.save_to_disk(cur_iter=1, output_folder=self.test_dir) for i in range(2): self.assert_img_equal( - os.path.join(self.test_dir, f"000{i}.png"), + os.path.join(self.test_dir, f"bbox_batched_{i}.png"), get_test_file(f"bbox_batched_{i}.png"), ) self.vis.reset() diff --git a/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151603547590.png b/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151603547590.png index 0f3289438..e791daf2e 100644 Binary files a/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151603547590.png and b/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151603547590.png differ diff --git a/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151604048025.png b/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151604048025.png index b57dda62b..e380e26e4 100644 Binary files a/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151604048025.png and b/tests/vis/image/testcases/bbox3d/BEV/n008-2018-08-01-15-16-36-0400__LIDAR_TOP__1533151604048025.png differ diff --git a/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151603528113.png b/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151603528113.png index a2ca50d8c..69231fb89 100644 Binary files a/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151603528113.png and b/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151603528113.png differ diff --git a/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151604028370.png b/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151604028370.png index 9bc1a6cc0..412dd904f 100644 Binary files a/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151604028370.png and b/tests/vis/image/testcases/bbox3d/CAM_BACK_RIGHT/n008-2018-08-01-15-16-36-0400__CAM_BACK_RIGHT__1533151604028370.png differ diff --git a/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png b/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png index ca93529ed..7c8700621 100644 Binary files a/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png and b/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png differ diff --git a/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png b/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png index 35b23be4e..884d56579 100644 Binary files a/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png and b/tests/vis/image/testcases/bbox3d/CAM_FRONT/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png differ diff --git a/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png b/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png index ca93529ed..7c8700621 100644 Binary files a/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png and b/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151603512404.png differ diff --git a/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png b/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png index 35b23be4e..884d56579 100644 Binary files a/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png and b/tests/vis/image/testcases/bbox3d/n008-2018-08-01-15-16-36-0400__CAM_FRONT__1533151604012404.png differ diff --git a/tests/vis/image/testcases/bbox_batched_0.png b/tests/vis/image/testcases/bbox_batched_0.png index 8edb8944b..66ed551ba 100644 Binary files a/tests/vis/image/testcases/bbox_batched_0.png and b/tests/vis/image/testcases/bbox_batched_0.png differ diff --git a/tests/vis/image/testcases/bbox_batched_1.png b/tests/vis/image/testcases/bbox_batched_1.png index b1de0499b..ea18ded49 100644 Binary files a/tests/vis/image/testcases/bbox_batched_1.png and b/tests/vis/image/testcases/bbox_batched_1.png differ diff --git a/tests/vis/image/testcases/bbox_with_c_target.png b/tests/vis/image/testcases/bbox_with_c_target.png index dc5661876..c6b4eb5a8 100644 Binary files a/tests/vis/image/testcases/bbox_with_c_target.png and b/tests/vis/image/testcases/bbox_with_c_target.png differ diff --git a/tests/vis/image/testcases/bbox_with_cs_target.png b/tests/vis/image/testcases/bbox_with_cs_target.png index 759ea270b..4bad5459a 100644 Binary files a/tests/vis/image/testcases/bbox_with_cs_target.png and b/tests/vis/image/testcases/bbox_with_cs_target.png differ diff --git a/tests/vis/image/testcases/bbox_with_cts_target.png b/tests/vis/image/testcases/bbox_with_cts_target.png index 8edb8944b..66ed551ba 100644 Binary files a/tests/vis/image/testcases/bbox_with_cts_target.png and b/tests/vis/image/testcases/bbox_with_cts_target.png differ diff --git a/tests/vis4d-test-data b/tests/vis4d-test-data index 1e52c1948..f21d3e316 160000 --- a/tests/vis4d-test-data +++ b/tests/vis4d-test-data @@ -1 +1 @@ -Subproject commit 1e52c194859cfc09ee2f10af595303da3646d7d3 +Subproject commit f21d3e3163f7f2d4010ceab8f0202023f9ce5881 diff --git a/tests/zoo/cc_3dt_test.py b/tests/zoo/cc_3dt_test.py index 93f04e18a..67475a8de 100644 --- a/tests/zoo/cc_3dt_test.py +++ b/tests/zoo/cc_3dt_test.py @@ -96,6 +96,18 @@ def test_nusc_vis(self) -> None: ) ) + def test_nusc_test(self) -> None: + """Test the config.""" + cfg_gt = f"{self.gt_config_path}/cc_3dt_nusc_test.yaml" + + self.assertTrue( + compare_configs( + f"{self.config_prefix}.cc_3dt_nusc_test", + cfg_gt, + self.varying_keys, + ) + ) + def test_bevformer_base_velo_lstm_nusc(self) -> None: """Test the config.""" cfg_gt = ( diff --git a/tests/zoo/util.py b/tests/zoo/util.py index bc6a276d3..591b39341 100644 --- a/tests/zoo/util.py +++ b/tests/zoo/util.py @@ -2,17 +2,8 @@ from __future__ import annotations -import importlib - from tests.util import content_equal -from vis4d.config.typing import ExperimentConfig - - -def get_config_for_name(config_name: str) -> ExperimentConfig: - """Get config for name.""" - module = importlib.import_module("vis4d.zoo." + config_name) - - return module.get_config() +from vis4d.zoo.util import get_config_for_name def compare_configs( diff --git a/vis4d/__init__.py b/vis4d/__init__.py index 279ebc78c..39b70481b 100644 --- a/vis4d/__init__.py +++ b/vis4d/__init__.py @@ -7,7 +7,7 @@ import logging -__version__ = "0.1.2" +__version__ = "0.1.3" _root_logger = logging.getLogger() _logger = logging.getLogger(__name__) diff --git a/vis4d/common/__init__.py b/vis4d/common/__init__.py index 4f017c2aa..72476b68b 100644 --- a/vis4d/common/__init__.py +++ b/vis4d/common/__init__.py @@ -16,7 +16,6 @@ NDArrayUI8, TorchCheckpoint, TorchLossFunc, - TrainingModule, ) __all__ = [ @@ -35,5 +34,4 @@ "TorchLossFunc", "GenericFunc", "ListAny", - "TrainingModule", ] diff --git a/vis4d/common/array.py b/vis4d/common/array.py index 16c745caa..c520155ec 100644 --- a/vis4d/common/array.py +++ b/vis4d/common/array.py @@ -10,67 +10,74 @@ from vis4d.common.typing import ( ArrayLike, NDArrayBool, - NDArrayFloat, - NDArrayInt, + NDArrayF32, + NDArrayF64, + NDArrayI32, + NDArrayI64, NDArrayNumber, - NDArrayUInt, - NumpyBool, - NumpyFloat, - NumpyInt, - NumpyUInt, + NDArrayUI8, + NDArrayUI16, + NDArrayUI32, ) +# Bool dtypes @overload def array_to_numpy( - data: ArrayLike, n_dims: int | None, dtype: type[NumpyBool] + data: ArrayLike, n_dims: int | None, dtype: type[np.bool_] ) -> NDArrayBool: ... +# Float dtypes @overload def array_to_numpy( - data: ArrayLike, n_dims: int | None, dtype: type[NumpyFloat] -) -> NDArrayFloat: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.float32] +) -> NDArrayF32: ... @overload def array_to_numpy( - data: ArrayLike, n_dims: int | None, dtype: type[NumpyInt] -) -> NDArrayInt: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.float64] +) -> NDArrayF64: ... +# Int dtypes @overload def array_to_numpy( - data: ArrayLike, n_dims: int | None, dtype: type[NumpyUInt] -) -> NDArrayUInt: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.int32] +) -> NDArrayI32: ... @overload def array_to_numpy( - data: ArrayLike | None, n_dims: int | None, dtype: type[NumpyBool] -) -> NDArrayBool | None: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.int64] +) -> NDArrayI64: ... +# UInt dtypes @overload def array_to_numpy( - data: ArrayLike | None, n_dims: int | None, dtype: type[NumpyFloat] -) -> NDArrayFloat | None: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.uint8] +) -> NDArrayUI8: ... @overload def array_to_numpy( - data: ArrayLike | None, n_dims: int | None, dtype: type[NumpyInt] -) -> NDArrayInt | None: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.uint16] +) -> NDArrayUI16: ... @overload def array_to_numpy( - data: ArrayLike | None, n_dims: int | None, dtype: type[NumpyUInt] -) -> NDArrayUInt | None: ... + data: ArrayLike | None, n_dims: int | None, dtype: type[np.uint32] +) -> NDArrayUI32: ... +# Union of all dtypes @overload -def array_to_numpy(data: ArrayLike, n_dims: int | None) -> NDArrayNumber: ... +def array_to_numpy( + data: ArrayLike | None, n_dims: int | None +) -> NDArrayNumber: ... @overload @@ -81,7 +88,14 @@ def array_to_numpy( data: ArrayLike | None, n_dims: int | None = None, dtype: ( - type[NumpyBool] | type[NumpyFloat] | type[NumpyInt] | type[NumpyUInt] + type[np.bool_] + | type[np.float32] + | type[np.float64] + | type[np.int32] + | type[np.int64] + | type[np.uint8] + | type[np.uint16] + | type[np.uint32] ) = np.float32, ) -> NDArrayNumber | None: """Converts a given array like object to a numpy array. @@ -92,14 +106,14 @@ def array_to_numpy( If the argument is None, None will be returned. Examples: - >>> convert_to_array([1,2,3]) - >>> # -> array([1,2,3]) - >>> convert_to_array(None) - >>> # -> None - >>> convert_to_array(torch.tensor([1,2,3]).cuda()) - >>> # -> array([1,2,3]) - >>> convert_to_array([1,2,3], n_dims = 2).shape - >>> # -> [1, 3] + >>> convert_to_array([1,2,3]) + >>> # -> array([1,2,3]) + >>> convert_to_array(None) + >>> # -> None + >>> convert_to_array(torch.tensor([1,2,3]).cuda()) + >>> # -> array([1,2,3]) + >>> convert_to_array([1,2,3], n_dims = 2).shape + >>> # -> [1, 3] Args: data (ArrayLike | None): ArrayLike object that should be converted @@ -110,9 +124,8 @@ def array_to_numpy( squeezed or exanded (from the left). If it still does not match, an error is raised. - dtype (type[NumpyBool] | type[NumpyFloat] | type[NumpyInt] | - type[NumpyUInt], optional): Target dtype of the array. Defaults to - np.float32. + dtype (SUPPORTED_DTYPES, optional): Target dtype of the array. Defaults + to np.float32. Raises: ValueError: If the provied array like objects can not be converted @@ -150,61 +163,4 @@ def array_to_numpy( f"have {n_dims} dimensions." ) - # hardcode next type check since mypy can not resolve this correctly - typed_arr: NDArrayNumber = array.astype(dtype) # type: ignore - return typed_arr - - -@overload -def arrays_to_numpy( - *args: ArrayLike, n_dims: int | None, dtype: type[NumpyBool] -) -> tuple[NDArrayBool, ...]: ... - - -@overload -def arrays_to_numpy( - *args: ArrayLike, n_dims: int | None, dtype: type[NumpyFloat] -) -> tuple[NDArrayFloat, ...]: ... - - -@overload -def arrays_to_numpy( - *args: ArrayLike, n_dims: int | None, dtype: type[NumpyInt] -) -> tuple[NDArrayInt, ...]: ... - - -@overload -def arrays_to_numpy( - *args: ArrayLike, n_dims: int | None, dtype: type[NumpyUInt] -) -> tuple[NDArrayUInt, ...]: ... - - -def arrays_to_numpy( - *args: ArrayLike | None, - n_dims: int | None = None, - dtype: ( - type[NumpyBool] | type[NumpyFloat] | type[NumpyInt] | type[NumpyUInt] - ) = np.float32, -) -> tuple[NDArrayNumber | None, ...]: - """Converts a given sequence of optional ArrayLike objects to numpy. - - Args: - args (ArrayLike | None): Provided arguments. - n_dims (int | None, optional): Target number of dimension of the array. - If the provided array does not have this shape, it will be - squeezed or exanded (from the left). If it still does not match, - an error is Raised. - dtype (type[NumpyBool] | type[NumpyFloat] | type[NumpyInt] | - type[NumpyUInt], optional): Target dtype of the array. Defaults to - np.float32. - - Raises: - ValueError: If the provied array like objects can not be converted - with the target dimensions. - - Returns: - tuple[NDArrayNumber | None]: The converted arguments as numpy array. - """ - # Ignore mypy check due to 'Not all union combinations were tried because - # there are too many unions' - return tuple(array_to_numpy(arg, n_dims, dtype) for arg in args) # type: ignore # pylint: disable=line-too-long + return array.astype(dtype) # type: ignore diff --git a/vis4d/common/ckpt.py b/vis4d/common/ckpt.py index fdd5b6e4a..d5e36faa3 100644 --- a/vis4d/common/ckpt.py +++ b/vis4d/common/ckpt.py @@ -205,7 +205,9 @@ def load_from_local( filename = osp.expanduser(filename) if not osp.isfile(filename): raise FileNotFoundError(f"{filename} can not be found.") - checkpoint = torch.load(filename, map_location=map_location) + checkpoint = torch.load( + filename, weights_only=True, map_location=map_location + ) return checkpoint @@ -326,7 +328,7 @@ def load(module: nn.Module, prefix: str = "") -> None: # recursively check parallel module in case that the model has a # complicated structure, e.g., nn.Module(nn.Module(DDP)) if is_module_wrapper(module): - module = module.module + module = module.module # type: ignore local_metadata = ( {} if metadata is None else metadata.get(prefix[:-1], {}) ) diff --git a/vis4d/common/distributed.py b/vis4d/common/distributed.py index 7241c67cd..6511dc444 100644 --- a/vis4d/common/distributed.py +++ b/vis4d/common/distributed.py @@ -11,52 +11,15 @@ from functools import wraps from typing import Any -import cloudpickle import torch import torch.distributed as dist -from torch import nn +from torch import Tensor, nn from torch.distributed import broadcast_object_list from torch.nn.parallel import DataParallel, DistributedDataParallel from vis4d.common import ArgsType, DictStrAny, GenericFunc -class PicklableWrapper: # mypy: disable=line-too-long - """Wrap an object to make it more picklable. - - Note that it uses heavy weight serialization libraries that are slower than - pickle. It's best to use it only on closures (which are usually not - picklable). This is a simplified version of - https://github.com/joblib/joblib/blob/master/joblib/externals/loky/cloudpickle_wrapper.py - """ - - def __init__(self, obj: Any | PicklableWrapper) -> None: # type: ignore - """Creates an instance of the class.""" - while isinstance(obj, PicklableWrapper): - # Wrapping an object twice is no-op - obj = obj._obj - self._obj: Any = obj - - def __reduce__(self) -> tuple[Any, tuple[bytes]]: - """Reduce.""" - s = cloudpickle.dumps(self._obj) - return cloudpickle.loads, (s,) - - def __call__(self, *args: ArgsType, **kwargs: ArgsType) -> Any: - """Call.""" - return self._obj(*args, **kwargs) - - def __getattr__(self, attr: str) -> Any: - """Get attribute. - - Ensure that the wrapped object can be used seamlessly as the previous - object. - """ - if attr not in ["_obj"]: - return getattr(self._obj, attr) - return getattr(self, attr) - - # no coverage for these functions, since we don't unittest distributed setting def get_world_size() -> int: # pragma: no cover """Get the world size (number of processes) of torch.distributed. @@ -128,7 +91,7 @@ def synchronize() -> None: # pragma: no cover dist.barrier(group=dist.group.WORLD, device_ids=[get_local_rank()]) -def broadcast(obj: Any, src: int = 0) -> Any: # pragma: no cover +def broadcast(obj: Any, src: int = 0) -> Any: # type: ignore """Broadcast an object from a source to all processes.""" if not distributed_available(): return obj @@ -140,14 +103,14 @@ def broadcast(obj: Any, src: int = 0) -> Any: # pragma: no cover return obj[0] -def serialize_to_tensor(data: Any) -> torch.Tensor: # pragma: no cover - """Serialize arbitrary picklable data to a torch.Tensor. +def serialize_to_tensor(data: Any) -> Tensor: # type: ignore + """Serialize arbitrary picklable data to a Tensor. Args: data (Any): The data to serialize. Returns: - torch.Tensor: The serialized data as a torch.Tensor. + Tensor: The serialized data as a Tensor. Raises: AssertionError: If the backend of torch.distributed is not gloo or @@ -186,7 +149,7 @@ def rank_zero_only(func: GenericFunc) -> GenericFunc: """ @wraps(func) - def wrapped_fn(*args: ArgsType, **kwargs: ArgsType) -> Any: + def wrapped_fn(*args: ArgsType, **kwargs: ArgsType) -> Any: # type: ignore rank = get_rank() if rank == 0: return func(*args, **kwargs) @@ -196,8 +159,8 @@ def wrapped_fn(*args: ArgsType, **kwargs: ArgsType) -> Any: def pad_to_largest_tensor( - tensor: torch.Tensor, -) -> tuple[list[int], torch.Tensor]: # pragma: no cover + tensor: Tensor, +) -> tuple[list[int], Tensor]: # pragma: no cover """Pad tensor to largest size among the tensors in each process. Args: @@ -251,7 +214,7 @@ def all_gather_object_gpu( # type: ignore tensor_list = [tensor.clone() for _ in range(world_size)] dist.all_gather_object(tensor_list, tensor) # (world_size, N) - if rank_zero_return_only and not rank == 0: + if rank_zero_return_only and rank != 0: return None # decode @@ -264,7 +227,7 @@ def all_gather_object_gpu( # type: ignore def create_tmpdir( - rank: int, tmpdir: None | str = None + rank: int, tmpdir: None | str = None, use_system_tmp: bool = True ) -> str: # pragma: no cover """Create and distribute a temporary directory across all processes.""" if tmpdir is not None: @@ -273,10 +236,10 @@ def create_tmpdir( if rank == 0: # create a temporary directory default_tmpdir = tempfile.gettempdir() - if default_tmpdir is not None: + if default_tmpdir is not None and use_system_tmp: dist_tmpdir = os.path.join(default_tmpdir, ".dist_tmp") else: - dist_tmpdir = ".dist_tmp" + dist_tmpdir = os.path.join("vis4d-workspace", ".dist_tmp") os.makedirs(dist_tmpdir, exist_ok=True) tmpdir = tempfile.mkdtemp(dir=dist_tmpdir) else: @@ -288,13 +251,15 @@ def all_gather_object_cpu( # type: ignore data: Any, tmpdir: None | str = None, rank_zero_return_only: bool = True, + use_system_tmp: bool = False, ) -> list[Any] | None: # pragma: no cover """Share arbitrary picklable data via file system caching. Args: data: any picklable object. tmpdir: Save path for temporary files. If None, safely create tmpdir. - rank_zero_return_only: if results should only be returned on rank 0 + rank_zero_return_only: if results should only be returned on rank 0. + use_system_tmp: if use system tmpdir or not. Returns: list[Any]: list of data gathered from each process. @@ -304,14 +269,14 @@ def all_gather_object_cpu( # type: ignore return [data] # make tmp dir - tmpdir = create_tmpdir(rank, tmpdir) + tmpdir = create_tmpdir(rank, tmpdir, use_system_tmp) # encode & save with open(os.path.join(tmpdir, f"part_{rank}.pkl"), "wb") as f: pickle.dump(data, f) synchronize() - if rank_zero_return_only and not rank == 0: + if rank_zero_return_only and rank != 0: return None # load & decode @@ -330,18 +295,18 @@ def all_gather_object_cpu( # type: ignore return data_list -def reduce_mean(tensor: torch.Tensor) -> torch.Tensor: +def reduce_mean(tensor: Tensor) -> Tensor: """Obtain the mean of tensor on different GPUs.""" if not (dist.is_available() and dist.is_initialized()): return tensor tensor = tensor.clone() - dist.all_reduce(tensor.div_(dist.get_world_size()), op=dist.ReduceOp.SUM) + dist.all_reduce(tensor.div_(get_world_size()), op=dist.ReduceOp.SUM) return tensor -def obj2tensor( +def obj2tensor( # type: ignore pyobj: Any, device: torch.device = torch.device("cuda") -) -> torch.Tensor: +) -> Tensor: """Serialize picklable python object to tensor. Args: @@ -352,11 +317,11 @@ def obj2tensor( return torch.ByteTensor(storage).to(device=device) -def tensor2obj(tensor: torch.Tensor) -> Any: +def tensor2obj(tensor: Tensor) -> Any: # type: ignore """Deserialize tensor to picklable python object. Args: - tensor (torch.Tensor): Tensor to be deserialized. + tensor (Tensor): Tensor to be deserialized. """ return pickle.loads(tensor.cpu().numpy().tobytes()) diff --git a/vis4d/common/typing.py b/vis4d/common/typing.py index 5e78f0c1f..0a23a96bb 100644 --- a/vis4d/common/typing.py +++ b/vis4d/common/typing.py @@ -16,13 +16,6 @@ Tensor, ) -NumpyBool = np.bool_ -NumpyFloat = Union[np.float32, np.float64] -NumpyInt = Union[np.int32, np.int64] -NumpyUInt = Union[ # pylint: disable=invalid-name - np.uint8, np.uint16, np.uint32 -] - NDArrayBool = npt.NDArray[np.bool_] NDArrayF32 = npt.NDArray[np.float32] NDArrayF64 = npt.NDArray[np.float64] @@ -47,7 +40,6 @@ LossesType = Dict[str, Tensor] TorchLossFunc = Callable[..., Any] # type: ignore GenericFunc = Callable[..., Any] # type: ignore -TrainingModule = Any # type: ignore ArrayIterableFloat = Iterable[Union[float, "ArrayIterableFloat"]] ArrayIterableBool = Iterable[Union[bool, "ArrayIterableBool"]] @@ -63,3 +55,20 @@ ArrayLike = Union[ArrayLikeBool, ArrayLikeFloat, ArrayLikeInt, ArrayLikeUInt] ListAny = list[Any] # type: ignore + + +# Trick mypy into not applying contravariance rules to inputs by defining +# forward as a value, rather than a function. See also +# https://github.com/python/mypy/issues/8795 +def unimplemented(self, *args: Any) -> None: # type: ignore + r"""Define the computation performed at every call. + + Should be overridden by all subclasses. + + .. note:: + Although the recipe for forward pass needs to be defined within + this function, one should call the :class:`Module` instance afterwards + instead of this since the former takes care of running the + registered hooks while the latter silently ignores them. + """ + raise NotImplementedError() diff --git a/vis4d/common/util.py b/vis4d/common/util.py index 75d51ce4c..38dca5d7d 100644 --- a/vis4d/common/util.py +++ b/vis4d/common/util.py @@ -73,7 +73,7 @@ def set_tf32(use_tf32: bool, precision: str) -> None: # pragma: no cover def init_random_seed() -> int: """Initialize random seed for the experiment.""" - return np.random.randint(2**31) + return int(np.random.randint(2**31)) def set_random_seed(seed: int, deterministic: bool = False) -> None: diff --git a/vis4d/config/replicator.py b/vis4d/config/replicator.py deleted file mode 100644 index a12c1d40d..000000000 --- a/vis4d/config/replicator.py +++ /dev/null @@ -1,324 +0,0 @@ -"""Replication methods to perform different parameters sweeps.""" - -from __future__ import annotations - -import re -from collections.abc import Callable, Generator, Iterable -from queue import Queue -from typing import Any - -from ml_collections import ConfigDict - -from vis4d.common.typing import ArgsType - - -def iterable_sampler( # type: ignore - samples: Iterable[Any], -) -> Callable[[], Generator[Any, None, None]]: - """Creates a sampler from an iterable. - - This fuction returns a method that returns a generator that iterates - over all values provided in the 'samples' iterable. - - Args: - samples (Iterable[Any]): Iterable over which to sample. - - Returns: - Callable[[], Generator[Any, None, None]]: Function that - returns a generator which iterates over all elements in the given - iterable. - """ - - def _sampler() -> Generator[float, None, None]: - yield from samples - - return _sampler - - -def linspace_sampler( - min_value: float, max_value: float, n_steps: int = 1 -) -> Callable[[], Generator[float, None, None]]: - """Creates a linear space sampler. - - This fuction returns a method that returns a generator that iterates - from min_value to max_value in n_steps. - - Args: - min_value (float): Lower value bound - max_value (float): Upper value bound - n_steps (int, optional): Number of steps. Defaults to 1. - - Returns: - Callable[[], Generator[float, None, None]]: Function that - returns a generator which iterates from min to max in n_steps. - """ - - def _sampler() -> Generator[float, None, None]: - for i in range(n_steps): - yield min_value + i / max(n_steps - 1, 1) * (max_value - min_value) - - return _sampler - - -def logspace_sampler( - min_exponent: float, - max_exponent: float, - n_steps: int = 1, - base: float = 10, -) -> Callable[[], Generator[float, None, None]]: - """Creates a logarithmic space sampler. - - This fuction returns a method that returns a generator that iterates - from base^min_exponent to base^max_exponent in n_steps. - - Args: - min_exponent (float): Lower value bound - max_exponent (float): Upper value bound - n_steps (int, optional): Number of steps. Defaults to 1. - base (float): Base value for exponential calculation. Defaults to 10. - - Returns: - Callable[[], Generator[float, None, None]]: Function that - returns a generator which iterates from 10^min to 10^max in n_steps. - """ - - def _sampler() -> Generator[float, None, None]: - for exp in linspace_sampler(min_exponent, max_exponent, n_steps)(): - yield base**exp - - return _sampler - - -def replicate_config( # type: ignore - configuration: ConfigDict, - sampling_args: list[ - tuple[str, Callable[[], Generator[Any, None, None]] | Iterable[Any]] - ], - method: str = "grid", - fstring="", -) -> Generator[ConfigDict, None, None]: - """Function used to replicate a config. - - This function takes a ConfigDict and a dict with (key: generator) entries. - It will yield, multiple modified config dicts assigned with different - values defined in the sampling_args dictionary. - - Example: - >>> config = ConfigDict({"trainer": {"lr": 0.2, "bs": 2}}) - >>> replicated_config = replicate_config(config, - >>> sampling_args = [("trainer.lr", linspace_sampler(0.01, 0.1, 3))], - >>> method = "grid" - >>> ) - >>> for c in replicated_config: - >>> print(c) - - Will print: - trainer: bs: 2 lr: 0.01 - trainer: bs: 2 lr: 0.055 - trainer: bs: 2 lr: 0.1 - - NOTE, the config dict instance that will be returned will be mutable and - continuously updated to preserve references. - In the code above, executing - >>> print(list(replicated_config)) - Prints: - trainer: bs: 2 lr: 0.1 - trainer: bs: 2 lr: 0.1 - trainer: bs: 2 lr: 0.1 - - Please resolve the reference and copy the dict if you need a list: - >>> print([c.copy_and_resolve_references() for c in replicated_config]) - - - Args: - configuration (ConfigDict): Configuration to replicate - sampling_args (dict[str, Callable[[], Any]]): The queue, - that contains (key, iterator) pairs where the iterator - yields the values which should be assigned to the key. - method (str): What replication method to use. Currently only - 'grid' and 'linear' is supported. - Grid combines the sampling arguments in a grid wise fashion - ([1,2],[3,4] -> [1,3],[1,4],[2,3],[2,4]) whereas 'linear' will - only select elements at the same index ([1,2],[3,4]->[1,3],[2,4]). - fstring (str): Format string to use for the experiment name. Defaults - to an empty string. The format string will be resolved with the - values of the config dict. For example, if the config dict - contains a key 'trainer.lr' with value 0.1, the format string - '{trainer.lr}' will be resolved to '0.1'. - - Raises: - ValueError: if the replication method is unknown. - """ - sampling_queue: Queue[ # type: ignore - tuple[str, Callable[[], Generator[Any, None, None]]] - ] = Queue() - - for key, value in sampling_args: - # Convert Iterable to a callable generator - if isinstance(value, Iterable): - generator: Callable[[], Generator[ArgsType, None, None]] = ( - lambda value=value: (i for i in value) # type: ignore - ) - sampling_queue.put((key, generator)) - else: - sampling_queue.put((key, value)) - - if method == "grid": - replicated = _replicate_config_grid(configuration, sampling_queue) - elif method == "linear": - replicated = _replicate_config_linear(configuration, sampling_queue) - else: - raise ValueError(f"Unknown replication method {method}") - - original_name = configuration.experiment_name - - for config in replicated: - # Update config name - config.experiment_name = ( - f"{original_name}_{_resolve_fstring(fstring, config)}" - ) - yield config - - -def _resolve_fstring(fstring: str, config: ConfigDict) -> str: - """Resolves a format string with the values from the config. - - This function takes a format string and replaces all the keys - with the values from the config. The keys are expected to be - in the format {key} or {key:format}. - This function may fail if the format string contains a key that - is not present in the config. It will also fail if the format - string contains a key that is not a valid python identifier. - - Args: - fstring (str): The format string. E.g. "lr_{params.lr}". - config (ConfigDict): The config dict. E.g. {"params": {"lr": 0.1}}. - - Returns: - str: The resolved format string. E.g. "lr_0.1 - """ - # match everything between { and ':' or '}' - pattern = re.compile(r"{([^:}]+)") - required_params = {p.strip() for p in pattern.findall(fstring)} - - format_dict: dict[str, str] = {} - for param in required_params: - # Maks out '.' which is invalid for .format() call - new_param_name = param.replace(".", "_") - format_dict[new_param_name] = getattr(config, param) - fstring = fstring.replace(param, new_param_name) - - # apply formatting - return fstring.format(**format_dict) - - -def _replicate_config_grid( # type: ignore - configuration: ConfigDict, - sampling_queue: Queue[ - tuple[str, Callable[[], Generator[Any, None, None]]] - ], -) -> Generator[ConfigDict, None, None]: - """Internal function used to replicate a config. - - This function takes a ConfigDict and a queue with (key, generator) entries. - It will then recursively call itself and yield the ConfigDict with - updated values for every key in the sampling_queue. Each key combination - will be yielded exactly once, resulting in prod(len(generator)) entires. - - - For example, a parameter sweep using 'lr: [0,1], bs: [8, 16]' will yield - [0, 8], [0, 16], [0, 8], [1, 16] as combinations. - - Args: - configuration (ConfigDict): Configuration to replicate - sampling_queue (Queue[tuple[str, Callable[[], Any]]]): The queue, - that contains (key, iterator) pairs where the iterator - yields the values which should be assigned to the key. - - Yields: - ConfigDict: Replicated configuration with updated key values. - """ - # Termination criterion reached, We processed all samplers - if sampling_queue.empty(): - yield configuration - return - - # Get next key we want to replicate - (key_name, sampler) = sampling_queue.get() - - # Iterate over all possible assignement values for this key - for value in sampler(): - # Update value ignoring type errors - # (e.g. user set default lr to 1 instead 1.0 would - # otherwise give a type error (float != int)) - with configuration.ignore_type(): - configuration.update_from_flattened_dict({key_name: value}) - - # Let the other samplers process the remaining keys - yield from _replicate_config_grid(configuration, sampling_queue) - - # Add back this sampler for next round - sampling_queue.put((key_name, sampler)) - - -def _replicate_config_linear( # type: ignore - configuration: ConfigDict, - sampling_queue: Queue[ - tuple[str, Callable[[], Generator[Any, None, None]]] - ], - current_position: int | None = None, -) -> Generator[ConfigDict, None, None]: - """Internal function used to replicate a config in a linear way. - - This function takes a ConfigDict and a queue with (key, generator) entries. - It will then recursively call itself and yield the ConfigDict with - updated values for every key in the sampling_queue. - - For example, a parameter sweep using 'lr: [0,1], bs: [8, 16]' will yield - [0, 8], [1, 16] as combinations. - - Args: - configuration (ConfigDict): Configuration to replicate - sampling_queue (Queue[tuple[str, Callable[[], Any]]]): The queue, - that contains (key, iterator) pairs where the iterator - yields the values which should be assigned to the key. - current_position (int, optional): Current position of the top level - sampling module. Used and updated internally. - - Yields: - ConfigDict: Replicated configuration with updated key values. - """ - # Termination criterion reached, We processed all samplers - if sampling_queue.empty(): - yield configuration - return - - # Get next key we want to replicate - (key_name, sampler) = sampling_queue.get() - - is_top_level = False - if current_position is None: - is_top_level = True # This is the top level call. - current_position = 0 - - # Iterate over all possible assignement values for this key - for idx, value in enumerate(sampler()): - if not is_top_level and idx != current_position: - continue # only yield entry that matches requested position - - # Update value ignoring type errors - # (e.g. user set default lr to 1 instead 1.0 would - # otherwise give a type error (float != int)) - with configuration.ignore_type(): - configuration.update_from_flattened_dict({key_name: value}) - - # Let the other samplers process the remaining keys - yield from _replicate_config_linear( - configuration, sampling_queue, current_position - ) - - if is_top_level: - current_position += 1 - - # Add back this sampler for next round - sampling_queue.put((key_name, sampler)) diff --git a/vis4d/data/cbgs.py b/vis4d/data/cbgs.py index 226b25d09..d087ad541 100644 --- a/vis4d/data/cbgs.py +++ b/vis4d/data/cbgs.py @@ -113,7 +113,9 @@ def _get_sample_indices(self) -> list[int]: sample_indices = [] frac = 1.0 / len(self.cat2id) - ratios = [frac / v for v in class_distribution.values()] + ratios = [ + frac / v if v > 0 else 1 for v in class_distribution.values() + ] for cls_inds, ratio in zip( list(class_sample_indices.values()), ratios ): diff --git a/vis4d/data/const.py b/vis4d/data/const.py index 603a0306e..fb26daad3 100644 --- a/vis4d/data/const.py +++ b/vis4d/data/const.py @@ -18,21 +18,21 @@ class AxisMode(Enum): """Enum for choosing among different coordinate frame conventions. ROS: The coordinate frame aligns with the right hand rule: - x axis points forward - y axis points left - z axis points up + - x axis points forward. + - y axis points left. + - z axis points up. See also: https://www.ros.org/reps/rep-0103.html#axis-orientation OpenCV: The coordinate frame aligns with a camera coordinate system: - x axis points right - y axis points down - z axis points forward + - x axis points right. + - y axis points down. + - z axis points forward. See also: https://docs.opencv.org/3.4/d9/d0c/group__calib3d.html LiDAR: The coordinate frame aligns with a LiDAR coordinate system: - x axis points right - y axis points forward - z axis points up + - x axis points right. + - y axis points forward. + - z axis points up. See also: https://www.nuscenes.org/nuscenes#data-collection """ @@ -49,88 +49,131 @@ class CommonKeys: keys where we expect a pre-defined format to enable the usage of common data pre-processing operations among different datasets. - images (NDArrayF32): Image of shape [1, H, W, C]. - input_hw (Tuple[int, int]): Shape of image in (height, width) after - transformations. - original_images (NDArrayF32): Original image of shape [1, H, W, C]. - original_hw (Tuple[int, int]): Original shape of image in (height, width). - - sample_names (str): Name of the current sample. - sequence_names (str): If the dataset contains videos, this field indicates - the name of the current sequence. - frame_ids (int): If the dataset contains videos, this field indicates the - temporal frame index of the current image / sample. - - categories (NDArrayF32): Class labels of shape [C, ]. - - boxes2d (NDArrayF32): 2D bounding boxes of shape [N, 4] in xyxy format. - boxes2d_classes (NDArrayI64): Semantic classes of 2D bounding boxes, shape - [N,]. - boxes2d_track_ids (NDArrayI64): Tracking IDs of 2D bounding boxes, - shape [N,]. - instance_masks (NDArrayUI8): Instance segmentation masks of shape - [N, H, W]. - seg_masks (NDArrayUI8): Semantic segmentation masks [H, W]. - panoptic_masks (NDArrayI64): Panoptic segmentation masks [H, W]. - deph_maps (NDArrayF32): Depth maps of shape [H, W]. - - intrinsics (NDArrayF32): Intrinsic sensor calibration. Shape [3, 3]. - extrinsics (NDArrayF32): Extrinsic sensor calibration, transformation of - sensor to world coordinate frame. Shape [4, 4]. - axis_mode (AxisMode): Coordinate convention of the current sensor. - timestamp (int): Sensor timestamp in Unix format. - - points3d (NDArrayF32): 3D pointcloud data, assumed to be [N, 3] and in - sensor frame. - colors3d (NDArrayF32): Associated color values for each point, [N, 3]. - - semantics3d: TODO complete - instances3d: TODO complete - boxes3d (NDArrayF32): [N, 10], each row consists of center (XYZ), - dimensions (WLH), and orientation quaternion (WXYZ). - boxes3d_classes (NDArrayI64): Associated semantic classes of 3D bounding - boxes, [N,]. + General Info: + - sample_names (str): Name of the sample. + + If the dataset contains videos: + - sequence_names (str): The name of the sequence. + - frame_ids (int): The temporal frame index of the sample. + + Image Based Inputs: + - images (NDArrayF32): Image of shape [1, H, W, C]. + - input_hw (Tuple[int, int]): Shape of image in (height, width) after + transformations. + - original_images (NDArrayF32): Original image of shape [1, H, W, C]. + - original_hw (Tuple[int, int]): Shape of original image in + (height, width). + + Image Classification: + - categories (NDArrayI64): Class labels of shape [1, ]. + + 2D Object Detection: + - boxes2d (NDArrayF32): 2D bounding boxes of shape [N, 4] in xyxy + format. + - boxes2d_classes (NDArrayI64): Classes of 2D bounding boxes of shape + [N,]. + - boxes2d_names (List[str]): Names of 2D bounding box classes, same + order as `boxes2d_classes`. + + 2D Object Tracking: + - boxes2d_track_ids (NDArrayI64): Tracking IDs of 2D bounding boxes of + shape [N,]. + + Segmentation: + - masks (NDArrayUI8): Segmentation masks of shape [N, H, W]. + - seg_masks (NDArrayUI8): Semantic segmentation masks [H, W]. + - instance_masks (NDArrayUI8): Instance segmentation masks of shape + [N, H, W]. + - panoptic_masks (NDArrayI64): Panoptic segmentation masks [H, W]. + + Depth Estimation: + - depth_maps (NDArrayF32): Depth maps of shape [H, W]. + + Optical Flow: + - optical_flows (NDArrayF32): Optical flow maps of shape [H, W, 2]. + + Sensor Calibration: + - intrinsics (NDArrayF32): Intrinsic sensor calibration. Shape [3, 3]. + - extrinsics (NDArrayF32): Extrinsic sensor calibration, transformation + of sensor to world coordinate frame. Shape [4, 4]. + - axis_mode (AxisMode): Coordinate convention of the current sensor. + - timestamp (int): Sensor timestamp in Unix format. + + 3D Point Cloud Data: + - points3d (NDArrayF32): 3D pointcloud data, assumed to be [N, 3] and + in sensor frame. + - colors3d (NDArrayF32): Associated color values for each point [N, 3]. + + 3D Point Cloud Annotations: + - semantics3d (NDArrayI64): Semantic classes of 3D points [N, 1]. + - instances3d (NDArrayI64): Instance IDs of 3D points [N, 1]. + + 3D Object Detection: + - boxes3d (NDArrayF32): 3D bounding boxes of shape [N, 10], each + consists of center (XYZ), dimensions (WLH), and orientation + quaternion (WXYZ). + - boxes3d_classes (NDArrayI64): Associated semantic classes of 3D + bounding boxes of shape [N,]. + - boxes3d_names (List[str]): Names of 3D bounding box classes, same + order as `boxes3d_classes`. + - boxes3d_track_ids (NDArrayI64): Associated tracking IDs of 3D + bounding boxes of shape [N,]. + - boxes3d_velocities (NDArrayF32): Associated velocities of 3D bounding + boxes of shape [N, 3], where each velocity is in the form of + (vx, vy, vz). """ + # General Info + sample_names = "sample_names" + sequence_names = "sequence_names" + frame_ids = "frame_ids" + # image based inputs images = "images" input_hw = "input_hw" original_images = "original_images" original_hw = "original_hw" - # General Info - sample_names = "sample_names" # Sample name for each sample - sequence_names = "sequence_names" # Sequence name for each sample - frame_ids = "frame_ids" + # Image Classification + categories = "categories" - # 2D annotations + # 2D Object Detection boxes2d = "boxes2d" boxes2d_classes = "boxes2d_classes" + boxes2d_names = "boxes2d_names" + + # 2D Object Tracking boxes2d_track_ids = "boxes2d_track_ids" - instance_masks = "instance_masks" + + # Segmentation + masks = "masks" seg_masks = "seg_masks" + instance_masks = "instance_masks" panoptic_masks = "panoptic_masks" + + # Depth Estimation depth_maps = "depth_maps" - optical_flows = "optical_flows" - # Image Classification - categories = "categories" + # Optical Flow + optical_flows = "optical_flows" - # sensor calibration + # Sensor Calibration intrinsics = "intrinsics" extrinsics = "extrinsics" axis_mode = "axis_mode" timestamp = "timestamp" - # 3D data + # 3D Point Cloud Data points3d = "points3d" colors3d = "colors3d" - # 3D annotation + # 3D Point Cloud Annotations semantics3d = "semantics3d" instances3d = "instances3d" + + # 3D Object Detection boxes3d = "boxes3d" boxes3d_classes = "boxes3d_classes" + boxes3d_names = "boxes3d_names" boxes3d_track_ids = "boxes3d_track_ids" boxes3d_velocities = "boxes3d_velocities" - occupancy3d = "occupancy3d" diff --git a/vis4d/data/datasets/coco.py b/vis4d/data/datasets/coco.py index 4e7a963ce..30f183002 100644 --- a/vis4d/data/datasets/coco.py +++ b/vis4d/data/datasets/coco.py @@ -16,7 +16,7 @@ from vis4d.data.typing import DictData from .base import Dataset -from .util import CacheMappingMixin, im_decode +from .util import CacheMappingMixin, get_category_names, im_decode # COCO detection coco_det_map = { @@ -204,6 +204,11 @@ def __init__( cached_file_path=cached_file_path, ) + if self.use_pascal_voc_cats: + self.category_names = get_category_names(coco_seg_map) + else: + self.category_names = get_category_names(coco_det_map) + def __repr__(self) -> str: """Concise representation of the dataset.""" return ( @@ -355,4 +360,6 @@ def __getitem__(self, idx: int) -> DictData: seg_masks[mask_tensor.sum(0) > 1] = 255 # discard overlapped dict_data[K.seg_masks] = seg_masks[None] + dict_data[K.boxes2d_names] = self.category_names + return dict_data diff --git a/vis4d/data/datasets/scalabel.py b/vis4d/data/datasets/scalabel.py index 215c07ffd..746cd2be8 100644 --- a/vis4d/data/datasets/scalabel.py +++ b/vis4d/data/datasets/scalabel.py @@ -41,7 +41,9 @@ poly2ds_to_mask, rle_to_mask, ) - from scalabel.label.typing import Config + from scalabel.label.typing import ( + Config, + ) from scalabel.label.typing import Dataset as ScalabelData from scalabel.label.typing import ( Extrinsics, diff --git a/vis4d/data/datasets/util.py b/vis4d/data/datasets/util.py index b4ffbcb97..cdb232259 100644 --- a/vis4d/data/datasets/util.py +++ b/vis4d/data/datasets/util.py @@ -54,8 +54,9 @@ def im_decode( "L", }, f"{mode} not supported for image decoding!" if backend == "PIL": - pil_img = Image.open(BytesIO(bytearray(im_bytes))) - pil_img = ImageOps.exif_transpose(pil_img) # type: ignore + pil_img_file = Image.open(BytesIO(bytearray(im_bytes))) + pil_img = ImageOps.exif_transpose(pil_img_file) + assert pil_img is not None, "Image could not be loaded!" if pil_img.mode == "L": # pragma: no cover if mode == "L": img: NDArrayUI8 = np.array(pil_img)[..., None] @@ -359,3 +360,8 @@ def short_name(name: str) -> str: f"Distribution of instances among all {num_classes} categories:\n" + colored(table, "cyan") ) + + +def get_category_names(det_mapping: dict[str, int]) -> list[str]: + """Get category names from a mapping of category names to ids.""" + return sorted(det_mapping, key=det_mapping.get) # type: ignore diff --git a/vis4d/data/io/hdf5.py b/vis4d/data/io/hdf5.py index 151077700..c7eee55a7 100644 --- a/vis4d/data/io/hdf5.py +++ b/vis4d/data/io/hdf5.py @@ -6,12 +6,10 @@ from __future__ import annotations -import argparse import os from typing import Literal import numpy as np -from tqdm import tqdm from vis4d.common.imports import H5PY_AVAILABLE @@ -139,13 +137,13 @@ def _get_client(self, hdf5_path: str, mode: str) -> File: File: the hdf5 file. """ if hdf5_path not in self.db_cache: - client = File(hdf5_path, mode) + client = File(hdf5_path, mode, swmr=True, libver="latest") self.db_cache[hdf5_path] = [client, mode] else: client, current_mode = self.db_cache[hdf5_path] if current_mode != mode: client.close() - client = File(hdf5_path, mode) + client = File(hdf5_path, mode, swmr=True, libver="latest") self.db_cache[hdf5_path] = [client, mode] return client @@ -242,63 +240,3 @@ def close(self) -> None: for client, _ in self.db_cache.values(): client.close() self.db_cache.clear() - - -def convert_dataset(source_dir: str) -> None: - """Convert a dataset to HDF5 format. - - This function converts an arbitary dictionary to an HDF5 file. The keys - inside the HDF5 file preserve the directory structure of the original. - - As an example, if you convert "/path/to/dataset" to HDF5, the resulting - file will be: "/path/to/dataset.hdf5". The file "relative/path/to/file" - will be stored at "relative/path/to/file" inside /path/to/dataset.hdf5. - - Args: - source_dir (str): The path to the dataset to convert. - """ - if not os.path.exists(source_dir): - raise FileNotFoundError(f"No such file or directory: {source_dir}") - - source_dir = os.path.join(source_dir, "") # must end with trailing slash - hdf5_path = source_dir.rstrip("/") + ".hdf5" - if os.path.exists(hdf5_path): - print(f"File {hdf5_path} already exists! Skipping {source_dir}") - return - - print(f"Converting dataset at: {source_dir}") - hdf5_file = h5py.File(hdf5_path, mode="w") - sub_dirs = list(os.walk(source_dir)) - file_count = sum(len(files) for (_, _, files) in sub_dirs) - - with tqdm(total=file_count) as pbar: - for root, _, files in sub_dirs: - g_name = root.replace(source_dir, "") - g = hdf5_file.create_group(g_name) if g_name else hdf5_file - for f in files: - filepath = os.path.join(root, f) - if os.path.isfile(filepath): - with open(filepath, "rb") as fp: - file_content = fp.read() - g.create_dataset( - f, data=np.frombuffer(file_content, dtype="uint8") - ) - pbar.update() - - hdf5_file.close() - print("done.") - - -if __name__ == "__main__": # pragma: no cover - parser = argparse.ArgumentParser( - description="Converts a dataset at the specified path to hdf5. The " - "local directory structure is preserved in the hdf5 file." - ) - parser.add_argument( - "-p", - "--path", - required=True, - help="path to the root folder of a specific dataset to convert", - ) - args = parser.parse_args() - convert_dataset(args.path) diff --git a/vis4d/data/io/to_hdf5.py b/vis4d/data/io/to_hdf5.py new file mode 100644 index 000000000..4a2161a5b --- /dev/null +++ b/vis4d/data/io/to_hdf5.py @@ -0,0 +1,76 @@ +"""Script to convert a dataset to hdf5 format.""" + +from __future__ import annotations + +import argparse +import os + +import numpy as np +from tqdm import tqdm + +from vis4d.common.imports import H5PY_AVAILABLE + +if H5PY_AVAILABLE: + import h5py +else: + raise ImportError("Please install h5py to enable HDF5Backend.") + + +def convert_dataset(source_dir: str) -> None: + """Convert a dataset to HDF5 format. + + This function converts an arbitary dictionary to an HDF5 file. The keys + inside the HDF5 file preserve the directory structure of the original. + + As an example, if you convert "/path/to/dataset" to HDF5, the resulting + file will be: "/path/to/dataset.hdf5". The file "relative/path/to/file" + will be stored at "relative/path/to/file" inside /path/to/dataset.hdf5. + + Args: + source_dir (str): The path to the dataset to convert. + """ + if not os.path.exists(source_dir): + raise FileNotFoundError(f"No such file or directory: {source_dir}") + + source_dir = os.path.join(source_dir, "") # must end with trailing slash + hdf5_path = source_dir.rstrip("/") + ".hdf5" + if os.path.exists(hdf5_path): + print(f"File {hdf5_path} already exists! Skipping {source_dir}") + return + + print(f"Converting dataset at: {source_dir}") + hdf5_file = h5py.File(hdf5_path, mode="w") + sub_dirs = list(os.walk(source_dir)) + file_count = sum(len(files) for (_, _, files) in sub_dirs) + + with tqdm(total=file_count) as pbar: + for root, _, files in sub_dirs: + g_name = root.replace(source_dir, "") + g = hdf5_file.create_group(g_name) if g_name else hdf5_file + for f in files: + filepath = os.path.join(root, f) + if os.path.isfile(filepath): + with open(filepath, "rb") as fp: + file_content = fp.read() + g.create_dataset( + f, data=np.frombuffer(file_content, dtype="uint8") + ) + pbar.update() + + hdf5_file.close() + print("done.") + + +if __name__ == "__main__": # pragma: no cover + parser = argparse.ArgumentParser( + description="Converts a dataset at the specified path to hdf5. The " + "local directory structure is preserved in the hdf5 file." + ) + parser.add_argument( + "-p", + "--path", + required=True, + help="path to the root folder of a specific dataset to convert", + ) + args = parser.parse_args() + convert_dataset(args.path) diff --git a/vis4d/data/loader.py b/vis4d/data/loader.py index 482697e55..44bf1f3ca 100644 --- a/vis4d/data/loader.py +++ b/vis4d/data/loader.py @@ -8,15 +8,21 @@ import numpy as np import torch -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import ( + DataLoader, + Dataset, + RandomSampler, + SequentialSampler, +) from torch.utils.data.distributed import DistributedSampler +from torch.utils.data.sampler import Sampler from vis4d.common.distributed import get_rank, get_world_size from .const import CommonKeys as K from .data_pipe import DataPipe from .datasets import VideoDataset -from .samplers import VideoInferenceSampler +from .samplers import AspectRatioBatchSampler, VideoInferenceSampler from .transforms import compose from .transforms.to_tensor import ToTensor from .typing import DictData, DictDataOrList @@ -121,8 +127,11 @@ def build_train_dataloader( collate_keys: Sequence[str] = DEFAULT_COLLATE_KEYS, sensors: Sequence[str] | None = None, pin_memory: bool = True, - shuffle: bool = True, + shuffle: bool | None = True, + drop_last: bool = False, seed: int | None = None, + aspect_ratio_grouping: bool = False, + sampler: Sampler | None = None, # type: ignore disable_subprocess_warning: bool = False, ) -> DataLoader[DictDataOrList]: """Build training dataloader.""" @@ -164,10 +173,31 @@ def _worker_init_fn(worker_id: int) -> None: if disable_subprocess_warning and worker_id != 0: warnings.simplefilter("ignore") - sampler = None - if get_world_size() > 1: - sampler = DistributedSampler(dataset, shuffle=shuffle) - shuffle = False + if sampler is None: + if get_world_size() > 1: + assert isinstance( + shuffle, bool + ), "When using distributed training, shuffle must be a boolean." + sampler = DistributedSampler( + dataset, shuffle=shuffle, drop_last=drop_last + ) + shuffle = False + drop_last = False + elif shuffle: + sampler = RandomSampler(dataset) + shuffle = False + else: + sampler = SequentialSampler(dataset) + + batch_sampler = None + if aspect_ratio_grouping: + batch_sampler = AspectRatioBatchSampler( + sampler, batch_size=samples_per_gpu, drop_last=drop_last + ) + samples_per_gpu = 1 + shuffle = None + drop_last = False + sampler = None dataloader = DataLoader( dataset, @@ -177,10 +207,12 @@ def _worker_init_fn(worker_id: int) -> None: _collate_fn_multi if dataset.has_reference else _collate_fn_single ), sampler=sampler, + batch_sampler=batch_sampler, worker_init_fn=_worker_init_fn, persistent_workers=workers_per_gpu > 0, pin_memory=pin_memory, shuffle=shuffle, + drop_last=drop_last, ) return dataloader @@ -216,17 +248,17 @@ def _collate_fn(data: list[DictData]) -> DictData: dataloaders = [] for dataset in datasets_: - if isinstance(dataset, DataPipe): - assert ( - len(dataset.datasets) == 1 - ), "Inference needs a single dataset per DataPipe." - current_dataset = dataset.datasets[0] - else: - current_dataset = dataset - sampler: DistributedSampler[list[int]] | None if get_world_size() > 1: if video_based_inference: + if isinstance(dataset, DataPipe): + assert ( + len(dataset.datasets) == 1 + ), "DDP Vdieo Inference only support a single dataset." + current_dataset = dataset.datasets[0] + else: + current_dataset = dataset + assert isinstance( current_dataset, VideoDataset ), "Video based inference needs a VideoDataset." diff --git a/vis4d/data/samplers.py b/vis4d/data/samplers.py index ae3d00a10..b7401fcc2 100644 --- a/vis4d/data/samplers.py +++ b/vis4d/data/samplers.py @@ -7,6 +7,9 @@ import numpy as np from torch.utils.data import Dataset from torch.utils.data.distributed import DistributedSampler +from torch.utils.data.sampler import BatchSampler, Sampler + +from vis4d.data.const import CommonKeys as K from .datasets.base import VideoDataset from .typing import DictDataOrList @@ -76,3 +79,78 @@ def __iter__(self) -> Iterator[list[int]]: def __len__(self) -> int: """Return length of sampler instance.""" return len(self._local_idcs) + + +class AspectRatioBatchSampler(BatchSampler): + """A sampler wrapper for grouping images with similar aspect ratio. + + Moidified from: + https://github.com/open-mmlab/mmdetection/blob/main/mmdet/datasets/samplers/batch_sampler.py + + Args: + sampler (Sampler): Base sampler. + batch_size (int): Size of mini-batch. + drop_last (bool): If ``True``, the sampler will drop the last batch if + its size would be less than ``batch_size``. + """ + + def __init__( + self, + sampler: Sampler, # type: ignore + batch_size: int, + drop_last: bool = False, + ) -> None: + """Creates an instance of the class.""" + if not isinstance(sampler, Sampler): + raise TypeError( + "sampler should be an instance of ``Sampler``, " + f"but got {sampler}" + ) + + super().__init__(sampler, batch_size, drop_last) + + # two groups for w < h and w >= h + self._aspect_ratio_buckets: list[list[int]] = [[] for _ in range(2)] + + def __iter__(self) -> Iterator[list[int]]: + """Iteration method.""" + for idx in self.sampler: + if hasattr(self.sampler, "dataset"): + data_dict = self.sampler.dataset[idx] + elif hasattr(self.sampler, "data_source"): + data_dict = self.sampler.data_source[idx] + else: + raise ValueError( + "sampler should have dataset or data_source attribute" + ) + height, width = data_dict[K.input_hw] + bucket_id = 0 if width < height else 1 + bucket = self._aspect_ratio_buckets[bucket_id] + bucket.append(idx) + # yield a batch of indices in the same aspect ratio group + if len(bucket) == self.batch_size: + yield bucket[:] + del bucket[:] + + # yield the rest data and reset the bucket + left_data = ( + self._aspect_ratio_buckets[0] + self._aspect_ratio_buckets[1] + ) + self._aspect_ratio_buckets = [[] for _ in range(2)] + while len(left_data) > 0: + if len(left_data) <= self.batch_size: + if not self.drop_last: + yield left_data[:] + left_data = [] + else: + yield left_data[: self.batch_size] + left_data = left_data[self.batch_size :] + + def __len__(self) -> int: + """Return length of sampler instance.""" + if self.drop_last: + return len(self.sampler) // self.batch_size # type: ignore + + return ( + len(self.sampler) + self.batch_size - 1 # type: ignore + ) // self.batch_size diff --git a/vis4d/data/transforms/autoaugment.py b/vis4d/data/transforms/autoaugment.py index f0bbd7812..6e47ad9f9 100644 --- a/vis4d/data/transforms/autoaugment.py +++ b/vis4d/data/transforms/autoaugment.py @@ -1,5 +1,7 @@ """A wrap for timm transforms.""" +from __future__ import annotations + from typing import Union import numpy as np @@ -45,7 +47,7 @@ class _AutoAug: """Apply Timm's AutoAugment to a image array.""" def __init__(self) -> None: - self.aug_op = None + self.aug_op: AugOp | None = None def _create(self, policy: str, hparams: dict[str, float]) -> AugOp: """Create augmentation op.""" diff --git a/vis4d/data/transforms/photometric.py b/vis4d/data/transforms/photometric.py index 0d19eb9a1..a0abd454e 100644 --- a/vis4d/data/transforms/photometric.py +++ b/vis4d/data/transforms/photometric.py @@ -330,9 +330,9 @@ def __call__(self, images: list[NDArrayF32]) -> list[NDArrayF32]: for i, image in enumerate(images): image = image[0].astype(np.uint8) if self.image_channel_mode == "BGR": - image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # type: ignore + image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) else: - image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV) # type: ignore + image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV) image = image.astype(np.int16) hsv_gains = np.random.uniform(-1, 1, 3) * [ self.hue_delta, diff --git a/vis4d/data/transforms/resize.py b/vis4d/data/transforms/resize.py index c613bc0e1..c829f7a09 100644 --- a/vis4d/data/transforms/resize.py +++ b/vis4d/data/transforms/resize.py @@ -50,6 +50,7 @@ def __init__( align_long_edge: bool = False, resize_short_edge: bool = False, allow_overflow: bool = False, + fixed_scale: bool = False, ) -> None: """Creates an instance of the class. @@ -78,14 +79,63 @@ def __init__( to the smallest size such that it is no smaller than shape. Otherwise, we scale the image to the largest size such that it is no larger than shape. Defaults to False. + fixed_scale (bool, optional): If set to True, we scale the image + without offset. Defaults to False. """ self.shape = shape self.keep_ratio = keep_ratio + + assert multiscale_mode in {"list", "range"} self.multiscale_mode = multiscale_mode + + assert ( + scale_range[0] <= scale_range[1] + ), f"Invalid scale range: {scale_range[1]} < {scale_range[0]}" self.scale_range = scale_range + self.align_long_edge = align_long_edge self.resize_short_edge = resize_short_edge self.allow_overflow = allow_overflow + self.fixed_scale = fixed_scale + + def _get_target_shape( + self, input_shape: tuple[int, int] + ) -> tuple[int, int]: + """Generate possibly random target shape.""" + if self.multiscale_mode == "range": + assert isinstance( + self.shape, tuple + ), "Specify shape as tuple when using multiscale mode range." + if self.scale_range[0] < self.scale_range[1]: # do multi-scale + w_scale = ( + random.uniform(0, 1) + * (self.scale_range[1] - self.scale_range[0]) + + self.scale_range[0] + ) + h_scale = ( + random.uniform(0, 1) + * (self.scale_range[1] - self.scale_range[0]) + + self.scale_range[0] + ) + else: + h_scale = w_scale = 1.0 + + shape = int(self.shape[0] * h_scale), int(self.shape[1] * w_scale) + else: + assert isinstance( + self.shape, list + ), "Specify shape as list when using multiscale mode list." + shape = random.choice(self.shape) + + return get_resize_shape( + input_shape, + shape, + self.keep_ratio, + self.align_long_edge, + self.resize_short_edge, + self.allow_overflow, + self.fixed_scale, + ) def __call__( self, images: list[NDArrayF32] @@ -94,16 +144,7 @@ def __call__( image = images[0] im_shape = (image.shape[1], image.shape[2]) - target_shape = get_target_shape( - im_shape, - self.shape, - self.keep_ratio, - self.multiscale_mode, - self.scale_range, - self.align_long_edge, - self.resize_short_edge, - self.allow_overflow, - ) + target_shape = self._get_target_shape(im_shape) scale_factor = ( target_shape[1] / im_shape[1], target_shape[0] / im_shape[0], @@ -117,6 +158,66 @@ def __call__( return resize_params, target_shapes +def get_resize_shape( + original_shape: tuple[int, int], + new_shape: tuple[int, int], + keep_ratio: bool = True, + align_long_edge: bool = False, + resize_short_edge: bool = False, + allow_overflow: bool = False, + fixed_scale: bool = False, +) -> tuple[int, int]: + """Get shape for resize, considering keep_ratio and align_long_edge. + + Args: + original_shape (tuple[int, int]): Original shape in [H, W]. + new_shape (tuple[int, int]): New shape in [H, W]. + keep_ratio (bool, optional): Whether to keep the aspect ratio. + Defaults to True. + align_long_edge (bool, optional): Whether to align the long edge of + the original shape with the long edge of the new shape. + Defaults to False. + resize_short_edge (bool, optional): Whether to resize according to the + short edge. Defaults to False. + allow_overflow (bool, optional): Whether to allow overflow. + Defaults to False. + fixed_scale (bool, optional): Whether to use fixed scale. + + Returns: + tuple[int, int]: The new shape in [H, W]. + """ + h, w = original_shape + new_h, new_w = new_shape + + if keep_ratio: + if allow_overflow: + comp_fn = max + else: + comp_fn = min + + if align_long_edge: + long_edge, short_edge = max(new_shape), min(new_shape) + scale_factor = comp_fn( + long_edge / max(h, w), short_edge / min(h, w) + ) + elif resize_short_edge: + short_edge = min(original_shape) + new_short_edge = min(new_shape) + scale_factor = new_short_edge / short_edge + else: + scale_factor = comp_fn(new_w / w, new_h / h) + + if fixed_scale: + offset = 0.0 + else: + offset = 0.5 + + new_h = int(h * scale_factor + offset) + new_w = int(w * scale_factor + offset) + + return new_h, new_w + + @Transform([K.images, "transforms.resize.target_shape"], K.images) class ResizeImages: """Resize Images.""" @@ -166,6 +267,36 @@ def __call__( return images +def resize_image( + inputs: NDArrayF32, + shape: tuple[int, int], + interpolation: str = "bilinear", + antialias: bool = False, + backend: str = "torch", +) -> NDArrayF32: + """Resize image.""" + if backend == "torch": + image = torch.from_numpy(inputs).permute(0, 3, 1, 2) + image = resize_tensor(image, shape, interpolation, antialias) + return image.permute(0, 2, 3, 1).numpy() + + if backend == "cv2": + cv2_interp_codes = { + "nearest": INTER_NEAREST, + "bilinear": INTER_LINEAR, + "bicubic": INTER_CUBIC, + "area": INTER_AREA, + "lanczos": INTER_LANCZOS4, + } + return cv2.resize( # pylint: disable=no-member, unsubscriptable-object + inputs[0].astype(np.uint8), + (shape[1], shape[0]), + interpolation=cv2_interp_codes[interpolation], + )[None, ...].astype(np.float32) + + raise ValueError(f"Invalid imresize backend: {backend}") + + @Transform([K.boxes2d, "transforms.resize.scale_factor"], K.boxes2d) class ResizeBoxes2D: """Resize list of 2D bounding boxes.""" @@ -309,7 +440,7 @@ def __call__( optical_flow_[:, :, 0] *= scale_factor[0] optical_flow_[:, :, 1] *= scale_factor[1] optical_flows[i] = optical_flow_.numpy() - return optical_flow_.numpy() + return optical_flows @Transform( @@ -389,34 +520,6 @@ def __call__( return intrinsics -def resize_image( - inputs: NDArrayF32, - shape: tuple[int, int], - interpolation: str = "bilinear", - antialias: bool = False, - backend: str = "torch", -) -> NDArrayF32: - """Resize image.""" - if backend == "torch": - image = torch.from_numpy(inputs).permute(0, 3, 1, 2) - image = resize_tensor(image, shape, interpolation, antialias) - return image.permute(0, 2, 3, 1).numpy() - if backend == "cv2": - cv2_interp_codes = { - "nearest": INTER_NEAREST, - "bilinear": INTER_LINEAR, - "bicubic": INTER_CUBIC, - "area": INTER_AREA, - "lanczos": INTER_LANCZOS4, - } - return cv2.resize( # pylint: disable=no-member, unsubscriptable-object - inputs[0].astype(np.uint8), - (shape[1], shape[0]), - interpolation=cv2_interp_codes[interpolation], - )[None, ...].astype(np.float32) - raise ValueError(f"Invalid imresize backend: {backend}") - - def resize_tensor( inputs: Tensor, shape: tuple[int, int], @@ -434,107 +537,3 @@ def resize_tensor( antialias=antialias, ) return output - - -def get_resize_shape( - original_shape: tuple[int, int], - new_shape: tuple[int, int], - keep_ratio: bool = True, - align_long_edge: bool = False, - resize_short_edge: bool = False, - allow_overflow: bool = False, -) -> tuple[int, int]: - """Get shape for resize, considering keep_ratio and align_long_edge. - - Args: - original_shape (tuple[int, int]): Original shape in [H, W]. - new_shape (tuple[int, int]): New shape in [H, W]. - keep_ratio (bool, optional): Whether to keep the aspect ratio. - Defaults to True. - align_long_edge (bool, optional): Whether to align the long edge of - the original shape with the long edge of the new shape. - Defaults to False. - resize_short_edge (bool, optional): Whether to resize according to the - short edge. Defaults to False. - allow_overflow (bool, optional): Whether to allow overflow. - Defaults to False. - - Returns: - tuple[int, int]: The new shape in [H, W]. - """ - h, w = original_shape - new_h, new_w = new_shape - if keep_ratio: - if allow_overflow: - comp_fn = max - else: - comp_fn = min - if align_long_edge: - long_edge, short_edge = max(new_shape), min(new_shape) - scale_factor = comp_fn( - long_edge / max(h, w), short_edge / min(h, w) - ) - elif resize_short_edge: - short_edge = min(original_shape) - new_short_edge = min(new_shape) - scale_factor = new_short_edge / short_edge - else: - scale_factor = comp_fn(new_w / w, new_h / h) - new_h = int(h * scale_factor + 0.5) - new_w = int(w * scale_factor + 0.5) - return new_h, new_w - - -def get_target_shape( - input_shape: tuple[int, int], - shape: tuple[int, int] | list[tuple[int, int]], - keep_ratio: bool = False, - multiscale_mode: str = "range", - scale_range: tuple[float, float] = (1.0, 1.0), - align_long_edge: bool = False, - resize_short_edge: bool = False, - allow_overflow: bool = False, -) -> tuple[int, int]: - """Generate possibly random target shape.""" - assert multiscale_mode in {"list", "range"} - if multiscale_mode == "list": - assert isinstance( - shape, list - ), "Specify shape as list when using multiscale mode list." - assert len(shape) >= 1 - else: - assert isinstance( - shape, tuple - ), "Specify shape as tuple when using multiscale mode range." - assert ( - scale_range[0] <= scale_range[1] - ), f"Invalid scale range: {scale_range[1]} < {scale_range[0]}" - - if multiscale_mode == "range": - assert isinstance(shape, tuple) - if scale_range[0] < scale_range[1]: # do multi-scale - w_scale = ( - random.uniform(0, 1) * (scale_range[1] - scale_range[0]) - + scale_range[0] - ) - h_scale = ( - random.uniform(0, 1) * (scale_range[1] - scale_range[0]) - + scale_range[0] - ) - else: - h_scale = w_scale = 1.0 - - shape = int(shape[0] * h_scale), int(shape[1] * w_scale) - else: - assert isinstance(shape, list) - shape = random.choice(shape) - - shape = get_resize_shape( - input_shape, - shape, - keep_ratio, - align_long_edge, - resize_short_edge, - allow_overflow, - ) - return shape diff --git a/vis4d/engine/callbacks/__init__.py b/vis4d/engine/callbacks/__init__.py index 9cfea9a67..e95dc07ee 100644 --- a/vis4d/engine/callbacks/__init__.py +++ b/vis4d/engine/callbacks/__init__.py @@ -1,11 +1,10 @@ """Callback modules.""" from .base import Callback -from .checkpoint import CheckpointCallback from .ema import EMACallback from .evaluator import EvaluatorCallback from .logging import LoggingCallback -from .trainer_state import TrainerState +from .scheduler import LRSchedulerCallback from .visualizer import VisualizerCallback from .yolox_callbacks import ( YOLOXModeSwitchCallback, @@ -15,12 +14,11 @@ __all__ = [ "Callback", - "CheckpointCallback", "EMACallback", "EvaluatorCallback", "LoggingCallback", - "TrainerState", "VisualizerCallback", + "LRSchedulerCallback", "YOLOXModeSwitchCallback", "YOLOXSyncNormCallback", "YOLOXSyncRandomResizeCallback", diff --git a/vis4d/engine/callbacks/base.py b/vis4d/engine/callbacks/base.py index 67e3ae1d4..0ce9cea86 100644 --- a/vis4d/engine/callbacks/base.py +++ b/vis4d/engine/callbacks/base.py @@ -2,17 +2,15 @@ from __future__ import annotations -from torch import Tensor, nn +import lightning.pytorch as pl +from torch import Tensor -from vis4d.common.typing import DictStrArrNested, MetricLogs +from vis4d.common.typing import DictStrArrNested from vis4d.data.typing import DictData from vis4d.engine.connectors import CallbackConnector -from vis4d.engine.loss_module import LossModule -from .trainer_state import TrainerState - -class Callback: +class Callback(pl.Callback): """Base class for Callbacks.""" def __init__( @@ -37,7 +35,9 @@ def __init__( self.train_connector = train_connector self.test_connector = test_connector - def setup(self) -> None: + def setup( + self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str + ) -> None: """Setup callback.""" def get_train_callback_inputs( @@ -83,110 +83,3 @@ def get_test_callback_inputs( assert self.test_connector is not None, "Test connector is None." return self.test_connector(outputs, batch) - - def on_train_batch_start( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - batch: DictData, - batch_idx: int, - ) -> None: - """Hook to run at the start of a training batch. - - Args: - trainer_state (TrainerState): Trainer state. - model: Model that is being trained. - loss_module (LossModule): Loss module. - batch (DictData): Dataloader output data batch. - batch_idx (int): Index of the batch. - """ - - def on_train_epoch_start( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - ) -> None: - """Hook to run at the beginning of a training epoch. - - Args: - trainer_state (TrainerState): Trainer state. - model (nn.Module): Model that is being trained. - loss_module (LossModule): Loss module. - """ - - def on_train_batch_end( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - outputs: DictData, - batch: DictData, - batch_idx: int, - ) -> None | MetricLogs: - """Hook to run at the end of a training batch. - - Args: - trainer_state (TrainerState): Trainer state. - model: Model that is being trained. - loss_module (LossModule): Loss module. - outputs (DictData): Model prediction output. - batch (DictData): Dataloader output data batch. - batch_idx (int): Index of the batch. - """ - - def on_train_epoch_end( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - ) -> None: - """Hook to run at the end of a training epoch. - - Args: - trainer_state (TrainerState): Trainer state. - model (nn.Module): Model that is being trained. - loss_module (LossModule): Loss module. - """ - - def on_test_epoch_start( - self, trainer_state: TrainerState, model: nn.Module - ) -> None: - """Hook to run at the beginning of a testing epoch. - - Args: - trainer_state (TrainerState): Trainer state. - model (nn.Module): Model that is being trained. - """ - - def on_test_batch_end( - self, - trainer_state: TrainerState, - model: nn.Module, - outputs: DictData, - batch: DictData, - batch_idx: int, - dataloader_idx: int = 0, - ) -> None: - """Hook to run at the end of a testing batch. - - Args: - trainer_state (TrainerState): Trainer state. - model: Model that is being trained. - outputs (DictData): Model prediction output. - batch (DictData): Dataloader output data batch. - batch_idx (int): Index of the batch. - dataloader_idx (int, optional): Index of the dataloader. Defaults - to 0. - """ - - def on_test_epoch_end( - self, trainer_state: TrainerState, model: nn.Module - ) -> None | MetricLogs: - """Hook to run at the end of a testing epoch. - - Args: - trainer_state (TrainerState): Trainer state. - model (nn.Module): Model that is being trained. - """ diff --git a/vis4d/engine/callbacks/checkpoint.py b/vis4d/engine/callbacks/checkpoint.py deleted file mode 100644 index 9d31ce1f1..000000000 --- a/vis4d/engine/callbacks/checkpoint.py +++ /dev/null @@ -1,106 +0,0 @@ -"""This module contains utilities for callbacks.""" - -from __future__ import annotations - -import os - -import torch -from torch import nn - -from vis4d.common import ArgsType -from vis4d.common.distributed import broadcast, rank_zero_only -from vis4d.data.typing import DictData -from vis4d.engine.callbacks.trainer_state import TrainerState -from vis4d.engine.loss_module import LossModule - -from .base import Callback -from .trainer_state import TrainerState - - -class CheckpointCallback(Callback): - """Callback for model checkpointing.""" - - def __init__( - self, - *args: ArgsType, - save_prefix: str, - checkpoint_period: int = 1, - **kwargs: ArgsType, - ) -> None: - """Init callback. - - Args: - save_prefix (str): Prefix of checkpoint path for saving. - checkpoint_period (int, optional): Checkpoint period. Defaults to - 1. - """ - super().__init__(*args, **kwargs) - self.output_dir = f"{save_prefix}/checkpoints" - self.checkpoint_period = checkpoint_period - - def setup(self) -> None: # pragma: no cover - """Setup callback.""" - self.output_dir = broadcast(self.output_dir) - os.makedirs(self.output_dir, exist_ok=True) - - @rank_zero_only - def _save_checkpoint( - self, trainer_state: TrainerState, model: nn.Module - ) -> None: - """Save checkpoint.""" - epoch = trainer_state["current_epoch"] - step = trainer_state["global_step"] - ckpt_dict = { - "epoch": epoch, - "global_step": step, - "state_dict": model.state_dict(), - } - - if "optimizers" in trainer_state: - ckpt_dict["optimizers"] = [ - optimizer.state_dict() - for optimizer in trainer_state["optimizers"] - ] - - if "lr_schedulers" in trainer_state: - ckpt_dict["lr_schedulers"] = [ - lr_scheduler.state_dict() - for lr_scheduler in trainer_state["lr_schedulers"] - ] - - torch.save( - ckpt_dict, - f"{self.output_dir}/epoch={epoch}-step={step}.ckpt", - ) - - torch.save(ckpt_dict, f"{self.output_dir}/last.ckpt") - - def on_train_batch_end( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - outputs: DictData, - batch: DictData, - batch_idx: int, - ) -> None: - """Hook to run at the end of a training batch.""" - if ( - not self.epoch_based - and trainer_state["global_step"] % self.checkpoint_period == 0 - ): - self._save_checkpoint(trainer_state, model) - - def on_train_epoch_end( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - ) -> None: - """Hook to run at the end of a training epoch.""" - if ( - self.epoch_based - and (trainer_state["current_epoch"] + 1) % self.checkpoint_period - == 0 - ): - self._save_checkpoint(trainer_state, model) diff --git a/vis4d/engine/callbacks/ema.py b/vis4d/engine/callbacks/ema.py index f876d14cf..fb732fd99 100644 --- a/vis4d/engine/callbacks/ema.py +++ b/vis4d/engine/callbacks/ema.py @@ -2,38 +2,38 @@ from __future__ import annotations -from torch import nn +import lightning.pytorch as pl from vis4d.common.distributed import is_module_wrapper -from vis4d.common.typing import MetricLogs from vis4d.data.typing import DictData -from vis4d.engine.loss_module import LossModule from vis4d.model.adapter import ModelEMAAdapter from .base import Callback -from .trainer_state import TrainerState +from .util import get_model class EMACallback(Callback): """Callback for EMA.""" - def on_train_batch_end( # pylint: disable=useless-return + def on_train_batch_end( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, + trainer: pl.Trainer, + pl_module: pl.LightningModule, outputs: DictData, batch: DictData, batch_idx: int, - ) -> None | MetricLogs: + ) -> None: """Hook to run at the end of a training batch.""" + model = get_model(pl_module) + if is_module_wrapper(model): module = model.module else: module = model + assert isinstance(module, ModelEMAAdapter), ( "Model should be wrapped with ModelEMAAdapter when using " "EMACallback." ) - module.update(trainer_state["global_step"]) - return None + + module.update(trainer.global_step) diff --git a/vis4d/engine/callbacks/evaluator.py b/vis4d/engine/callbacks/evaluator.py index 989b005cb..e79e746a0 100644 --- a/vis4d/engine/callbacks/evaluator.py +++ b/vis4d/engine/callbacks/evaluator.py @@ -3,8 +3,9 @@ from __future__ import annotations import os +from typing import Any -from torch import nn +import lightning.pytorch as pl from vis4d.common import ArgsType, MetricLogs from vis4d.common.distributed import ( @@ -18,21 +19,10 @@ from vis4d.eval.base import Evaluator from .base import Callback -from .trainer_state import TrainerState class EvaluatorCallback(Callback): - """Callback for model evaluation. - - Args: - evaluator (Evaluator): Evaluator. - metrics_to_eval (list[str], Optional): Metrics to evaluate. If None, - all metrics in the evaluator will be evaluated. Defaults to None. - save_predictions (bool): If the predictions should be saved. Defaults - to False. - save_prefix (str, Optional): Output directory for saving the - evaluation results. Defaults to None. - """ + """Callback for model evaluation.""" def __init__( self, @@ -41,9 +31,23 @@ def __init__( metrics_to_eval: list[str] | None = None, save_predictions: bool = False, save_prefix: None | str = None, + output_dir: str | None = None, **kwargs: ArgsType, ) -> None: - """Init callback.""" + """Init callback. + + Args: + evaluator (Evaluator): Evaluator. + metrics_to_eval (list[str], Optional): Metrics to evaluate. If + None, all metrics in the evaluator will be evaluated. Defaults + to None. + save_predictions (bool): If the predictions should be saved. + Defaults to False. + save_prefix (str, Optional): Output directory for saving the + evaluation results. Defaults to None. + output_dir (str, Optional): Output directory for saving the + evaluation results. + """ super().__init__(*args, **kwargs) self.evaluator = evaluator self.save_predictions = save_predictions @@ -51,23 +55,63 @@ def __init__( if self.save_predictions: assert ( - save_prefix is not None + output_dir is not None ), "If save_predictions is True, save_prefix must be provided." - self.output_dir = save_prefix - def setup(self) -> None: # pragma: no cover + output_dir = os.path.join(output_dir, "eval") + + self.output_dir = output_dir + self.save_prefix = save_prefix + + def setup( + self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str + ) -> None: # pragma: no cover """Setup callback.""" if self.save_predictions: self.output_dir = broadcast(self.output_dir) + + if self.save_prefix is not None: + self.output_dir = os.path.join( + self.output_dir, self.save_prefix + ) + for metric in self.metrics_to_eval: output_dir = os.path.join(self.output_dir, metric) os.makedirs(output_dir, exist_ok=True) self.evaluator.reset() - def on_test_batch_end( + def on_validation_batch_end( # type: ignore + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + """Hook to run at the end of a validation batch.""" + self.on_test_batch_end( + trainer=trainer, + pl_module=pl_module, + outputs=outputs, + batch=batch, + batch_idx=batch_idx, + dataloader_idx=dataloader_idx, + ) + + def on_validation_epoch_end( + self, trainer: pl.Trainer, pl_module: pl.LightningModule + ) -> None: + """Wait for on_validation_epoch_end PL hook to call 'evaluate'.""" + log_dict = self.run_eval() + + for k, v in log_dict.items(): + pl_module.log(f"val/{k}", v, sync_dist=True, rank_zero_only=True) + + def on_test_batch_end( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, + trainer: pl.Trainer, + pl_module: pl.LightningModule, outputs: DictData, batch: DictData, batch_idx: int, @@ -84,19 +128,39 @@ def on_test_batch_end( self.evaluator.save_batch(metric, output_dir) def on_test_epoch_end( - self, trainer_state: TrainerState, model: nn.Module - ) -> None | MetricLogs: + self, trainer: pl.Trainer, pl_module: pl.LightningModule + ) -> None: """Hook to run at the end of a testing epoch.""" + log_dict = self.run_eval() + + for k, v in log_dict.items(): + pl_module.log(f"test/{k}", v, sync_dist=True, rank_zero_only=True) + + def run_eval(self) -> MetricLogs: + """Run evaluation for the given evaluator.""" self.evaluator.gather(all_gather_object_cpu) synchronize() - log_dict = self.evaluate() - log_dict = broadcast(log_dict) + self.process() + + log_dict: MetricLogs = {} + for metric in self.metrics_to_eval: + metric_dict = self.evaluate(metric) + metric_dict = broadcast(metric_dict) + assert isinstance(metric_dict, dict) + log_dict.update(metric_dict) + self.evaluator.reset() + return log_dict @rank_zero_only - def evaluate(self) -> MetricLogs: + def process(self) -> None: + """Process the evaluator.""" + self.evaluator.process() + + @rank_zero_only + def evaluate(self, metric: str) -> MetricLogs: """Evaluate the performance after processing all input/output pairs. Returns: @@ -104,26 +168,26 @@ def evaluate(self) -> MetricLogs: keys are formatted as {metric_name}/{key_name}, and the values are the corresponding evaluated values. """ - rank_zero_info("Running evaluator %s...", str(self.evaluator)) - self.evaluator.process() - + rank_zero_info( + f"Running evaluator {str(self.evaluator)} with {metric} metric... " + ) log_dict = {} - for metric in self.metrics_to_eval: - # Save output predictions. This is done here instead of - # on_test_batch_end because the evaluator may not have processed - # all batches yet. - if self.save_predictions: - output_dir = os.path.join(self.output_dir, metric) - self.evaluator.save(metric, output_dir) - # Evaluate metric - metric_dict, metric_str = self.evaluator.evaluate(metric) - for k, v in metric_dict.items(): - log_k = metric + "/" + k - rank_zero_info("%s: %.4f", log_k, v) - log_dict[f"{metric}/{k}"] = v - - rank_zero_info("Showing results for metric: %s", metric) - rank_zero_info(metric_str) + # Save output predictions. This is done here instead of + # on_test_batch_end because the evaluator may not have processed + # all batches yet. + if self.save_predictions: + output_dir = os.path.join(self.output_dir, metric) + self.evaluator.save(metric, output_dir) + + # Evaluate metric + metric_dict, metric_str = self.evaluator.evaluate(metric) + for k, v in metric_dict.items(): + log_k = metric + "/" + k + rank_zero_info("%s: %.4f", log_k, v) + log_dict[f"{metric}/{k}"] = v + + rank_zero_info("Showing results for metric: %s", metric) + rank_zero_info(metric_str) return log_dict diff --git a/vis4d/engine/callbacks/logging.py b/vis4d/engine/callbacks/logging.py index 2b44245c9..01d7b70b9 100644 --- a/vis4d/engine/callbacks/logging.py +++ b/vis4d/engine/callbacks/logging.py @@ -3,18 +3,16 @@ from __future__ import annotations from collections import defaultdict +from typing import Any -from torch import nn +import lightning.pytorch as pl from vis4d.common import ArgsType, MetricLogs from vis4d.common.logging import rank_zero_info from vis4d.common.progress import compose_log_str from vis4d.common.time import Timer -from vis4d.data.typing import DictData -from vis4d.engine.loss_module import LossModule from .base import Callback -from .trainer_state import TrainerState class LoggingCallback(Callback): @@ -31,68 +29,61 @@ def __init__( self.test_timer = Timer() self.last_step = 0 - def on_train_batch_start( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - batch: DictData, - batch_idx: int, - ) -> None: - """Hook to run at the start of a training batch.""" - if not self.epoch_based and self.train_timer.paused: - self.train_timer.resume() - def on_train_epoch_start( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, + self, trainer: pl.Trainer, pl_module: pl.LightningModule ) -> None: """Hook to run at the start of a training epoch.""" if self.epoch_based: self.train_timer.reset() self.last_step = 0 self._metrics.clear() - elif trainer_state["global_step"] == 0: + elif trainer.global_step == 0: self.train_timer.reset() - def on_train_batch_end( + def on_train_batch_start( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - outputs: DictData, - batch: DictData, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + batch: Any, batch_idx: int, - ) -> None | MetricLogs: + ) -> None: + """Hook to run at the start of a training batch.""" + if self.train_timer.paused: + self.train_timer.resume() + + def on_train_batch_end( # type: ignore + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, + batch_idx: int, + ) -> None: """Hook to run at the end of a training batch.""" - if "metrics" in trainer_state: - for k, v in trainer_state["metrics"].items(): + if "metrics" in outputs: + for k, v in outputs["metrics"].items(): self._metrics[k].append(v) if self.epoch_based: cur_iter = batch_idx + 1 - total_iters = ( - trainer_state["num_train_batches"] - if trainer_state["num_train_batches"] is not None - else -1 - ) + # Resolve float("inf") to -1 + if isinstance(trainer.num_training_batches, float): + total_iters = -1 + else: + total_iters = trainer.num_training_batches else: - # After optimizer.step(), global_step is already incremented by 1. - cur_iter = trainer_state["global_step"] - total_iters = trainer_state["num_steps"] + cur_iter = trainer.global_step + 1 + total_iters = trainer.max_steps - log_dict: None | MetricLogs = None if cur_iter % self._refresh_rate == 0 and cur_iter != self.last_step: prefix = ( - f"Epoch {trainer_state['current_epoch'] + 1}" + f"Epoch {pl_module.current_epoch + 1}" if self.epoch_based else "Iter" ) - log_dict = { + log_dict: MetricLogs = { k: sum(v) / len(v) if len(v) > 0 else float("NaN") for k, v in self._metrics.items() } @@ -106,32 +97,65 @@ def on_train_batch_end( self._metrics.clear() self.last_step = cur_iter - return log_dict + for k, v in log_dict.items(): + pl_module.log(f"train/{k}", v, rank_zero_only=True) + + def on_validation_epoch_start( + self, trainer: pl.Trainer, pl_module: pl.LightningModule + ) -> None: + """Hook to run at the start of a validation epoch.""" + self.test_timer.reset() + self.train_timer.pause() + + def on_validation_batch_end( # type: ignore + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + """Wait for on_validation_batch_end PL hook to call 'process'.""" + cur_iter = batch_idx + 1 + + # Resolve float("inf") to -1 + if isinstance(trainer.num_val_batches[dataloader_idx], int): + total_iters = int(trainer.num_val_batches[dataloader_idx]) + else: + total_iters = -1 + + if cur_iter % self._refresh_rate == 0: + rank_zero_info( + compose_log_str( + "Validation", cur_iter, total_iters, self.test_timer + ) + ) def on_test_epoch_start( - self, trainer_state: TrainerState, model: nn.Module + self, trainer: pl.Trainer, pl_module: pl.LightningModule ) -> None: """Hook to run at the start of a testing epoch.""" self.test_timer.reset() - if not self.epoch_based: - self.train_timer.pause() + self.train_timer.pause() - def on_test_batch_end( + def on_test_batch_end( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, - outputs: DictData, - batch: DictData, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, batch_idx: int, dataloader_idx: int = 0, ) -> None: """Hook to run at the end of a testing batch.""" cur_iter = batch_idx + 1 - total_iters = ( - trainer_state["num_test_batches"][dataloader_idx] - if trainer_state["num_test_batches"] is not None - else -1 - ) + + # Resolve float("inf") to -1 + if isinstance(trainer.num_test_batches[dataloader_idx], int): + total_iters = int(trainer.num_test_batches[dataloader_idx]) + else: + total_iters = -1 if cur_iter % self._refresh_rate == 0: rank_zero_info( diff --git a/vis4d/pl/callbacks/scheduler.py b/vis4d/engine/callbacks/scheduler.py similarity index 57% rename from vis4d/pl/callbacks/scheduler.py rename to vis4d/engine/callbacks/scheduler.py index 456ca48a1..c8cb0a541 100644 --- a/vis4d/pl/callbacks/scheduler.py +++ b/vis4d/engine/callbacks/scheduler.py @@ -9,10 +9,17 @@ from vis4d.engine.optim.scheduler import LRSchedulerWrapper +from .base import Callback -class LRSchedulerCallback(pl.Callback): + +class LRSchedulerCallback(Callback): """Callback to configure learning rate during training.""" + def __init__(self) -> None: + """Initialize the callback.""" + super().__init__() + self.last_step = 0 + def on_train_batch_end( # type: ignore self, trainer: pl.Trainer, @@ -27,8 +34,11 @@ def on_train_batch_end( # type: ignore if not isinstance(schedulers, Iterable): schedulers = [schedulers] # type: ignore - for scheduler in schedulers: - if scheduler is None: - continue - assert isinstance(scheduler, LRSchedulerWrapper) - scheduler.step_on_batch(trainer.global_step) + if trainer.global_step != self.last_step: + for scheduler in schedulers: + if scheduler is None: + continue + assert isinstance(scheduler, LRSchedulerWrapper) + scheduler.step_on_batch(trainer.global_step) + + self.last_step = trainer.global_step diff --git a/vis4d/engine/callbacks/trainer_state.py b/vis4d/engine/callbacks/trainer_state.py deleted file mode 100644 index 2c28c7656..000000000 --- a/vis4d/engine/callbacks/trainer_state.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Trainer state for callbacks.""" - -from __future__ import annotations - -from typing import TypedDict - -from torch.optim.optimizer import Optimizer -from torch.utils.data import DataLoader -from typing_extensions import NotRequired - -from vis4d.common import TrainingModule -from vis4d.data.typing import DictData -from vis4d.engine.optim import LRSchedulerWrapper - - -class TrainerState(TypedDict): - """State of the trainer. - - Attributes: - current_epoch (int): Current epoch. - num_epochs (int): Total number of the training epochs. - global_step (int): Global step. - num_steps (int): Total number of the training steps. - train_dataloader (DataLoader[DictData] | None): Training dataloader. - num_train_batches (int | None): Number of training batches. - test_dataloader (list[DataLoader[DictData]] | None): List of test - dataloaders. - num_test_batches (list[int] | None): List of number of test batches. - optimizers (NotRequired[list[Optimizer]]): List of optimizers. - metrics (NotRequired[dict[str, float]]): Metrics for the logging. - """ - - current_epoch: int - num_epochs: int - global_step: int - num_steps: int - train_dataloader: DataLoader[DictData] | None - num_train_batches: int | None - test_dataloader: list[DataLoader[DictData]] | None - num_test_batches: list[int] | None - optimizers: NotRequired[list[Optimizer]] - lr_schedulers: NotRequired[list[LRSchedulerWrapper]] - metrics: NotRequired[dict[str, float]] - train_module: NotRequired[TrainingModule] - train_engine: NotRequired[str] diff --git a/vis4d/engine/callbacks/util.py b/vis4d/engine/callbacks/util.py new file mode 100644 index 000000000..e46038a8e --- /dev/null +++ b/vis4d/engine/callbacks/util.py @@ -0,0 +1,24 @@ +"""PyTorch Lightning callbacks utilities.""" + +from __future__ import annotations + +import lightning.pytorch as pl +from torch import nn + +from vis4d.engine.loss_module import LossModule +from vis4d.engine.training_module import TrainingModule + + +def get_model(model: pl.LightningModule) -> nn.Module: + """Get model from pl module.""" + if isinstance(model, TrainingModule): + return model.model + return model + + +def get_loss_module(loss_module: pl.LightningModule) -> LossModule: + """Get loss_module from pl module.""" + assert hasattr(loss_module, "loss_module") and isinstance( + loss_module.loss_module, LossModule + ), "Loss module is not set in the training module." + return loss_module.loss_module diff --git a/vis4d/engine/callbacks/visualizer.py b/vis4d/engine/callbacks/visualizer.py index 8b12159a1..e6b24ef55 100644 --- a/vis4d/engine/callbacks/visualizer.py +++ b/vis4d/engine/callbacks/visualizer.py @@ -3,17 +3,15 @@ from __future__ import annotations import os +from typing import Any -from torch import nn +import lightning.pytorch as pl from vis4d.common import ArgsType -from vis4d.common.distributed import broadcast -from vis4d.data.typing import DictData -from vis4d.engine.loss_module import LossModule +from vis4d.common.distributed import broadcast, synchronize from vis4d.vis.base import Visualizer from .base import Callback -from .trainer_state import TrainerState class VisualizerCallback(Callback): @@ -27,6 +25,7 @@ def __init__( show: bool = False, save_to_disk: bool = True, save_prefix: str | None = None, + output_dir: str | None = None, **kwargs: ArgsType, ) -> None: """Init callback. @@ -35,11 +34,13 @@ def __init__( visualizer (Visualizer): Visualizer. visualize_train (bool): If the training data should be visualized. Defaults to False. - save_prefix (str): Output directory for saving the visualizations. show (bool): If the visualizations should be shown. Defaults to False. save_to_disk (bool): If the visualizations should be saved to disk. Defaults to True. + save_prefix (str): Output directory prefix for distinguish + different visualizations. + output_dir (str): Output directory for saving the visualizations. """ super().__init__(*args, **kwargs) self.visualizer = visualizer @@ -50,22 +51,27 @@ def __init__( if self.save_to_disk: assert ( - save_prefix is not None - ), "If save_to_disk is True, save_prefix must be provided." - self.output_dir = f"{self.save_prefix}/vis" + output_dir is not None + ), "If save_to_disk is True, output_dir must be provided." - def setup(self) -> None: # pragma: no cover + output_dir = os.path.join(output_dir, "vis") + + self.output_dir = output_dir + self.save_prefix = save_prefix + + def setup( + self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str + ) -> None: # pragma: no cover """Setup callback.""" if self.save_to_disk: self.output_dir = broadcast(self.output_dir) - def on_train_batch_end( + def on_train_batch_end( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - outputs: DictData, - batch: DictData, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, batch_idx: int, ) -> None: """Hook to run at the end of a training batch.""" @@ -81,20 +87,41 @@ def on_train_batch_end( self.visualizer.show(cur_iter=cur_iter) if self.save_to_disk: - os.makedirs(f"{self.output_dir}/train", exist_ok=True) - self.visualizer.save_to_disk( - cur_iter=cur_iter, - output_folder=f"{self.output_dir}/train", - ) + self.save(cur_iter=cur_iter, stage="train") self.visualizer.reset() - def on_test_batch_end( + def on_validation_batch_end( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, - outputs: DictData, - batch: DictData, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + """Hook to run at the end of a validation batch.""" + cur_iter = batch_idx + 1 + + self.visualizer.process( + cur_iter=cur_iter, + **self.get_test_callback_inputs(outputs, batch), + ) + + if self.show: + self.visualizer.show(cur_iter=cur_iter) + + if self.save_to_disk: + self.save(cur_iter=cur_iter, stage="val") + + self.visualizer.reset() + + def on_test_batch_end( # type: ignore + self, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + outputs: Any, + batch: Any, batch_idx: int, dataloader_idx: int = 0, ) -> None: @@ -110,10 +137,29 @@ def on_test_batch_end( self.visualizer.show(cur_iter=cur_iter) if self.save_to_disk: - os.makedirs(f"{self.output_dir}/test", exist_ok=True) - self.visualizer.save_to_disk( - cur_iter=cur_iter, - output_folder=f"{self.output_dir}/test", - ) + self.save(cur_iter=cur_iter, stage="test") self.visualizer.reset() + + def save(self, cur_iter: int, stage: str) -> None: + """Save the visualizer state.""" + output_folder = os.path.join(self.output_dir, stage) + + if self.save_prefix is not None: + output_folder = os.path.join(output_folder, self.save_prefix) + + os.makedirs(output_folder, exist_ok=True) + + self.visualizer.save_to_disk( + cur_iter=cur_iter, output_folder=output_folder + ) + + # TODO: Add support for logging images to WandB. + # if get_rank() == 0: + # if isinstance(trainer.logger, WandbLogger) and image is not None: + # trainer.logger.log_image( + # key=f"{self.visualizer}/{cur_iter}", + # images=[image], + # ) + + synchronize() diff --git a/vis4d/engine/callbacks/yolox_callbacks.py b/vis4d/engine/callbacks/yolox_callbacks.py index 55d9201b5..dfd63ad3b 100644 --- a/vis4d/engine/callbacks/yolox_callbacks.py +++ b/vis4d/engine/callbacks/yolox_callbacks.py @@ -4,7 +4,9 @@ import random from collections import OrderedDict +from typing import Any +import lightning.pytorch as pl import torch import torch.nn.functional as F from torch import nn @@ -22,44 +24,38 @@ from vis4d.common.logging import rank_zero_info, rank_zero_warn from vis4d.data.const import CommonKeys as K from vis4d.data.data_pipe import DataPipe -from vis4d.data.typing import DictDataOrList -from vis4d.engine.loss_module import LossModule from vis4d.op.detect.yolox import YOLOXHeadLoss from vis4d.op.loss.common import l1_loss from .base import Callback -from .trainer_state import TrainerState +from .util import get_loss_module, get_model class YOLOXModeSwitchCallback(Callback): - """Callback for switching the mode of YOLOX training. - - Args: - switch_epoch (int): Epoch to switch the mode. - """ + """Callback for switching the mode of YOLOX training.""" def __init__( self, *args: ArgsType, switch_epoch: int, **kwargs: ArgsType ) -> None: - """Init callback.""" + """Init callback. + + Args: + switch_epoch (int): Epoch to switch the mode. + """ super().__init__(*args, **kwargs) self.switch_epoch = switch_epoch self.switched = False def on_train_epoch_end( - self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, + self, trainer: pl.Trainer, pl_module: pl.LightningModule ) -> None: """Hook to run at the end of a training epoch.""" - if ( - trainer_state["current_epoch"] < self.switch_epoch - 1 - or self.switched - ): + if pl_module.current_epoch < self.switch_epoch - 1 or self.switched: # TODO: Make work with resume. return + loss_module = get_loss_module(pl_module) + found_loss = False for loss in loss_module.losses: if isinstance(loss["loss"], YOLOXHeadLoss): @@ -77,10 +73,10 @@ def on_train_epoch_end( rank_zero_warn("YOLOXHeadLoss should be in LossModule.") # Set data pipeline to default DataPipe to skip strong augs. # Switch to checking validation every epoch. - dataloader = trainer_state["train_dataloader"] + dataloader = trainer.train_dataloader assert dataloader is not None new_dataloader = DataLoader( - DataPipe(dataloader.dataset.datasets), # type: ignore + DataPipe(dataloader.dataset.datasets), batch_size=dataloader.batch_size, num_workers=dataloader.num_workers, collate_fn=dataloader.collate_fn, @@ -88,25 +84,18 @@ def on_train_epoch_end( persistent_workers=dataloader.persistent_workers, pin_memory=dataloader.pin_memory, ) - train_module = trainer_state["train_module"] - train_module.check_val_every_n_epoch = 1 - if trainer_state["train_engine"] == "vis4d": - # Directly modify the train dataloader. - train_module.train_dataloader = new_dataloader - elif trainer_state["train_engine"] == "pl": - # Override train_dataloader method in PL datamodule. - # Set reload_dataloaders_every_n_epochs to 1 to use the new - # dataloader. - def train_dataloader() -> DataLoader: # type: ignore - """Return dataloader for training.""" - return new_dataloader - - train_module.datamodule.train_dataloader = train_dataloader - train_module.reload_dataloaders_every_n_epochs = self.switch_epoch - else: - raise ValueError( - f"Unsupported training engine {trainer_state['train_engine']}." - ) + + pl_module.check_val_every_n_epoch = 1 # type: ignore + + # Override train_dataloader method in PL datamodule. + # Set reload_dataloaders_every_n_epochs to 1 to use the new + # dataloader. + def train_dataloader() -> DataLoader: # type: ignore + """Return dataloader for training.""" + return new_dataloader + + pl_module.datamodule.train_dataloader = train_dataloader # type: ignore # pylint: disable=line-too-long + pl_module.reload_dataloaders_every_n_epochs = self.switch_epoch # type: ignore # pylint: disable=line-too-long self.switched = True @@ -129,22 +118,17 @@ class YOLOXSyncNormCallback(Callback): """Callback for syncing the norm states of YOLOX training.""" def on_test_epoch_start( - self, trainer_state: TrainerState, model: nn.Module + self, trainer: pl.Trainer, pl_module: pl.LightningModule ) -> None: - """Hook to run at the beginning of a testing epoch. + """Hook to run at the beginning of a testing epoch.""" + if get_world_size() > 1: + model = get_model(pl_module) + norm_states = get_norm_states(model) - Args: - trainer_state (TrainerState): Trainer state. - model (nn.Module): Model that is being trained. - """ - rank_zero_info("Synced norm states across all processes.") - if get_world_size() == 1: - return - norm_states = get_norm_states(model) - if len(norm_states) == 0: - return - norm_states = all_reduce_dict(norm_states, reduce_op="mean") - model.load_state_dict(norm_states, strict=False) + if len(norm_states) > 0: + rank_zero_info("Synced norm states across all processes.") + norm_states = all_reduce_dict(norm_states, reduce_op="mean") + model.load_state_dict(norm_states, strict=False) class YOLOXSyncRandomResizeCallback(Callback): @@ -173,18 +157,17 @@ def _get_random_shape(self, device: torch.device) -> tuple[int, int]: shape_tensor = broadcast(shape_tensor, 0) return (int(shape_tensor[0].item()), int(shape_tensor[1].item())) - def on_train_batch_start( + def on_train_batch_start( # type: ignore self, - trainer_state: TrainerState, - model: nn.Module, - loss_module: LossModule, - batch: DictDataOrList, + trainer: pl.Trainer, + pl_module: pl.LightningModule, + batch: Any, batch_idx: int, ) -> None: """Hook to run at the start of a training batch.""" if not isinstance(batch, list): batch = [batch] - if (trainer_state["global_step"] + 1) % self.interval == 0: + if (trainer.global_step + 1) % self.interval == 0: self.random_shape = self._get_random_shape( batch[0][K.images].device ) diff --git a/vis4d/pl/data_module.py b/vis4d/engine/data_module.py similarity index 100% rename from vis4d/pl/data_module.py rename to vis4d/engine/data_module.py diff --git a/vis4d/engine/experiment.py b/vis4d/engine/experiment.py deleted file mode 100644 index a0781ab12..000000000 --- a/vis4d/engine/experiment.py +++ /dev/null @@ -1,262 +0,0 @@ -"""Implementation of a single experiment. - -Helper functions to execute a single experiment. - -This will be called for each experiment configuration. -""" - -from __future__ import annotations - -import logging -import os - -import torch -from torch.distributed import destroy_process_group, init_process_group -from torch.nn.parallel import DistributedDataParallel as DDP -from torch.utils.collect_env import get_pretty_env_info - -from vis4d.common.ckpt import load_model_checkpoint -from vis4d.common.distributed import ( - broadcast, - get_local_rank, - get_rank, - get_world_size, -) -from vis4d.common.logging import ( - _info, - dump_config, - rank_zero_info, - rank_zero_warn, - setup_logger, -) -from vis4d.common.slurm import init_dist_slurm -from vis4d.common.util import init_random_seed, set_random_seed, set_tf32 -from vis4d.config import instantiate_classes -from vis4d.config.typing import ExperimentConfig - -from .optim import set_up_optimizers -from .parser import pprints_config -from .trainer import Trainer - - -def ddp_setup( - torch_distributed_backend: str = "nccl", slurm: bool = False -) -> None: - """Setup DDP environment and init processes. - - Args: - torch_distributed_backend (str): Backend to use (`nccl` or `gloo`) - slurm (bool): If set, setup slurm running jobs. - """ - if slurm: - init_dist_slurm() - - global_rank = get_rank() - world_size = get_world_size() - _info( - f"Initializing distributed: " - f"GLOBAL_RANK: {global_rank}, MEMBER: {global_rank + 1}/{world_size}" - ) - init_process_group( - torch_distributed_backend, rank=global_rank, world_size=world_size - ) - - # On rank=0 let everyone know training is starting - rank_zero_info( - f"{'-' * 100}\n" - f"distributed_backend={torch_distributed_backend}\n" - f"All distributed processes registered. " - f"Starting with {world_size} processes\n" - f"{'-' * 100}\n" - ) - - local_rank = get_local_rank() - all_gpu_ids = ",".join(str(x) for x in range(torch.cuda.device_count())) - devices = os.getenv("CUDA_VISIBLE_DEVICES", all_gpu_ids) - - torch.cuda.set_device(local_rank) - - _info(f"LOCAL_RANK: {local_rank} - CUDA_VISIBLE_DEVICES: [{devices}]") - - -def run_experiment( - config: ExperimentConfig, - mode: str, - num_gpus: int = 0, - show_config: bool = False, - use_slurm: bool = False, - ckpt_path: str | None = None, - resume: bool = False, -) -> None: - """Entry point for running a single experiment. - - Args: - config (ExperimentConfig): Configuration dictionary. - mode (str): Mode to run the experiment in. Either `fit` or `test`. - num_gpus (int): Number of GPUs to use. - show_config (bool): If set, prints the configuration. - use_slurm (bool): If set, setup slurm running jobs. This will set the - required environment variables for slurm. - ckpt_path (str | None): Path to a checkpoint to load. - resume (bool): If set, resume training from the checkpoint. - - Raises: - ValueError: If `mode` is not `fit` or `test`. - """ - # Setup logging - logger_vis4d = logging.getLogger("vis4d") - log_dir = os.path.join(config.output_dir, f"log_{config.timestamp}.txt") - setup_logger(logger_vis4d, log_dir) - - # Dump config - config_file = os.path.join( - config.output_dir, f"config_{config.timestamp}.yaml" - ) - dump_config(config, config_file) - - rank_zero_info("Environment info: %s", get_pretty_env_info()) - - # PyTorch Setting - set_tf32(config.use_tf32, config.tf32_matmul_precision) - torch.hub.set_dir(f"{config.work_dir}/.cache/torch/hub") - torch.backends.cudnn.benchmark = config.benchmark - - if show_config: - rank_zero_info(pprints_config(config)) - - # Instantiate classes - model = instantiate_classes(config.model) - - if config.get("sync_batchnorm", False): - if num_gpus > 1: - rank_zero_info( - "SyncBN enabled, converting BatchNorm layers to" - " SyncBatchNorm layers." - ) - model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) - else: - rank_zero_warn( - "use_sync_bn is True, but not in a distributed setting." - " BatchNorm layers are not converted." - ) - - # Callbacks - callbacks = [instantiate_classes(cb) for cb in config.callbacks] - - # Setup DDP & seed - seed = init_random_seed() if config.seed == -1 else config.seed - - if num_gpus > 1: - ddp_setup(slurm=use_slurm) - - # broadcast seed to all processes - seed = broadcast(seed) - - # Setup Dataloaders & seed - if mode == "fit": - set_random_seed(seed) - _info(f"[rank {get_rank()}] Global seed set to {seed}") - train_dataloader = instantiate_classes( - config.data.train_dataloader, seed=seed - ) - train_data_connector = instantiate_classes(config.train_data_connector) - optimizers, lr_schedulers = set_up_optimizers( - config.optimizers, [model], len(train_dataloader) - ) - loss = instantiate_classes(config.loss) - else: - train_dataloader = None - train_data_connector = None - - if config.data.test_dataloader is not None: - test_dataloader = instantiate_classes(config.data.test_dataloader) - else: - test_dataloader = None - - if config.test_data_connector is not None: - test_data_connector = instantiate_classes(config.test_data_connector) - else: - test_data_connector = None - - # Setup Model - if num_gpus == 0: - device = torch.device("cpu") - else: - rank = get_local_rank() - device = torch.device(f"cuda:{rank}") - - model.to(device) - - if num_gpus > 1: - model = DDP( # pylint: disable=redefined-variable-type - model, device_ids=[rank] - ) - - # Setup Callbacks - for cb in callbacks: - cb.setup() - - # Resume training - if resume: - assert mode == "fit", "Resume is only supported in fit mode" - if ckpt_path is None: - ckpt_path = os.path.join( - config.output_dir, "checkpoints/last.ckpt" - ) - rank_zero_info( - f"Restoring states from the checkpoint path at {ckpt_path}" - ) - ckpt = torch.load(ckpt_path, map_location="cpu") - - epoch = ckpt["epoch"] + 1 - global_step = ckpt["global_step"] - - for i, optimizer in enumerate( - optimizers # pylint:disable=possibly-used-before-assignment - ): - optimizer.load_state_dict(ckpt["optimizers"][i]) - - for i, lr_scheduler in enumerate( - lr_schedulers # pylint:disable=possibly-used-before-assignment - ): - lr_scheduler.load_state_dict(ckpt["lr_schedulers"][i]) - else: - epoch = 0 - global_step = 0 - - if ckpt_path is not None: - load_model_checkpoint( - model, - ckpt_path, - rev_keys=[(r"^model\.", ""), (r"^module\.", "")], - ) - - trainer = Trainer( - device=device, - output_dir=config.output_dir, - train_dataloader=train_dataloader, - test_dataloader=test_dataloader, - train_data_connector=train_data_connector, - test_data_connector=test_data_connector, - callbacks=callbacks, - num_epochs=config.params.get("num_epochs", -1), - num_steps=config.params.get("num_steps", -1), - epoch=epoch, - global_step=global_step, - check_val_every_n_epoch=config.get("check_val_every_n_epoch", 1), - val_check_interval=config.get("val_check_interval", None), - log_every_n_steps=config.get("log_every_n_steps", 50), - ) - - if resume: - rank_zero_info( - f"Restored all states from the checkpoint at {ckpt_path}" - ) - - if mode == "fit": - trainer.fit(model, optimizers, lr_schedulers, loss) - elif mode == "test": - trainer.test(model) - - if num_gpus > 1: - destroy_process_group() diff --git a/vis4d/engine/flag.py b/vis4d/engine/flag.py index f127f1909..fa4fe3769 100644 --- a/vis4d/engine/flag.py +++ b/vis4d/engine/flag.py @@ -21,6 +21,11 @@ _SLURM = flags.DEFINE_bool( "slurm", default=False, help="If set, setup slurm running jobs." ) +_VIS = flags.DEFINE_bool( + "vis", + default=False, + help="If set, running visualization using visualizer callback.", +) __all__ = [ @@ -31,4 +36,5 @@ "_SHOW_CONFIG", "_SWEEP", "_SLURM", + "_VIS", ] diff --git a/vis4d/engine/optim/scheduler.py b/vis4d/engine/optim/scheduler.py index 2e1e63370..0271f3e94 100644 --- a/vis4d/engine/optim/scheduler.py +++ b/vis4d/engine/optim/scheduler.py @@ -5,8 +5,8 @@ from typing import TypedDict -from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler +from torch.optim.optimizer import Optimizer from vis4d.common.typing import DictStrAny from vis4d.config import copy_and_resolve_references, instantiate_classes @@ -78,21 +78,23 @@ def _instantiate_lr_scheduler( "epoch_based": lr_scheduler_cfg["epoch_based"], } - def get_lr(self) -> list[float]: # type: ignore + def get_lr(self) -> list[float]: """Get current learning rate.""" - return [ - lr_scheduler["scheduler"].get_lr() - for lr_scheduler in self.lr_schedulers.values() - ] + lr = [] + for lr_scheduler in self.lr_schedulers.values(): + lr.extend(lr_scheduler["scheduler"].get_lr()) + return lr - def state_dict(self) -> dict[int, DictStrAny]: # type: ignore + def state_dict(self) -> dict[int, DictStrAny]: """Get state dict.""" state_dict = {} for scheduler_idx, lr_scheduler in self.lr_schedulers.items(): state_dict[scheduler_idx] = lr_scheduler["scheduler"].state_dict() return state_dict - def load_state_dict(self, state_dict: dict[int, DictStrAny]) -> None: # type: ignore # pylint: disable=line-too-long + def load_state_dict( + self, state_dict: dict[int, DictStrAny] # type: ignore + ) -> None: """Load state dict.""" for scheduler_idx, _state_dict in state_dict.items(): # Instantiate the lr scheduler if it is not instantiated yet @@ -107,7 +109,7 @@ def load_state_dict(self, state_dict: dict[int, DictStrAny]) -> None: # type: i def _step_lr(self, lr_scheduler: LRSchedulerDict, step: int) -> None: """Step the learning rate.""" if lr_scheduler["begin"] <= step and ( - lr_scheduler["end"] == -1 or lr_scheduler["end"] > step + lr_scheduler["end"] == -1 or lr_scheduler["end"] >= step ): lr_scheduler["scheduler"].step() @@ -126,15 +128,13 @@ def step(self, epoch: int | None = None) -> None: def step_on_batch(self, step: int) -> None: """Step on training batch end.""" - # Minus 1 because the step is called after the optimizer.step() - step -= 1 for lr_scheduler in self.lr_schedulers.values(): if not lr_scheduler["epoch_based"]: self._step_lr(lr_scheduler, step) for i, lr_scheduler_cfg in enumerate(self.lr_schedulers_cfg): if not lr_scheduler_cfg["epoch_based"] and ( - lr_scheduler_cfg["begin"] == step + 1 + lr_scheduler_cfg["begin"] == step ): self._instantiate_lr_scheduler(i, lr_scheduler_cfg) @@ -161,9 +161,9 @@ def __init__( self.factor = factor super().__init__(optimizer, last_epoch) - def get_lr(self) -> list[float]: # type: ignore + def get_lr(self) -> list[float]: """Compute current learning rate.""" - step_count = self._step_count - 1 # type: ignore + step_count = self._step_count - 1 if step_count == 0: return [ group["lr"] * self.factor @@ -211,9 +211,9 @@ def __init__( self.min_lr = min_lr super().__init__(optimizer, last_epoch) - def get_lr(self) -> list[float]: # type: ignore + def get_lr(self) -> list[float]: """Compute current learning rate.""" - step_count = self._step_count - 1 # type: ignore + step_count = self._step_count - 1 if step_count == 0 or step_count > self.max_steps: return [group["lr"] for group in self.optimizer.param_groups] decay_factor = ( @@ -245,9 +245,9 @@ def __init__( self.max_steps = max_steps super().__init__(optimizer, last_epoch) - def get_lr(self) -> list[float]: # type: ignore + def get_lr(self) -> list[float]: """Compute current learning rate.""" - step_count = self._step_count - 1 # type: ignore + step_count = self._step_count - 1 if step_count >= self.max_steps: return self.base_lrs factors = [ diff --git a/vis4d/engine/parser.py b/vis4d/engine/parser.py index 03e2189db..6cb728b57 100644 --- a/vis4d/engine/parser.py +++ b/vis4d/engine/parser.py @@ -110,7 +110,7 @@ def DEFINE_config_file( # pylint: disable=invalid-name help_string: str = "path to config file [.py |.yaml].", lock_config: bool = False, method_name: str = "get_config", -) -> flags.FlagHolder: +) -> flags.FlagHolder: # type: ignore """Registers a new flag for a config file. Args: diff --git a/vis4d/engine/run.py b/vis4d/engine/run.py index c0033ab4d..7c8205bac 100644 --- a/vis4d/engine/run.py +++ b/vis4d/engine/run.py @@ -1,74 +1,161 @@ -"""CLI interface.""" +"""CLI interface using PyTorch Lightning.""" from __future__ import annotations +import logging +import os.path as osp + +import torch from absl import app # pylint: disable=no-name-in-module +from torch.utils.collect_env import get_pretty_env_info from vis4d.common import ArgsType -from vis4d.common.logging import rank_zero_info +from vis4d.common.logging import dump_config, rank_zero_info, setup_logger +from vis4d.common.util import set_tf32 from vis4d.config import instantiate_classes -from vis4d.config.replicator import replicate_config from vis4d.config.typing import ExperimentConfig - -from .experiment import run_experiment -from .flag import _CKPT, _CONFIG, _GPUS, _RESUME, _SHOW_CONFIG, _SLURM, _SWEEP +from vis4d.engine.callbacks import ( + Callback, + LRSchedulerCallback, + VisualizerCallback, +) +from vis4d.engine.data_module import DataModule +from vis4d.engine.flag import ( + _CKPT, + _CONFIG, + _GPUS, + _RESUME, + _SHOW_CONFIG, + _VIS, +) +from vis4d.engine.parser import pprints_config +from vis4d.engine.trainer import PLTrainer +from vis4d.engine.training_module import TrainingModule def main(argv: ArgsType) -> None: """Main entry point for the CLI. Example to run this script: - >>> python -m vis4d.engine.run --config vis4d/zoo/faster_rcnn/faster_rcnn_coco.py - With parameter sweep config: - >>> python -m vis4d.engine.run fit --config vis4d/zoo/faster_rcnn/faster_rcnn_coco.py --sweep vis4d/zoo/faster_rcnn/faster_rcnn_coco.py + >>> python -m vis4d.pl.run fit --config configs/faster_rcnn/faster_rcnn_coco.py """ # Get config - assert len(argv) > 1, "Mode must be specified: `fit` or `test`" mode = argv[1] assert mode in {"fit", "test"}, f"Invalid mode: {mode}" - experiment_config: ExperimentConfig = _CONFIG.value + config: ExperimentConfig = _CONFIG.value + num_gpus = _GPUS.value + + # Setup logging + logger_vis4d = logging.getLogger("vis4d") + logger_pl = logging.getLogger("pytorch_lightning") + log_file = osp.join(config.output_dir, f"log_{config.timestamp}.txt") + setup_logger(logger_vis4d, log_file) + setup_logger(logger_pl, log_file) + + # Dump config + config_file = osp.join( + config.output_dir, f"config_{config.timestamp}.yaml" + ) + dump_config(config, config_file) + + rank_zero_info("Environment info: %s", get_pretty_env_info()) + + # PyTorch Setting + set_tf32(config.use_tf32, config.tf32_matmul_precision) + torch.hub.set_dir(f"{config.work_dir}/.cache/torch/hub") + + # Setup device + if num_gpus > 0: + config.pl_trainer.accelerator = "gpu" + config.pl_trainer.devices = num_gpus + else: + config.pl_trainer.accelerator = "cpu" + config.pl_trainer.devices = 1 + + trainer_args = instantiate_classes(config.pl_trainer).to_dict() + + if _SHOW_CONFIG.value: + rank_zero_info(pprints_config(config)) + + # Instantiate classes + if mode == "fit": + train_data_connector = instantiate_classes(config.train_data_connector) + loss = instantiate_classes(config.loss) + else: + train_data_connector = None + loss = None + + if config.test_data_connector is not None: + test_data_connector = instantiate_classes(config.test_data_connector) + else: + test_data_connector = None - if _SWEEP.value is not None: - # Perform parameter sweep - rank_zero_info( - "Found Parameter Sweep in config file. Running Parameter Sweep..." + # Callbacks + vis = _VIS.value + + callbacks: list[Callback] = [] + for cb in config.callbacks: + callback = instantiate_classes(cb) + + assert isinstance(callback, Callback), ( + "Callback must be a subclass of Callback. " + f"Provided callback: {cb} is not!" ) - experiment_config = _CONFIG.value - sweep_config = instantiate_classes(_SWEEP.value) - - for run_id, config in enumerate( - replicate_config( - experiment_config, - method=sweep_config.method, - sampling_args=sweep_config.sampling_args, - fstring=sweep_config.get("suffix", ""), - ) - ): + + if not vis and isinstance(callback, VisualizerCallback): rank_zero_info( - "Running experiment #%d: %s", - run_id, - config.experiment_name, - ) - # Run single experiment - run_experiment( - experiment_config, - mode, - _GPUS.value, - _SHOW_CONFIG.value, - _SLURM.value, + f"{callback.visualizer} is not used." + "Please set --vis=True to use it." ) + continue + + callbacks.append(callback) + + # Add needed callbacks + callbacks.append(LRSchedulerCallback()) + # Checkpoint path + ckpt_path = _CKPT.value + + # Resume training + resume = _RESUME.value + if resume: + if ckpt_path is None: + resume_ckpt_path = osp.join( + config.output_dir, "checkpoints/last.ckpt" + ) + else: + resume_ckpt_path = ckpt_path else: - # Run single experiment - run_experiment( - experiment_config, - mode, - _GPUS.value, - _SHOW_CONFIG.value, - _SLURM.value, - _CKPT.value, - _RESUME.value, + resume_ckpt_path = None + + trainer = PLTrainer(callbacks=callbacks, **trainer_args) + + hyper_params = trainer_args + + if config.get("params", None) is not None: + hyper_params.update(config.params.to_dict()) + + training_module = TrainingModule( + config.model, + config.optimizers, + loss, + train_data_connector, + test_data_connector, + hyper_params, + config.seed, + ckpt_path if not resume else None, + config.compute_flops, + config.check_unused_parameters, + ) + data_module = DataModule(config.data) + + if mode == "fit": + trainer.fit( + training_module, datamodule=data_module, ckpt_path=resume_ckpt_path ) + elif mode == "test": + trainer.test(training_module, datamodule=data_module, verbose=False) def entrypoint() -> None: diff --git a/vis4d/engine/trainer.py b/vis4d/engine/trainer.py index 8235e4b98..5aadf210d 100644 --- a/vis4d/engine/trainer.py +++ b/vis4d/engine/trainer.py @@ -1,409 +1,141 @@ -"""Trainer for running train and test.""" +"""Trainer for PyTorch Lightning.""" from __future__ import annotations -import torch -from torch import nn -from torch.optim.optimizer import Optimizer -from torch.utils.data import DataLoader -from torch.utils.data.distributed import DistributedSampler -from torch.utils.tensorboard.writer import SummaryWriter +import datetime +import os.path as osp -from vis4d.common.distributed import rank_zero_only -from vis4d.common.logging import rank_zero_info, rank_zero_warn -from vis4d.data.typing import DictData -from vis4d.engine.callbacks import Callback, TrainerState -from vis4d.engine.connectors import DataConnector -from vis4d.engine.loss_module import LossModule +from lightning.pytorch import Callback, Trainer +from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint +from lightning.pytorch.loggers import Logger, TensorBoardLogger +from lightning.pytorch.loggers.wandb import WandbLogger +from lightning.pytorch.strategies.ddp import DDPStrategy -from .optim import LRSchedulerWrapper -from .util import move_data_to_device +from vis4d.common import ArgsType +from vis4d.common.imports import TENSORBOARD_AVAILABLE +from vis4d.common.logging import rank_zero_info -class Trainer: - """Trainer class.""" +class PLTrainer(Trainer): + """Trainer for PyTorch Lightning.""" def __init__( self, - device: torch.device, - output_dir: str, - train_dataloader: DataLoader[DictData] | None, - test_dataloader: list[DataLoader[DictData]] | None, - train_data_connector: DataConnector | None, - test_data_connector: DataConnector | None, - callbacks: list[Callback], - num_epochs: int = 1000, - num_steps: int = -1, - epoch: int = 0, - global_step: int = 0, - check_val_every_n_epoch: int | None = 1, - val_check_interval: int | None = None, - log_every_n_steps: int = 50, + *args: ArgsType, + work_dir: str, + exp_name: str, + version: str, + epoch_based: bool = True, + find_unused_parameters: bool = False, + save_top_k: int = 1, + checkpoint_period: int = 1, + checkpoint_callback: ModelCheckpoint | None = None, + wandb: bool = False, + seed: int = -1, + timeout: int = 3600, + wandb_id: str | None = None, + **kwargs: ArgsType, ) -> None: - """Initialize the trainer. + """Perform some basic common setups at the beginning of a job. Args: - device (torch.device): Device that should be used for training. - output_dir (str): Output directory for saving tensorboard logs. - train_dataloader (DataLoader[DictData] | None, optional): - Dataloader for training. - test_dataloader (list[DataLoader[DictData]] | None, optional): - Dataloaders for testing. - train_data_connector (DataConnector | None): Data connector used - for generating training inputs from a batch of data. - test_data_connector (DataConnector | None): Data connector used for - generating testing inputs from a batch of data. - callbacks (list[Callback]): Callbacks that should be used during - training. - num_epochs (int, optional): Number of training epochs. Defaults to - 1000. - num_steps (int, optional): Number of training steps. Defaults to - -1. - epoch (int, optional): Starting epoch. Defaults to 0. - global_step (int, optional): Starting step. Defaults to 0. - check_val_every_n_epoch (int | None, optional): Evaluate the model - every n epochs during training. Defaults to 1. - val_check_interval (int | None, optional): Interval for evaluating - the model during training. Defaults to None. - log_every_n_steps (int, optional): Log the training status every n - steps. Defaults to 50. + work_dir: Specific directory to save checkpoints, logs, etc. + Integrates with exp_name and version to get output_dir. + exp_name: Name of current experiment. + version: Version of current experiment. + epoch_based: Use epoch-based / iteration-based training. Default is + True. + find_unused_parameters: Activates PyTorch checking for unused + parameters in DDP setting. Default: False, for better + performance. + save_top_k: Save top k checkpoints. Default: 1 (save last). + checkpoint_period: After N epochs / stpes, save out checkpoints. + Default: 1. + checkpoint_callback: Custom PL checkpoint callback. Default: None. + wandb: Use weights and biases logging instead of tensorboard. + Default: False. + seed (int, optional): The integer value seed for global random + state. Defaults to -1. If -1, a random seed will be generated. + This will be set by TrainingModule. + timeout: Timeout (seconds) for DDP connection. Default: 3600. + wandb_id: If using wandb, the id of the run. If None, a new run + will be created. Default: None. """ - self.device = device - self.output_dir = output_dir - self.train_dataloader = train_dataloader - self.test_dataloader = test_dataloader - self.train_data_connector = train_data_connector - self.test_data_connector = test_data_connector - self.callbacks = callbacks - - if num_epochs == -1 and num_steps == -1: - rank_zero_info( - "Neither `num_epochs` nor `num_steps` is specified. " - + "Training will run indefinitely." - ) - - self.num_epochs = num_epochs - self.num_steps = num_steps - - if check_val_every_n_epoch is None and val_check_interval is None: - rank_zero_warn("Validation is disabled during training.") - - self.check_val_every_n_epoch = check_val_every_n_epoch - self.val_check_interval = val_check_interval - self.log_every_n_steps = log_every_n_steps - - self.epoch = epoch - self.global_step = global_step - - self._setup_logger() - - @rank_zero_only - def _setup_logger(self) -> None: - """Setup trainer logger.""" - self.writer = SummaryWriter(self.output_dir) - - @rank_zero_only - def _teardown_logger(self) -> None: - """Teardown trainer logger.""" - self.writer.close() - - @rank_zero_only - def _log_scalar(self, tag: str, scalar_value: float) -> None: - """Setup trainer logger.""" - self.writer.add_scalar(tag, scalar_value, self.global_step) - - def _log_lr(self, optimizer: Optimizer) -> None: - """Log learning rate.""" - tag = f"lr-{optimizer.__class__.__name__}" - - if len(optimizer.param_groups) > 1: - for i, param_group in enumerate(optimizer.param_groups): - self._log_scalar(f"{tag}/pg{i+1}", param_group["lr"]) - else: - self._log_scalar(tag, optimizer.param_groups[0]["lr"]) - - def _run_test_on_epoch(self, epoch: int) -> bool: - """Return whether to run test on current training epoch. - - Args: - epoch (int): Current training epoch. - - Returns: - bool: Whether to run test. - """ - if self.check_val_every_n_epoch is None: - return False - return (epoch + 1) % self.check_val_every_n_epoch == 0 - - def _run_test_on_step(self, step: int) -> bool: - """Return whether to run test on current training step. - - Args: - step (int): Current training step. - - Returns: - bool: Whether to run test. - """ - if self.val_check_interval is None: - return False - return (step + 1) % self.val_check_interval == 0 - - def done(self) -> bool: - """Return whether training is done.""" - is_done = False - if _is_max_limit_reached(self.global_step, self.num_steps): - rank_zero_info( - f"`Trainer.fit` stopped: `num_steps={self.num_steps!r}` " - + "reached." - ) - is_done = True - elif _is_max_limit_reached(self.epoch, self.num_epochs): - rank_zero_info( - f"`Trainer.fit` stopped: `num_epochs={self.num_epochs!r}` " - + "reached." - ) - is_done = True - - if is_done: - self._teardown_logger() - - return is_done - - def fit( - self, - model: nn.Module, - optimizers: list[Optimizer], - lr_schedulers: list[LRSchedulerWrapper], - loss_module: LossModule, - ) -> None: - """Training loop. - - Args: - model (nn.Module): Model that should be trained. - optimizers (list[Optimizer]): Optimizers that should be used for - training. - lr_schedulers (list[LRSchedulerWrapper]): Learning rate schedulers - that should be used for training. - loss_module (LossModule): Loss module that should be used for - training. - - Raises: - TypeError: If the loss value is not a torch.Tensor or a dict of - torch.Tensor. - """ - assert ( - self.train_data_connector is not None - ), "No train data connector." - assert self.train_dataloader is not None, "No train dataloader." - - while True: - # Run callbacks for epoch begin - for callback in self.callbacks: - callback.on_train_epoch_start( - self.get_state(), model, loss_module + self.work_dir = work_dir + self.exp_name = exp_name + self.version = version + self.seed = seed + + self.output_dir = osp.join(work_dir, exp_name, version) + + # setup experiment logging + if "logger" not in kwargs or ( + isinstance(kwargs["logger"], bool) and kwargs["logger"] + ): + exp_logger: Logger | None = None + if wandb: # pragma: no cover + exp_logger = WandbLogger( + save_dir=work_dir, + project=exp_name, + name=version, + id=wandb_id, ) - - # Set model to train mode - model.train() - - # Set epoch for distributed sampler - if hasattr(self.train_dataloader, "sampler") and isinstance( - self.train_dataloader.sampler, DistributedSampler - ): - self.train_dataloader.sampler.set_epoch(self.epoch) - - # Training loop for one epoch - for batch_idx, data in enumerate(self.train_dataloader): - # Log epoch - if (self.global_step + 1) % self.log_every_n_steps == 0: - self._log_scalar("epoch", self.epoch) - - # Zero grad optimizers - for opt in optimizers: - opt.zero_grad() - - # Input data - data = move_data_to_device(data, self.device) - - for callback in self.callbacks: - callback.on_train_batch_start( - trainer_state=self.get_state(), - model=model, - loss_module=loss_module, - batch=data, - batch_idx=batch_idx, - ) - - # Forward + backward + optimize - output = model(**self.train_data_connector(data)) - - total_loss, metrics = loss_module(output, data) - - total_loss.backward() - - for optimizer in optimizers: - # Log learning rate - if (self.global_step + 1) % self.log_every_n_steps == 0: - self._log_lr(optimizer) - - # Step optimizers - optimizer.step() - self.global_step += 1 - - # Step learning rate schedulers - for lr_scheduler in lr_schedulers: - lr_scheduler.step_on_batch(self.global_step) - - for callback in self.callbacks: - log_dict = callback.on_train_batch_end( - trainer_state=self.get_state(metrics), - model=model, - loss_module=loss_module, - outputs=output, - batch=data, - batch_idx=batch_idx, - ) - - if log_dict is not None: - for k, v in log_dict.items(): - self._log_scalar(f"train/{k}", v) - - # Testing (step-based) - if ( - self._run_test_on_step(self.global_step) - and self.test_dataloader is not None - ): - self.test(model) - - # Set model back to train mode - model.train() - - if self.done(): - return - - # Update learning rate on epoch - for lr_scheduler in lr_schedulers: - lr_scheduler.step(self.epoch) - - # Run callbacks for epoch end - for callback in self.callbacks: - callback.on_train_epoch_end( - self.get_state( - optimizers=optimizers, lr_schedulers=lr_schedulers - ), - model, - loss_module, + elif TENSORBOARD_AVAILABLE: + exp_logger = TensorBoardLogger( + save_dir=work_dir, + name=exp_name, + version=version, + default_hp_metric=False, ) + else: + rank_zero_info( + "Neither `tensorboard` nor `tensorboardX` is " + "available. Running without experiment logger. To log " + "your experiments, try `pip install`ing either." + ) + kwargs["logger"] = exp_logger + + callbacks: list[Callback] = [] + + # add learning rate / GPU stats monitor (logs to tensorboard) + if TENSORBOARD_AVAILABLE or wandb: + callbacks += [LearningRateMonitor(logging_interval="step")] + + # Model checkpointer + if checkpoint_callback is None: + if epoch_based: + checkpoint_cb = ModelCheckpoint( + dirpath=osp.join(self.output_dir, "checkpoints"), + verbose=True, + save_last=True, + save_top_k=save_top_k, + every_n_epochs=checkpoint_period, + save_on_train_epoch_end=True, + ) + else: + checkpoint_cb = ModelCheckpoint( + dirpath=osp.join(self.output_dir, "checkpoints"), + verbose=True, + save_last=True, + save_top_k=save_top_k, + every_n_train_steps=checkpoint_period, + ) + else: + checkpoint_cb = checkpoint_callback + callbacks += [checkpoint_cb] + + kwargs["callbacks"] += callbacks + + # add distributed strategy + if kwargs["devices"] == 0: + kwargs["accelerator"] = "cpu" + kwargs["devices"] = "auto" + elif kwargs["devices"] > 1: # pragma: no cover + if kwargs["accelerator"] == "gpu": + ddp_plugin = DDPStrategy( + find_unused_parameters=find_unused_parameters, + timeout=datetime.timedelta(timeout), + ) + kwargs["strategy"] = ddp_plugin - # Testing (epoch-based) - if ( - self._run_test_on_epoch(self.epoch) - and self.test_dataloader is not None - ): - self.test(model, is_val=True) - - self.epoch += 1 - - if self.done(): - return - - @torch.no_grad() - def test(self, model: nn.Module, is_val: bool = False) -> None: - """Testing loop. - - Args: - model (nn.Module): Model that should be tested. - is_val (bool): Whether the test is run on the validation set during - training. - """ - assert self.test_data_connector is not None, "No test data connector." - assert self.test_dataloader is not None, "No test dataloader." - - model.eval() - - # run callbacks on test epoch begin - for callback in self.callbacks: - callback.on_test_epoch_start(self.get_state(), model) - - for i, test_loader in enumerate(self.test_dataloader): - for batch_idx, data in enumerate(test_loader): - data = move_data_to_device(data, self.device) - test_input = self.test_data_connector(data) - - # forward - output = model(**test_input) - - for callback in self.callbacks: - callback.on_test_batch_end( - trainer_state=self.get_state(), - model=model, - outputs=output, - batch=data, - batch_idx=batch_idx, - dataloader_idx=i, - ) - - # run callbacks on test epoch end - for callback in self.callbacks: - log_dict = callback.on_test_epoch_end(self.get_state(), model) - - if log_dict is not None: - for k, v in log_dict.items(): - key = f"val/{k}" if is_val else f"test/{k}" - self._log_scalar(key, v) - - def get_state( - self, - metrics: dict[str, float] | None = None, - optimizers: list[Optimizer] | None = None, - lr_schedulers: list[LRSchedulerWrapper] | None = None, - ) -> TrainerState: - """Get the state of the trainer.""" - num_train_batches = ( - len(self.train_dataloader) - if self.train_dataloader is not None - else None - ) - - num_test_batches = ( - [len(test_loader) for test_loader in self.test_dataloader] - if self.test_dataloader is not None - else None - ) - - trainer_state = TrainerState( - current_epoch=self.epoch, - num_epochs=self.num_epochs, - global_step=self.global_step, - num_steps=self.num_steps, - train_dataloader=self.train_dataloader, - num_train_batches=num_train_batches, - test_dataloader=self.test_dataloader, - num_test_batches=num_test_batches, - train_module=self, - train_engine="vis4d", - ) - - if metrics is not None: - trainer_state["metrics"] = metrics - - if optimizers is not None: - trainer_state["optimizers"] = optimizers - - if lr_schedulers is not None: - trainer_state["lr_schedulers"] = lr_schedulers - - return trainer_state - - -def _is_max_limit_reached(current: int, maximum: int = -1) -> bool: - """Check if the limit has been reached (if enabled). - - Args: - current: the current value - maximum: the maximum value (or -1 to disable limit) - - Returns: - bool: whether the limit has been reached - """ - return maximum != -1 and current >= maximum + super().__init__(*args, **kwargs) diff --git a/vis4d/pl/training_module.py b/vis4d/engine/training_module.py similarity index 89% rename from vis4d/pl/training_module.py rename to vis4d/engine/training_module.py index 8cc781535..99a3cf955 100644 --- a/vis4d/pl/training_module.py +++ b/vis4d/engine/training_module.py @@ -6,14 +6,16 @@ import lightning.pytorch as pl from lightning.pytorch import seed_everything +from lightning.pytorch.core.optimizer import LightningOptimizer from ml_collections import ConfigDict from torch import nn +from torch.optim.optimizer import Optimizer from vis4d.common.ckpt import load_model_checkpoint from vis4d.common.distributed import broadcast from vis4d.common.imports import FVCORE_AVAILABLE from vis4d.common.logging import rank_zero_info -from vis4d.common.typing import DictStrAny +from vis4d.common.typing import DictStrAny, GenericFunc from vis4d.common.util import init_random_seed from vis4d.config import instantiate_classes from vis4d.config.typing import OptimizerConfig @@ -45,6 +47,7 @@ def __init__( seed: int = -1, ckpt_path: None | str = None, compute_flops: bool = False, + check_unused_parameters: bool = False, ) -> None: """Initialize the TrainingModule. @@ -63,6 +66,8 @@ def __init__( Defaults to None. compute_flops (bool, optional): If to compute the FLOPs of the model. Defaults to False. + check_unused_parameters (bool, optional): If to check the + unused parameters. Defaults to False. """ super().__init__() self.model_cfg = model_cfg @@ -74,6 +79,7 @@ def __init__( self.seed = seed self.ckpt_path = ckpt_path self.compute_flops = compute_flops + self.check_unused_parameters = check_unused_parameters # Create model placeholder self.model: nn.Module @@ -187,3 +193,18 @@ def lr_scheduler_step( # type: ignore # pylint: disable=arguments-differ,line-t """Perform a step on the lr scheduler.""" # TODO: Support metric if needed scheduler.step(self.current_epoch) + + def optimizer_step( + self, + epoch: int, + batch_idx: int, + optimizer: Optimizer | LightningOptimizer, + optimizer_closure: GenericFunc | None = None, + ) -> None: + """Optimizer step.""" + if self.check_unused_parameters: + for name, param in self.model.named_parameters(): + if param.grad is None: + rank_zero_info(name) + + optimizer.step(closure=optimizer_closure) diff --git a/vis4d/engine/util.py b/vis4d/engine/util.py deleted file mode 100644 index 9d7c1a3c4..000000000 --- a/vis4d/engine/util.py +++ /dev/null @@ -1,221 +0,0 @@ -"""Run utilities.""" - -from __future__ import annotations - -import dataclasses -from abc import ABC -from collections import OrderedDict, defaultdict -from collections.abc import Callable, Mapping, Sequence -from copy import deepcopy -from typing import Any - -import torch -from torch import Tensor - -from vis4d.common.named_tuple import is_namedtuple - -_BLOCKING_DEVICE_TYPES = ("cpu", "mps") - - -class TransferableDataType(ABC): - """A custom type for data that can be moved to a torch device. - - Example: - >>> isinstance(dict, TransferableDataType) - False - >>> isinstance(torch.rand(2, 3), TransferableDataType) - True - >>> class CustomObject: - ... def __init__(self): - ... self.x = torch.rand(2, 2) - ... def to(self, device): - ... self.x = self.x.to(device) - ... return self - >>> isinstance(CustomObject(), TransferableDataType) - True - """ - - @classmethod - def __subclasshook__(cls, subclass: Any) -> bool | Any: # type: ignore - """Subclass hook.""" - if cls is TransferableDataType: - to = getattr(subclass, "to", None) - return callable(to) - return NotImplemented # pragma: no cover - - -def is_dataclass_instance(obj: object) -> bool: - """Check if obj is dataclass instance. - - https://docs.python.org/3/library/dataclasses.html#module-level-decorators-classes-and-functions - """ - return dataclasses.is_dataclass(obj) and not isinstance(obj, type) - - -def apply_to_collection( # type: ignore - data: Any, - dtype: type | Any | tuple[type | Any], - function: Callable[[Any], Any], - *args: Any, - wrong_dtype: None | type | tuple[type, ...] = None, - include_none: bool = True, - **kwargs: Any, -) -> Any: - """Recursively applies a function to all elements of a certain dtype. - - Args: - data: the collection to apply the function to - dtype: the given function will be applied to all elements of this dtype - function: the function to apply - *args: positional arguments (will be forwarded to calls of - ``function``) - wrong_dtype: the given function won't be applied if this type is - specified and the given collections is of the ``wrong_dtype`` even - if it is of type ``dtype`` - include_none: Whether to include an element if the output of - ``function`` is ``None``. - **kwargs: keyword arguments (will be forwarded to calls of - ``function``) - - Raises: - ValueError: If frozen dataclass is passed to `apply_to_collection`. - - Returns: - The resulting collection - """ - # Breaking condition - if isinstance(data, dtype) and ( - wrong_dtype is None or not isinstance(data, wrong_dtype) - ): - return function(data, *args, **kwargs) - - elem_type = type(data) - - # Recursively apply to collection items - if isinstance(data, Mapping): - out = [] - for k, v in data.items(): - v = apply_to_collection( - v, - dtype, - function, - *args, - wrong_dtype=wrong_dtype, - include_none=include_none, - **kwargs, - ) - if include_none or v is not None: - out.append((k, v)) - if isinstance(data, defaultdict): - return elem_type(data.default_factory, OrderedDict(out)) - return elem_type(OrderedDict(out)) - - is_namedtuple_ = is_namedtuple(data) - is_sequence = isinstance(data, Sequence) and not isinstance(data, str) - if is_namedtuple_ or is_sequence: - out = [] - for d in data: - v = apply_to_collection( - d, - dtype, - function, - *args, - wrong_dtype=wrong_dtype, - include_none=include_none, - **kwargs, - ) - if include_none or v is not None: - out.append(v) - return elem_type(*out) if is_namedtuple_ else elem_type(out) - - if is_dataclass_instance(data): - # make a deepcopy of the data, - # but do not deepcopy mapped fields since the computation would - # be wasted on values that likely get immediately overwritten - fields = {} - memo = {} - for field in dataclasses.fields(data): - field_value = getattr(data, field.name) - fields[field.name] = (field_value, field.init) - memo[id(field_value)] = field_value - result = deepcopy(data, memo=memo) - # apply function to each field - for field_name, (field_value, field_init) in fields.items(): - v = None - if field_init: - v = apply_to_collection( - field_value, - dtype, - function, - *args, - wrong_dtype=wrong_dtype, - include_none=include_none, - **kwargs, - ) - if not field_init or ( - not include_none and v is None - ): # retain old value - v = getattr(data, field_name) - try: - setattr(result, field_name, v) - except dataclasses.FrozenInstanceError as e: - raise ValueError( - "A frozen dataclass was passed to `apply_to_collection` " - "but this is not allowed." - ) from e - return result - - # data is neither of dtype, nor a collection - return data - - -def move_data_to_device( # type: ignore - batch: Any, - device: torch.device | str | int, - convert_to_numpy: bool = False, -) -> Any: - """Transfers a collection of data to the given device. - - Any object that defines a method ``to(device)`` will be moved and all other - objects in the collection will be left untouched. - - This implementation is modified from - https://github.com/Lightning-AI/lightning - - Args: - batch: A tensor or collection of tensors or anything that has a method - ``.to(...)``. See :func:`apply_to_collection` for a list of - supported collection types. - device: The device to which the data should be moved. - convert_to_numpy: Whether to convert from tensor to numpy array. - - Return: - The same collection but with all contained tensors residing on the new - device. - """ - if isinstance(device, str): - device = torch.device(device) - - def batch_to(data: Any) -> Any: # type: ignore[misc] - kwargs = {} - # Don't issue non-blocking transfers to CPU - # Same with MPS due to a race condition bug: - # https://github.com/pytorch/pytorch/issues/83015 - if ( - isinstance(data, Tensor) - and isinstance(device, torch.device) - and device.type not in _BLOCKING_DEVICE_TYPES - ): - kwargs["non_blocking"] = True - data_output = data.to(device, **kwargs) - if data_output is not None: - if convert_to_numpy: - data_output = data_output.numpy() - return data_output - # user wrongly implemented the `TransferableDataType` and forgot to - # return `self`. - return data - - return apply_to_collection( - batch, dtype=TransferableDataType, function=batch_to - ) diff --git a/vis4d/eval/base.py b/vis4d/eval/base.py index d152d951b..3b49e6f0a 100644 --- a/vis4d/eval/base.py +++ b/vis4d/eval/base.py @@ -2,7 +2,7 @@ from __future__ import annotations -from vis4d.common import ArgsType, GenericFunc, MetricLogs +from vis4d.common.typing import GenericFunc, MetricLogs, unimplemented class Evaluator: # pragma: no cover @@ -83,20 +83,11 @@ def reset(self) -> None: """ raise NotImplementedError - def process_batch(self, *args: ArgsType) -> None: - """Process a batch of data. - - Raises: - NotImplementedError: This is an abstract class method. - """ - raise NotImplementedError + # Process a batch of data. + process_batch: GenericFunc = unimplemented def process(self) -> None: - """Process all accumulated data at the end of an epoch, if any. - - Raises: - NotImplementedError: This is an abstract class method. - """ + """Process all accumulated data at the end of an epoch, if any.""" def evaluate(self, metric: str) -> tuple[MetricLogs, str]: """Evaluate all predictions according to given metric. diff --git a/vis4d/eval/bdd100k/seg.py b/vis4d/eval/bdd100k/seg.py index d445bda3e..a874d642b 100644 --- a/vis4d/eval/bdd100k/seg.py +++ b/vis4d/eval/bdd100k/seg.py @@ -69,7 +69,7 @@ def reset(self) -> None: """Reset the evaluator.""" self.frames = [] - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, data_names: list[str], masks_list: list[ArrayLike] ) -> None: """Process tracking results.""" diff --git a/vis4d/eval/coco/detect.py b/vis4d/eval/coco/detect.py index 0acc42651..dce40302e 100644 --- a/vis4d/eval/coco/detect.py +++ b/vis4d/eval/coco/detect.py @@ -13,21 +13,29 @@ from pycocotools.cocoeval import COCOeval from terminaltables import AsciiTable -from vis4d.common import DictStrAny, GenericFunc, MetricLogs, NDArrayNumber +from vis4d.common.array import array_to_numpy from vis4d.common.logging import rank_zero_warn +from vis4d.common.typing import ( + ArrayLike, + DictStrAny, + GenericFunc, + MetricLogs, + NDArrayF32, + NDArrayI64, +) from vis4d.data.datasets.coco import coco_det_map from ..base import Evaluator -def xyxy_to_xywh(boxes: NDArrayNumber) -> NDArrayNumber: +def xyxy_to_xywh(boxes: NDArrayF32) -> NDArrayF32: """Convert Tensor [N, 4] in xyxy format into xywh. Args: - boxes (NDArrayNumber): Bounding boxes in Vis4D format. + boxes (NDArrayF32): Bounding boxes in Vis4D format. Returns: - NDArrayNumber: COCO format bounding boxes. + NDArrayF32: COCO format bounding boxes. """ boxes[:, 2] = boxes[:, 2] - boxes[:, 0] boxes[:, 3] = boxes[:, 3] - boxes[:, 1] @@ -54,10 +62,10 @@ def predictions_to_coco( cat_map: dict[str, int], coco_id2name: dict[int, str], image_id: int, - boxes: NDArrayNumber, - scores: NDArrayNumber, - classes: NDArrayNumber, - masks: None | NDArrayNumber = None, + boxes: NDArrayF32, + scores: NDArrayF32, + classes: NDArrayI64, + masks: None | NDArrayF32 = None, ) -> list[DictStrAny]: """Convert Vis4D format predictions to COCO format. @@ -65,10 +73,10 @@ def predictions_to_coco( cat_map (dict[str, int]): COCO class name to class ID mapping. coco_id2name (dict[int, str]): COCO class ID to class name mapping. image_id (int): ID of image. - boxes (NDArrayNumber): Predicted bounding boxes. - scores (NDArrayNumber): Predicted scores for each box. - classes (NDArrayNumber): Predicted classes for each box. - masks (None | NDArrayNumber, optional): Predicted masks. Defaults to + boxes (NDArrayF32): Predicted bounding boxes. + scores (NDArrayF32): Predicted scores for each box. + classes (NDArrayI64): Predicted classes for each box. + masks (None | NDArrayF32, optional): Predicted masks. Defaults to None. Returns: @@ -76,8 +84,8 @@ def predictions_to_coco( """ predictions = [] boxes_xyxy = copy.deepcopy(boxes) - boxes_wywh = xyxy_to_xywh(boxes_xyxy) - for i, (box, score, cls) in enumerate(zip(boxes_wywh, scores, classes)): + boxes_xywh = xyxy_to_xywh(boxes_xyxy) + for i, (box, score, cls) in enumerate(zip(boxes_xywh, scores, classes)): mask = masks[i] if masks is not None else None xywh = box.tolist() area = float(xywh[2] * xywh[3]) @@ -91,7 +99,7 @@ def predictions_to_coco( } if mask is not None: annotation["segmentation"] = maskUtils.encode( - np.array(mask.cpu(), order="F", dtype="uint8") + np.array(mask, order="F", dtype="uint8") ) annotation["segmentation"]["counts"] = annotation["segmentation"][ "counts" @@ -131,7 +139,6 @@ def __init__( coco_gt_cats = self._coco_gt.loadCats(self._coco_gt.getCatIds()) self.cat_map = {c["name"]: c["id"] for c in coco_gt_cats} self._predictions: list[DictStrAny] = [] - self.coco_dt: COCO | None = None @property def metrics(self) -> list[str]: @@ -151,36 +158,45 @@ def gather(self, gather_func: GenericFunc) -> None: def reset(self) -> None: """Reset the saved predictions to start new round of evaluation.""" self._predictions = [] - self.coco_dt = None - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, coco_image_id: list[int], - pred_boxes: list[NDArrayNumber], - pred_scores: list[NDArrayNumber], - pred_classes: list[NDArrayNumber], - pred_masks: None | list[NDArrayNumber] = None, + pred_boxes: list[ArrayLike], + pred_scores: list[ArrayLike], + pred_classes: list[ArrayLike], + pred_masks: None | list[ArrayLike] = None, ) -> None: """Process sample and convert detections to coco format. coco_image_id (list[int]): COCO image ID. - pred_boxes (list[NDArrayNumber]): Predicted bounding boxes. - pred_scores (list[NDArrayNumber]): Predicted scores for each box. - pred_classes (list[NDArrayNumber]): Predicted classes for each box. - pred_masks (None | list[NDArrayNumber], optional): Predicted masks. + pred_boxes (list[ArrayLike]): Predicted bounding boxes. + pred_scores (list[ArrayLike]): Predicted scores for each box. + pred_classes (list[ArrayLike]): Predicted classes for each box. + pred_masks (None | list[ArrayLike], optional): Predicted masks. """ for i, (image_id, boxes, scores, classes) in enumerate( zip(coco_image_id, pred_boxes, pred_scores, pred_classes) ): - masks = pred_masks[i] if pred_masks else None + boxes_np = array_to_numpy(boxes, n_dims=None, dtype=np.float32) + scores_np = array_to_numpy(scores, n_dims=None, dtype=np.float32) + classes_np = array_to_numpy(classes, n_dims=None, dtype=np.int64) + + if pred_masks is not None: + masks_np = array_to_numpy( + pred_masks[i], n_dims=3, dtype=np.float32 + ) + else: + masks_np = None + coco_preds = predictions_to_coco( self.cat_map, self.coco_id2name, image_id, - boxes, - scores, - classes, - masks, + boxes_np, + scores_np, + classes_np, + masks_np, ) self._predictions.extend(coco_preds) diff --git a/vis4d/eval/common/binary.py b/vis4d/eval/common/binary.py index b9ea9db93..30f661d39 100644 --- a/vis4d/eval/common/binary.py +++ b/vis4d/eval/common/binary.py @@ -102,7 +102,7 @@ def reset(self) -> None: self.false_negatives = [] self.n_samples = [] - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, prediction: ArrayLike, groundtruth: ArrayLike, diff --git a/vis4d/eval/common/depth.py b/vis4d/eval/common/depth.py index 02d7db4b6..9d9b3244c 100644 --- a/vis4d/eval/common/depth.py +++ b/vis4d/eval/common/depth.py @@ -94,14 +94,14 @@ def _apply_mask( mask = (target > self.min_depth) & (target <= self.max_depth) return prediction[mask], target[mask] - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, prediction: ArrayLike, groundtruth: ArrayLike ) -> None: """Process a batch of data. Args: - prediction (np.array): Prediction optical flow, in shape (H, W, 2). - groundtruth (np.array): Target optical flow, in shape (H, W, 2). + prediction (np.array): Prediction optical flow, in shape (B, H, W). + groundtruth (np.array): Target optical flow, in shape (B, H, W). """ preds = ( array_to_numpy(prediction, n_dims=None, dtype=np.float32) diff --git a/vis4d/eval/common/flow.py b/vis4d/eval/common/flow.py index ec4780614..1a591587b 100644 --- a/vis4d/eval/common/flow.py +++ b/vis4d/eval/common/flow.py @@ -68,7 +68,7 @@ def _apply_mask( mask = np.sum(np.abs(target), axis=-1) <= self.max_flow return prediction[mask], target[mask] - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, prediction: ArrayLike, groundtruth: ArrayLike ) -> None: """Process a batch of data. diff --git a/vis4d/eval/common/seg.py b/vis4d/eval/common/seg.py index 8a6430da3..9d067acf4 100644 --- a/vis4d/eval/common/seg.py +++ b/vis4d/eval/common/seg.py @@ -86,7 +86,7 @@ def reset(self) -> None: """Reset the saved predictions to start new round of evaluation.""" self._confusion_matrix = None - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, prediction: ArrayLike, groundtruth: ArrayLike ) -> None: """Process sample and update confusion matrix. @@ -100,7 +100,7 @@ def process_batch( # type: ignore # pylint: disable=arguments-differ """ confusion_matrix = self.calc_confusion_matrix( array_to_numpy(prediction, n_dims=None, dtype=np.float32), - array_to_numpy(groundtruth, n_dims=None, dtype=np.int64), # type: ignore # pylint: disable=line-too-long + array_to_numpy(groundtruth, n_dims=None, dtype=np.int64), ) if self._confusion_matrix is None: diff --git a/vis4d/eval/nuscenes/detect3d.py b/vis4d/eval/nuscenes/detect3d.py index 178b9c3cb..c2817120e 100644 --- a/vis4d/eval/nuscenes/detect3d.py +++ b/vis4d/eval/nuscenes/detect3d.py @@ -258,7 +258,7 @@ def _process_detect_3d( annos.append(nusc_anno) self.detect_3d[token] = annos - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, tokens: list[str], boxes_3d: list[ArrayLike], diff --git a/vis4d/eval/nuscenes/track3d.py b/vis4d/eval/nuscenes/track3d.py index a1b113c69..77d7f8990 100644 --- a/vis4d/eval/nuscenes/track3d.py +++ b/vis4d/eval/nuscenes/track3d.py @@ -134,7 +134,7 @@ def _process_track_3d( annos.append(nusc_anno) self.tracks_3d[token] = annos - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, tokens: list[str], boxes_3d: list[ArrayLike], diff --git a/vis4d/eval/scalabel/base.py b/vis4d/eval/scalabel/base.py index 61a3e1534..d96f60990 100644 --- a/vis4d/eval/scalabel/base.py +++ b/vis4d/eval/scalabel/base.py @@ -60,12 +60,6 @@ def reset(self) -> None: """Reset the evaluator.""" self.frames = [] - def process_batch( # type: ignore # pragma: no cover - self, *args: Any, **kwargs: Any - ) -> None: - """Process sample and update confusion matrix.""" - raise NotImplementedError - def evaluate(self, metric: str) -> tuple[MetricLogs, str]: """Evaluate the dataset.""" raise NotImplementedError diff --git a/vis4d/eval/scalabel/detect.py b/vis4d/eval/scalabel/detect.py index dae776672..8e0a80444 100644 --- a/vis4d/eval/scalabel/detect.py +++ b/vis4d/eval/scalabel/detect.py @@ -44,7 +44,7 @@ def metrics(self) -> list[str]: """Supported metrics.""" return [self.METRICS_DET, self.METRICS_INS_SEG] - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, frame_ids: list[int], sample_names: list[str], diff --git a/vis4d/eval/scalabel/track.py b/vis4d/eval/scalabel/track.py index 7222b1eb9..fb80132c8 100644 --- a/vis4d/eval/scalabel/track.py +++ b/vis4d/eval/scalabel/track.py @@ -46,7 +46,7 @@ def metrics(self) -> list[str]: """Supported metrics.""" return [self.METRICS_TRACK, self.METRICS_SEG_TRACK] - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, frame_ids: list[int], sample_names: list[str], diff --git a/vis4d/eval/shift/multitask_writer.py b/vis4d/eval/shift/multitask_writer.py index 9d9187000..f94a01643 100644 --- a/vis4d/eval/shift/multitask_writer.py +++ b/vis4d/eval/shift/multitask_writer.py @@ -104,7 +104,7 @@ def _write_flow( """ raise NotImplementedError - def process_batch( # type: ignore # pylint: disable=arguments-differ + def process_batch( self, frame_ids: list[int], sample_names: list[str], diff --git a/vis4d/model/track3d/cc_3dt.py b/vis4d/model/track3d/cc_3dt.py index 1cc5e2319..972728ccd 100644 --- a/vis4d/model/track3d/cc_3dt.py +++ b/vis4d/model/track3d/cc_3dt.py @@ -1,7 +1,7 @@ """CC-3DT model implementation. This file composes the operations associated with CC-3DT -`https://arxiv.org/abs/2212.01247`_ into the full model implementation. +`https://arxiv.org/abs/2212.01247` into the full model implementation. """ from __future__ import annotations diff --git a/vis4d/op/base/dla.py b/vis4d/op/base/dla.py index 31d734494..98ff9d486 100644 --- a/vis4d/op/base/dla.py +++ b/vis4d/op/base/dla.py @@ -2,13 +2,21 @@ from __future__ import annotations +import math +from collections.abc import Sequence + import torch from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint + +from vis4d.common.ckpt import load_model_checkpoint from .base import BaseModel BN_MOMENTUM = 0.1 -DLA_MODEL_PREFIX = "http://dl.yf.io/dla/models/" + +DLA_MODEL_PREFIX = "http://dl.yf.io/dla/models/imagenet" + DLA_MODEL_MAPPING = { "dla34": "dla34-ba72cf86.pth", "dla46_c": "dla46_c-2bfd52c3.pth", @@ -21,6 +29,7 @@ "dla102x2": "dla102x2-262837b6.pth", "dla169": "dla169-0914e092.pth", } + DLA_ARCH_SETTINGS = { # pylint: disable=consider-using-namedtuple-or-dataclass "dla34": ( (1, 1, 1, 2, 2, 1), @@ -89,7 +98,12 @@ class BasicBlock(nn.Module): """BasicBlock.""" def __init__( - self, inplanes: int, planes: int, stride: int = 1, dilation: int = 1 + self, + inplanes: int, + planes: int, + stride: int = 1, + dilation: int = 1, + with_cp: bool = False, ) -> None: """Creates an instance of the class.""" super().__init__() @@ -115,22 +129,36 @@ def __init__( ) self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) self.stride = stride + self.with_cp = with_cp def forward( self, input_x: Tensor, residual: None | Tensor = None ) -> Tensor: """Forward.""" - if residual is None: - residual = input_x - out = self.conv1(input_x) - out = self.bn1(out) - out = self.relu(out) + def _inner_forward( + input_x: Tensor, residual: None | Tensor = None + ) -> Tensor: + if residual is None: + residual = input_x + out = self.conv1(input_x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + out += residual - out = self.conv2(out) - out = self.bn2(out) + return out + + if self.with_cp and input_x.requires_grad: + out = checkpoint( + _inner_forward, input_x, residual, use_reentrant=True + ) + else: + out = _inner_forward(input_x, residual) - out += residual out = self.relu(out) return out @@ -142,7 +170,12 @@ class Bottleneck(nn.Module): expansion = 2 def __init__( - self, inplanes: int, planes: int, stride: int = 1, dilation: int = 1 + self, + inplanes: int, + planes: int, + stride: int = 1, + dilation: int = 1, + with_cp: bool = False, ) -> None: """Creates an instance of the class.""" super().__init__() @@ -168,26 +201,41 @@ def __init__( self.bn3 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) self.relu = nn.ReLU(inplace=True) self.stride = stride + self.with_cp = with_cp def forward( self, input_x: Tensor, residual: None | Tensor = None ) -> Tensor: """Forward.""" - if residual is None: - residual = input_x - out = self.conv1(input_x) - out = self.bn1(out) - out = self.relu(out) + def _inner_forward( + input_x: Tensor, residual: None | Tensor = None + ) -> Tensor: + if residual is None: + residual = input_x - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) + out = self.conv1(input_x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + out += residual - out = self.conv3(out) - out = self.bn3(out) + return out + + if self.with_cp and input_x.requires_grad: + out = checkpoint( + _inner_forward, input_x, residual, use_reentrant=True + ) + else: + out = _inner_forward(input_x, residual) - out += residual out = self.relu(out) return out @@ -200,7 +248,12 @@ class BottleneckX(nn.Module): cardinality = 32 def __init__( - self, inplanes: int, planes: int, stride: int = 1, dilation: int = 1 + self, + inplanes: int, + planes: int, + stride: int = 1, + dilation: int = 1, + with_cp: bool = False, ) -> None: """Creates an instance of the class.""" super().__init__() @@ -227,26 +280,41 @@ def __init__( self.bn3 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM) self.relu = nn.ReLU(inplace=True) self.stride = stride + self.with_cp = with_cp def forward( self, input_x: Tensor, residual: None | Tensor = None ) -> Tensor: """Forward.""" - if residual is None: - residual = input_x - out = self.conv1(input_x) - out = self.bn1(out) - out = self.relu(out) + def _inner_forward( + input_x: Tensor, residual: None | Tensor = None + ) -> Tensor: + if residual is None: + residual = input_x - out = self.conv2(out) - out = self.bn2(out) - out = self.relu(out) + out = self.conv1(input_x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + out += residual - out = self.conv3(out) - out = self.bn3(out) + return out + + if self.with_cp and input_x.requires_grad: + out = checkpoint( + _inner_forward, input_x, residual, use_reentrant=True + ) + else: + out = _inner_forward(input_x, residual) - out += residual out = self.relu(out) return out @@ -261,6 +329,7 @@ def __init__( out_channels: int, kernel_size: int, residual: bool, + with_cp: bool = False, ) -> None: """Creates an instance of the class.""" super().__init__() @@ -272,17 +341,28 @@ def __init__( bias=False, padding=(kernel_size - 1) // 2, ) - self.bn1 = nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM) + self.bn = nn.BatchNorm2d( # pylint: disable=invalid-name + out_channels, momentum=BN_MOMENTUM + ) self.relu = nn.ReLU(inplace=True) self.residual = residual + self.with_cp = with_cp def forward(self, *input_x: Tensor) -> Tensor: """Forward.""" - children = input_x - feats = self.conv(torch.cat(input_x, 1)) - feats = self.bn1(feats) - if self.residual: - feats += children[0] + + def _inner_forward(*input_x: Tensor) -> Tensor: + feats = self.conv(torch.cat(input_x, 1)) + feats = self.bn(feats) + if self.residual: + feats += input_x[0] + return feats + + if self.with_cp and input_x[0].requires_grad: + feats = checkpoint(_inner_forward, *input_x, use_reentrant=True) + else: + feats = _inner_forward(*input_x) + feats = self.relu(feats) return feats @@ -303,6 +383,7 @@ def __init__( # pylint: disable=too-many-arguments root_kernel_size: int = 1, dilation: int = 1, root_residual: bool = False, + with_cp: bool = False, ) -> None: """Creates an instance of the class.""" super().__init__() @@ -320,13 +401,25 @@ def __init__( # pylint: disable=too-many-arguments root_dim += in_channels if levels == 1: self.tree1: Tree | BasicBlock = block_c( - in_channels, out_channels, stride, dilation=dilation + in_channels, + out_channels, + stride, + dilation=dilation, + with_cp=with_cp, ) self.tree2: Tree | BasicBlock = block_c( - out_channels, out_channels, 1, dilation=dilation + out_channels, + out_channels, + 1, + dilation=dilation, + with_cp=with_cp, ) self.root = Root( - root_dim, out_channels, root_kernel_size, root_residual + root_dim, + out_channels, + root_kernel_size, + root_residual, + with_cp=with_cp, ) else: self.tree1 = Tree( @@ -339,6 +432,7 @@ def __init__( # pylint: disable=too-many-arguments root_kernel_size=root_kernel_size, dilation=dilation, root_residual=root_residual, + with_cp=with_cp, ) self.tree2 = Tree( levels - 1, @@ -349,6 +443,7 @@ def __init__( # pylint: disable=too-many-arguments root_kernel_size=root_kernel_size, dilation=dilation, root_residual=root_residual, + with_cp=with_cp, ) self.level_root = level_root self.root_dim = root_dim @@ -369,7 +464,7 @@ def __init__( # pylint: disable=too-many-arguments stride=1, bias=False, ), - nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM), + nn.BatchNorm2d(out_channels), ) def forward( @@ -399,32 +494,21 @@ class DLA(BaseModel): def __init__( self, - name: None | str = None, - levels: tuple[int, int, int, int, int, int] = (1, 1, 1, 2, 2, 1), - channels: tuple[int, int, int, int, int, int] = ( - 16, - 32, - 64, - 128, - 256, - 512, - ), - block: str = "BasicBlock", - residual_root: bool = False, - cardinality: int = 32, + name: str, + out_indices: Sequence[int] = (0, 1, 2, 3), + with_cp: bool = False, + pretrained: bool = False, weights: None | str = None, - style: str = "imagenet", ) -> None: """Creates an instance of the class.""" super().__init__() - if name is not None: - assert name in DLA_ARCH_SETTINGS - arch_setting = DLA_ARCH_SETTINGS[name] - levels, channels, residual_root, block = arch_setting - if name == "dla102x2": # pragma: no cover - BottleneckX.cardinality = 64 - else: - BottleneckX.cardinality = cardinality + assert name in DLA_ARCH_SETTINGS, f"{name} is not supported!" + + levels, channels, residual_root, block = DLA_ARCH_SETTINGS[name] + + if name == "dla102x2": # pragma: no cover + BottleneckX.cardinality = 64 + self.base_layer = nn.Sequential( nn.Conv2d( 3, channels[0], kernel_size=7, stride=1, padding=3, bias=False @@ -446,6 +530,7 @@ def __init__( 2, level_root=False, root_residual=residual_root, + with_cp=with_cp, ) self.level3 = Tree( levels[3], @@ -455,6 +540,7 @@ def __init__( 2, level_root=True, root_residual=residual_root, + with_cp=with_cp, ) self.level4 = Tree( levels[4], @@ -464,6 +550,7 @@ def __init__( 2, level_root=True, root_residual=residual_root, + with_cp=with_cp, ) self.level5 = Tree( levels[5], @@ -473,19 +560,30 @@ def __init__( 2, level_root=True, root_residual=residual_root, + with_cp=with_cp, ) - self._out_channels = list(channels) + self.out_indices = out_indices + self._out_channels = [channels[i + 2] for i in out_indices] - if weights is not None: # pragma: no cover - if weights.startswith("dla://"): - weights_name = weights.split("dla://")[-1] - assert weights_name in DLA_MODEL_MAPPING - weights = ( - f"{DLA_MODEL_PREFIX}{style}/" - f"{DLA_MODEL_MAPPING[weights_name]}" - ) - self.load_pretrained_model(weights) + if pretrained: + if weights is None: # pragma: no cover + weights = f"{DLA_MODEL_PREFIX}/{DLA_MODEL_MAPPING[name]}" + + load_model_checkpoint(self, weights) + + else: + self._init_weights() + + def _init_weights(self) -> None: + """Initialize module weights.""" + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2.0 / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() @staticmethod def _make_conv_level( @@ -516,14 +614,6 @@ def _make_conv_level( inplanes = planes return nn.Sequential(*modules) - def load_pretrained_model(self, weights: str) -> None: - """Load pretrained weights.""" - if weights.startswith("http://") or weights.startswith("https://"): - model_weights = torch.hub.load_state_dict_from_url(weights) - else: # pragma: no cover - model_weights = torch.load(weights) - self.load_state_dict(model_weights, strict=False) - def forward(self, images: Tensor) -> list[Tensor]: """DLA forward. @@ -534,14 +624,17 @@ def forward(self, images: Tensor) -> list[Tensor]: Returns: fp (list[Tensor]): The output feature pyramid. The list index represents the level, which has a downsampling raio of 2^index. - fp[0] is a feature map with the image resolution instead of the - original image. """ input_x = self.base_layer(images) - outs: list[Tensor] = [] + + outs = [images, images] + for i in range(6): input_x = getattr(self, f"level{i}")(input_x) - outs.append(input_x) + + if i - 2 in self.out_indices: + outs.append(input_x) + return outs @property @@ -551,4 +644,4 @@ def out_channels(self) -> list[int]: Returns: list[int]: number of channels """ - return self._out_channels + return [3, 3] + self._out_channels diff --git a/vis4d/op/base/pointnet.py b/vis4d/op/base/pointnet.py index 67588cae9..39eb638c8 100644 --- a/vis4d/op/base/pointnet.py +++ b/vis4d/op/base/pointnet.py @@ -101,7 +101,7 @@ def __init__( self.activation_ = getattr(nn, activation_cls)() # Create norms - norm_fn: Callable[[int], nn.Module] = ( + norm_fn: Callable[[int], nn.Module] | None = ( getattr(nn, norm_cls) if norm_cls is not None else None ) @@ -249,7 +249,7 @@ def __init__( activation = getattr(nn, activation_cls)() # Create norms - norm_fn: Callable[[int], nn.Module] = ( + norm_fn: Callable[[int], nn.Module] | None = ( getattr(nn, norm_cls) if norm_cls is not None else None ) diff --git a/vis4d/op/base/pointnetpp.py b/vis4d/op/base/pointnetpp.py index 0cc40f280..d3f4105fa 100644 --- a/vis4d/op/base/pointnetpp.py +++ b/vis4d/op/base/pointnetpp.py @@ -229,13 +229,13 @@ def __init__( last_channel = in_channel # Create norms - norm_fn: Callable[[int], nn.Module] = ( + norm_fn: Callable[[int], nn.Module] | None = ( getattr(nn, norm_cls) if norm_cls is not None else None ) for out_channel in mlp: self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1)) - if norm_cls is not None: + if norm_fn is not None: self.mlp_bns.append(norm_fn(out_channel)) last_channel = out_channel self.group_all = group_all diff --git a/vis4d/op/base/resnet.py b/vis4d/op/base/resnet.py index 0e00bb026..10e9f51f3 100644 --- a/vis4d/op/base/resnet.py +++ b/vis4d/op/base/resnet.py @@ -81,7 +81,7 @@ def _inner_forward(x: Tensor) -> Tensor: return out if self.use_checkpoint and x.requires_grad: - out = checkpoint(_inner_forward, x) + out = checkpoint(_inner_forward, x, use_reentrant=True) else: out = _inner_forward(x) @@ -183,7 +183,7 @@ def _inner_forward(x: Tensor) -> Tensor: return out if self.use_checkpoint and x.requires_grad: - out = checkpoint(_inner_forward, x) + out = checkpoint(_inner_forward, x, use_reentrant=True) else: out = _inner_forward(x) diff --git a/vis4d/op/base/vit.py b/vis4d/op/base/vit.py index fdbdbd059..0948d2a05 100644 --- a/vis4d/op/base/vit.py +++ b/vis4d/op/base/vit.py @@ -3,7 +3,7 @@ from __future__ import annotations import torch -from timm.models.helpers import named_apply +from timm.models import named_apply from torch import nn from ..layer import PatchEmbed, TransformerBlock @@ -19,7 +19,7 @@ def _init_weights_vit_timm( # pylint: disable=unused-argument if module.bias is not None: nn.init.zeros_(module.bias) elif hasattr(module, "init_weights"): - module.init_weights() + module.init_weights() # type: ignore ViT_PRESET = { # pylint: disable=consider-using-namedtuple-or-dataclass diff --git a/vis4d/op/box/box2d.py b/vis4d/op/box/box2d.py index 3359c21ed..906f87643 100644 --- a/vis4d/op/box/box2d.py +++ b/vis4d/op/box/box2d.py @@ -420,7 +420,7 @@ def multiclass_nms( labels = torch.arange(num_classes, dtype=torch.long, device=scores.device) labels = labels.view(1, -1).expand_as(scores) - bboxes = bboxes.view(-1, 4) + bboxes = bboxes.reshape(-1, 4) scores = scores.reshape(-1) labels = labels.reshape(-1) diff --git a/vis4d/op/box/poolers/__init__.py b/vis4d/op/box/poolers/__init__.py index d6458ee94..0ef125f3e 100644 --- a/vis4d/op/box/poolers/__init__.py +++ b/vis4d/op/box/poolers/__init__.py @@ -1,6 +1,15 @@ """Init sampler module.""" from .base import RoIPooler -from .roi_pooler import MultiScaleRoIAlign, MultiScaleRoIPool +from .roi_pooler import ( + MultiScaleRoIAlign, + MultiScaleRoIPool, + MultiScaleRoIPooler, +) -__all__ = ["RoIPooler", "MultiScaleRoIAlign", "MultiScaleRoIPool"] +__all__ = [ + "RoIPooler", + "MultiScaleRoIAlign", + "MultiScaleRoIPool", + "MultiScaleRoIPooler", +] diff --git a/vis4d/op/detect/dense_anchor.py b/vis4d/op/detect/dense_anchor.py index 14f8a20d8..74816cbc8 100644 --- a/vis4d/op/detect/dense_anchor.py +++ b/vis4d/op/detect/dense_anchor.py @@ -29,7 +29,7 @@ class DetectorTargets(NamedTuple): def images_to_levels( targets: list[ tuple[list[Tensor], list[Tensor], list[Tensor], list[Tensor]] - ] + ], ) -> list[list[Tensor]]: """Convert targets by image to targets by feature level.""" targets_per_level = [] diff --git a/vis4d/op/detect3d/qd_3dt.py b/vis4d/op/detect3d/qd_3dt.py index 6d11b167a..1891f4621 100644 --- a/vis4d/op/detect3d/qd_3dt.py +++ b/vis4d/op/detect3d/qd_3dt.py @@ -11,7 +11,7 @@ from vis4d.common.typing import LossesType from vis4d.op.box.encoder.qd_3dt import QD3DTBox3DDecoder, QD3DTBox3DEncoder from vis4d.op.box.matchers import Matcher, MaxIoUMatcher -from vis4d.op.box.poolers import MultiScaleRoIAlign, RoIPooler +from vis4d.op.box.poolers import MultiScaleRoIAlign, MultiScaleRoIPooler from vis4d.op.box.samplers import ( CombinedSampler, Sampler, @@ -47,7 +47,7 @@ class QD3DTDet3DOut(NamedTuple): depth_uncertainty: list[Tensor] -def get_default_proposal_pooler() -> RoIPooler: +def get_default_proposal_pooler() -> MultiScaleRoIAlign: """Get default proposal pooler of QD-3DT bounding box 3D head.""" return MultiScaleRoIAlign( resolution=[7, 7], strides=[4, 8, 16, 32], sampling_ratio=0 @@ -101,10 +101,10 @@ def get_default_box_codec( class QD3DTBBox3DHead(nn.Module): """This class implements the QD-3DT bounding box 3D head.""" - def __init__( # pylint: disable=too-many-arguments + def __init__( # pylint: disable=too-many-arguments, too-many-positional-arguments, line-too-long self, num_classes: int, - proposal_pooler: None | RoIPooler = None, + proposal_pooler: None | MultiScaleRoIPooler = None, box_matcher: None | Matcher = None, box_sampler: None | Sampler = None, box_encoder: None | QD3DTBox3DEncoder = None, diff --git a/vis4d/op/fpp/fpn.py b/vis4d/op/fpp/fpn.py index a2856e116..0433437bb 100644 --- a/vis4d/op/fpp/fpn.py +++ b/vis4d/op/fpp/fpn.py @@ -14,7 +14,9 @@ from torchvision.ops.feature_pyramid_network import ( ExtraFPNBlock as _ExtraFPNBlock, ) -from torchvision.ops.feature_pyramid_network import LastLevelMaxPool +from torchvision.ops.feature_pyramid_network import ( + LastLevelMaxPool, +) from .base import FeaturePyramidProcessing diff --git a/vis4d/op/layer/attention.py b/vis4d/op/layer/attention.py index 4a539ee3e..f3ba881e5 100644 --- a/vis4d/op/layer/attention.py +++ b/vis4d/op/layer/attention.py @@ -100,6 +100,7 @@ def __init__( proj_drop: float = 0.0, dropout_layer: nn.Module | None = None, batch_first: bool = False, + need_weights: bool = False, **kwargs: ArgsType, ) -> None: """Init MultiheadAttention. @@ -116,10 +117,16 @@ def __init__( batch_first (bool): When it is True, Key, Query and Value are shape of (batch, n, embed_dim), otherwise (n, batch, embed_dim). Default to False. + need_weights (bool): Whether to return the attention weights. + If True, the output will be a tuple of (attn_output, + attn_output_weights) and not using FlashAttention. If False, + only the attn_output will be returned. Default to False. """ super().__init__() self.batch_first = batch_first self.embed_dims = embed_dims + self.num_heads = num_heads + self.need_weights = need_weights self.attn = nn.MultiheadAttention( embed_dims, num_heads, dropout=attn_drop, **kwargs @@ -193,8 +200,10 @@ def forward( key_pos = query_pos else: rank_zero_warn( - "position encoding of key is" - + f"missing in {self.__class__.__name__}." + f"Position encoding of key in {self.__class__.__name__}" + + "is missing, and positional encodeing of query has " + + "has different shape and cannot be usde for key. " + + "It it is not desired, please provide key_pos." ) if query_pos is not None: @@ -220,7 +229,11 @@ def forward( value=value, attn_mask=attn_mask, key_padding_mask=key_padding_mask, - )[0] + need_weights=self.need_weights, + ) + + if isinstance(out, tuple): + out = out[0] if self.batch_first: out = out.transpose(0, 1) diff --git a/vis4d/op/layer/ms_deform_attn.py b/vis4d/op/layer/ms_deform_attn.py index 9acb05566..1ee8a7532 100644 --- a/vis4d/op/layer/ms_deform_attn.py +++ b/vis4d/op/layer/ms_deform_attn.py @@ -61,7 +61,7 @@ def forward( # type: ignore @staticmethod @once_differentiable # type: ignore - def backward( + def backward( # type: ignore ctx, grad_output: Tensor ) -> tuple[Tensor, None, None, Tensor, Tensor, None]: """Backward pass.""" @@ -223,12 +223,15 @@ def __init__( is_power_of_2(d_model // n_heads) self.d_model = d_model - self.embed_dims = d_model self.n_levels = n_levels self.n_heads = n_heads self.n_points = n_points self.im2col_step = im2col_step + # Aligned Attributes to MHA + self.embed_dims = d_model + self.num_heads = n_heads + self.sampling_offsets = nn.Linear( d_model, n_heads * n_levels * n_points * 2 ) @@ -359,3 +362,22 @@ def forward( output = self.output_proj(output) return output + + def __call__( + self, + query: Tensor, + reference_points: Tensor, + input_flatten: Tensor, + input_spatial_shapes: Tensor, + input_level_start_index: Tensor, + input_padding_mask: Tensor | None = None, + ) -> Tensor: + """Type definition for call implementation.""" + return self._call_impl( + query, + reference_points, + input_flatten, + input_spatial_shapes, + input_level_start_index, + input_padding_mask, + ) diff --git a/vis4d/op/layer/patch_embed.py b/vis4d/op/layer/patch_embed.py index 5b91aa123..ba775a63c 100644 --- a/vis4d/op/layer/patch_embed.py +++ b/vis4d/op/layer/patch_embed.py @@ -1,6 +1,6 @@ """Image to Patch Embedding using Conv2d. -Modified from vision_transformer +Modified from vision_transformer (https://github.com/google-research/vision_transformer). """ diff --git a/vis4d/op/layer/positional_encoding.py b/vis4d/op/layer/positional_encoding.py index c931c4b5f..37ccfa7c0 100644 --- a/vis4d/op/layer/positional_encoding.py +++ b/vis4d/op/layer/positional_encoding.py @@ -3,6 +3,8 @@ Modified from mmdetection (https://github.com/open-mmlab/mmdetection). """ +from __future__ import annotations + import math import torch @@ -59,24 +61,45 @@ def __init__( self.eps = eps self.offset = offset - def forward(self, mask: Tensor) -> Tensor: + def forward( + self, mask: Tensor | None, inputs: Tensor | None = None + ) -> Tensor: """Forward function for `SinePositionalEncoding`. Args: - mask (Tensor): ByteTensor mask. Non-zero values representing + mask (Tensor | None): ByteTensor mask. Non-zero values representing ignored positions, while zero values means valid positions - for this image. Shape [bs, h, w]. + for this image. Shape [bs, h, w]. If None, it means single + image or batch image with no padding. + inputs (Tensor | None): The input tensor. It mask is None, this + input tensor is required to get the shape of the input image. Returns: pos (Tensor): Returned position embedding with shape [bs, num_feats*2, h, w]. """ - # For convenience of exporting to ONNX, it's required to convert - # `masks` from bool to int. - mask = mask.to(torch.int) - not_mask = 1 - mask # logical_not - y_embed = not_mask.cumsum(1, dtype=torch.float32) - x_embed = not_mask.cumsum(2, dtype=torch.float32) + if mask is not None: + # For convenience of exporting to ONNX, it's required to convert + # `masks` from bool to int. + mask = mask.to(torch.int) + b, h, w = mask.size() + device = mask.device + not_mask = 1 - mask # logical_not + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + else: + # single image or batch image with no padding + assert isinstance(inputs, Tensor) + b, _, h, w = inputs.shape + device = inputs.device + x_embed = torch.arange( + 1, w + 1, dtype=torch.float32, device=device + ) + x_embed = x_embed.view(1, 1, -1).repeat(b, h, 1) + y_embed = torch.arange( + 1, h + 1, dtype=torch.float32, device=device + ) + y_embed = y_embed.view(1, -1, 1).repeat(b, 1, w) if self.normalize: y_embed = ( (y_embed + self.offset) @@ -89,13 +112,13 @@ def forward(self, mask: Tensor) -> Tensor: * self.scale ) dim_t = torch.arange( - self.num_feats, dtype=torch.float32, device=mask.device + self.num_feats, dtype=torch.float32, device=device ) dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_feats) pos_x = x_embed[:, :, :, None] / dim_t pos_y = y_embed[:, :, :, None] / dim_t # use `view` instead of `flatten` for dynamically exporting to ONNX - b, h, w = mask.size() + pos_x = torch.stack( (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 ).view(b, h, w, -1) @@ -167,78 +190,3 @@ def forward(self, mask: Tensor) -> Tensor: .repeat(mask.shape[0], 1, 1, 1) ) return pos - - -class SinePositionalEncoding3D(SinePositionalEncoding): - """3D Position encoding with sine and cosine functions.""" - - def forward(self, mask: Tensor) -> Tensor: - """Forward function for `SinePositionalEncoding3D`. - - Args: - mask (Tensor): ByteTensor mask. Non-zero values representing - ignored positions, while zero values means valid positions - for this image. Shape [bs, t, h, w]. - - Returns: - pos (Tensor): Returned position embedding with shape - [bs, num_feats*2, h, w]. - """ - assert mask.dim() == 4, ( - f"{mask.shape} should be a 4-dimensional Tensor," - f" got {mask.dim()}-dimensional Tensor instead " - ) - # For convenience of exporting to ONNX, it's required to convert - # `masks` from bool to int. - mask = mask.to(torch.int) - not_mask = 1 - mask # logical_not - z_embed = not_mask.cumsum(1, dtype=torch.float32) - y_embed = not_mask.cumsum(2, dtype=torch.float32) - x_embed = not_mask.cumsum(3, dtype=torch.float32) - if self.normalize: - z_embed = ( - (z_embed + self.offset) - / (z_embed[:, -1:, :, :] + self.eps) - * self.scale - ) - y_embed = ( - (y_embed + self.offset) - / (y_embed[:, :, -1:, :] + self.eps) - * self.scale - ) - x_embed = ( - (x_embed + self.offset) - / (x_embed[:, :, :, -1:] + self.eps) - * self.scale - ) - dim_t = torch.arange( - self.num_feats, dtype=torch.float32, device=mask.device - ) - dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_feats) - - dim_t_z = torch.arange( - (self.num_feats * 2), dtype=torch.float32, device=mask.device - ) - dim_t_z = self.temperature ** ( - 2 * (dim_t_z // 2) / (self.num_feats * 2) - ) - - pos_x = x_embed[:, :, :, :, None] / dim_t - pos_y = y_embed[:, :, :, :, None] / dim_t - pos_z = z_embed[:, :, :, :, None] / dim_t_z - # use `view` instead of `flatten` for dynamically exporting to ONNX - b, t, h, w = mask.size() - pos_x = torch.stack( - (pos_x[:, :, :, :, 0::2].sin(), pos_x[:, :, :, :, 1::2].cos()), - dim=5, - ).view(b, t, h, w, -1) - pos_y = torch.stack( - (pos_y[:, :, :, :, 0::2].sin(), pos_y[:, :, :, :, 1::2].cos()), - dim=5, - ).view(b, t, h, w, -1) - pos_z = torch.stack( - (pos_z[:, :, :, :, 0::2].sin(), pos_z[:, :, :, :, 1::2].cos()), - dim=5, - ).view(b, t, h, w, -1) - pos = (torch.cat((pos_y, pos_x), dim=4) + pos_z).permute(0, 1, 4, 2, 3) - return pos diff --git a/vis4d/op/layer/transformer.py b/vis4d/op/layer/transformer.py index 08c7950dd..2600ade3d 100644 --- a/vis4d/op/layer/transformer.py +++ b/vis4d/op/layer/transformer.py @@ -212,6 +212,8 @@ def __init__( LayerScale. Default: 0.0 """ super().__init__() + self.embed_dims = embed_dims + layers: list[nn.Module] = [] in_channels = embed_dims for _ in range(num_fcs - 1): diff --git a/vis4d/op/mask/util.py b/vis4d/op/mask/util.py index 02a768b78..4e0b299ca 100644 --- a/vis4d/op/mask/util.py +++ b/vis4d/op/mask/util.py @@ -8,7 +8,7 @@ from torch import Tensor -def _do_paste_mask( +def _do_paste_mask( # type: ignore masks: Tensor, boxes: Tensor, img_h: int, diff --git a/vis4d/op/track/assignment.py b/vis4d/op/track/assignment.py index dd000fe2e..4136c4161 100644 --- a/vis4d/op/track/assignment.py +++ b/vis4d/op/track/assignment.py @@ -3,16 +3,18 @@ from __future__ import annotations import torch +from scipy.optimize import linear_sum_assignment +from torch import Tensor def greedy_assign( - detection_scores: torch.Tensor, - tracklet_ids: torch.Tensor, - affinity_scores: torch.Tensor, + detection_scores: Tensor, + tracklet_ids: Tensor, + affinity_scores: Tensor, match_score_thr: float = 0.5, obj_score_thr: float = 0.3, nms_conf_thr: None | float = None, -) -> torch.Tensor: +) -> Tensor: """Greedy assignment of detections to tracks given affinities.""" ids = torch.full( (len(detection_scores),), @@ -35,6 +37,40 @@ def greedy_assign( return ids +def hungarian_assign( + detection_scores: Tensor, + tracklet_ids: Tensor, + affinity_scores: Tensor, + match_score_thr: float = 0.5, + obj_score_thr: float = 0.3, + nms_conf_thr: None | float = None, +) -> Tensor: + """Hungarian assignment of detections to tracks given affinities.""" + ids = torch.full( + (len(detection_scores),), + -1, + dtype=torch.long, + device=detection_scores.device, + ) + + matched_indices = linear_sum_assignment(-affinity_scores.cpu().numpy()) + + for idx in range(len(matched_indices[0])): + i = matched_indices[0][idx] + memo_ind = matched_indices[1][idx] + conf = affinity_scores[i, memo_ind] + tid = tracklet_ids[memo_ind] + if conf > match_score_thr and tid > -1: + if detection_scores[i] > obj_score_thr: + ids[i] = tid + affinity_scores[:i, memo_ind] = 0 + affinity_scores[i + 1 :, memo_ind] = 0 + elif nms_conf_thr is not None and conf > nms_conf_thr: + ids[i] = -2 + + return ids + + class TrackIDCounter: """Global counter for track ids. @@ -52,7 +88,7 @@ def reset(cls) -> None: @classmethod def get_ids( cls, num_ids: int, device: torch.device = torch.device("cpu") - ) -> torch.Tensor: + ) -> Tensor: """Generate a num_ids number of new unique tracking ids. Args: @@ -61,7 +97,7 @@ def get_ids( to torch.device("cpu"). Returns: - torch.Tensor: Tensor of new contiguous track ids. + Tensor: Tensor of new contiguous track ids. """ new_ids = torch.arange(cls.count, cls.count + num_ids, device=device) cls.count = cls.count + num_ids diff --git a/vis4d/op/track/qdtrack.py b/vis4d/op/track/qdtrack.py index 480161b93..bbc46c2f2 100644 --- a/vis4d/op/track/qdtrack.py +++ b/vis4d/op/track/qdtrack.py @@ -10,7 +10,7 @@ from vis4d.op.box.box2d import bbox_iou from vis4d.op.box.matchers.max_iou import MaxIoUMatcher -from vis4d.op.box.poolers import MultiScaleRoIAlign, RoIPooler +from vis4d.op.box.poolers import MultiScaleRoIAlign, MultiScaleRoIPooler from vis4d.op.box.samplers import CombinedSampler, match_and_sample_proposals from vis4d.op.layer import add_conv_branch from vis4d.op.loss import EmbeddingDistanceLoss, MultiPosCrossEntropyLoss @@ -340,7 +340,7 @@ class QDSimilarityHead(nn.Module): def __init__( self, - proposal_pooler: None | RoIPooler = None, + proposal_pooler: None | MultiScaleRoIPooler = None, in_dim: int = 256, num_convs: int = 4, conv_out_dim: int = 256, @@ -355,8 +355,8 @@ def __init__( """Creates an instance of the class. Args: - proposal_pooler (None | RoIPooler, optional): RoI pooling module. - Defaults to None. + proposal_pooler (None | MultiScaleRoIPooler, optional): RoI pooling + module. Defaults to None. in_dim (int, optional): Input feature dimension. Defaults to 256. num_convs (int, optional): Number of convolutional layers inside the head. Defaults to 4. diff --git a/vis4d/pl/__init__.py b/vis4d/pl/__init__.py deleted file mode 100644 index 50e47fe37..000000000 --- a/vis4d/pl/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Run training and evaluation in PyTorch Lightening.""" diff --git a/vis4d/pl/__main__.py b/vis4d/pl/__main__.py deleted file mode 100644 index 67e5297f9..000000000 --- a/vis4d/pl/__main__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Entry point for the vis4d.pl package.""" - -from .run import entrypoint - -entrypoint() diff --git a/vis4d/pl/callbacks/__init__.py b/vis4d/pl/callbacks/__init__.py deleted file mode 100644 index 36c3ea271..000000000 --- a/vis4d/pl/callbacks/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Callbacks handling data related stuff (evaluation, visualization, etc).""" - -from .callback_wrapper import CallbackWrapper -from .scheduler import LRSchedulerCallback - -__all__ = ["CallbackWrapper", "LRSchedulerCallback"] diff --git a/vis4d/pl/callbacks/callback_wrapper.py b/vis4d/pl/callbacks/callback_wrapper.py deleted file mode 100644 index 1aa5e0640..000000000 --- a/vis4d/pl/callbacks/callback_wrapper.py +++ /dev/null @@ -1,232 +0,0 @@ -"""Wrapper to connect PyTorch Lightning callbacks.""" - -from __future__ import annotations - -from typing import Any - -import lightning.pytorch as pl -from torch import nn - -from vis4d.engine.callbacks import Callback, TrainerState -from vis4d.engine.loss_module import LossModule -from vis4d.pl.training_module import TrainingModule - - -def get_trainer_state( - trainer: pl.Trainer, pl_module: pl.LightningModule, val: bool = False -) -> TrainerState: - """Wrap pl.Trainer and pl.LightningModule into Trainer.""" - # Resolve float("inf") to -1 - if val: - test_dataloader = trainer.val_dataloaders - num_test_batches = [ - num_batch if isinstance(num_batch, int) else -1 - for num_batch in trainer.num_val_batches - ] - else: - test_dataloader = trainer.test_dataloaders - num_test_batches = [ - num_batch if isinstance(num_batch, int) else -1 - for num_batch in trainer.num_test_batches - ] - - # Map max_epochs=None to -1 - if trainer.max_epochs is None: - num_epochs = -1 - else: - num_epochs = trainer.max_epochs - - # Resolve float("inf") to -1 - if isinstance(trainer.num_training_batches, float): - num_train_batches = -1 - else: - num_train_batches = trainer.num_training_batches - - return TrainerState( - current_epoch=pl_module.current_epoch, - num_epochs=num_epochs, - global_step=trainer.global_step, - num_steps=trainer.max_steps, - train_dataloader=trainer.train_dataloader, - num_train_batches=num_train_batches, - test_dataloader=test_dataloader, - num_test_batches=num_test_batches, - train_module=trainer, - train_engine="pl", - ) - - -def get_model(model: pl.LightningModule) -> nn.Module: - """Get model from pl module.""" - if isinstance(model, TrainingModule): - return model.model - return model - - -def get_loss_module(loss_module: pl.LightningModule) -> LossModule: - """Get loss_module from pl module.""" - if isinstance(loss_module, TrainingModule): - assert loss_module.loss_module is not None - return loss_module.loss_module - return loss_module # type: ignore - - -class CallbackWrapper(pl.Callback): - """Wrapper to connect vis4d callbacks to pytorch lightning callbacks.""" - - def __init__(self, callback: Callback) -> None: - """Init class.""" - self.callback = callback - - def setup( - self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str - ) -> None: - """Setup callback.""" - self.callback.setup() - - def on_train_batch_start( # type: ignore - self, - trainer: pl.Trainer, - pl_module: pl.LightningModule, - batch: Any, - batch_idx: int, - ) -> None: - """Called when the train batch begins.""" - trainer_state = get_trainer_state(trainer, pl_module) - - self.callback.on_train_batch_start( - trainer_state=trainer_state, - model=get_model(pl_module), - loss_module=get_loss_module(pl_module), - batch=batch, - batch_idx=batch_idx, - ) - - def on_train_epoch_start( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - """Hook to run at the start of a training epoch.""" - self.callback.on_train_epoch_start( - get_trainer_state(trainer, pl_module), - get_model(pl_module), - get_loss_module(pl_module), - ) - - def on_train_batch_end( # type: ignore - self, - trainer: pl.Trainer, - pl_module: pl.LightningModule, - outputs: Any, - batch: Any, - batch_idx: int, - ) -> None: - """Hook to run at the end of a training batch.""" - trainer_state = get_trainer_state(trainer, pl_module) - trainer_state["metrics"] = outputs["metrics"] - - log_dict = self.callback.on_train_batch_end( - trainer_state=trainer_state, - model=get_model(pl_module), - loss_module=get_loss_module(pl_module), - outputs=outputs["predictions"], - batch=batch, - batch_idx=batch_idx, - ) - - if log_dict is not None: - for k, v in log_dict.items(): - pl_module.log(f"train/{k}", v, rank_zero_only=True) - - def on_train_epoch_end( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - """Hook to run at the end of a training epoch.""" - self.callback.on_train_epoch_end( - get_trainer_state(trainer, pl_module), - get_model(pl_module), - get_loss_module(pl_module), - ) - - def on_validation_epoch_start( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - """Hook to run at the start of a validation epoch.""" - self.callback.on_test_epoch_start( - get_trainer_state(trainer, pl_module, val=True), - get_model(pl_module), - ) - - def on_validation_batch_end( # type: ignore - self, - trainer: pl.Trainer, - pl_module: pl.LightningModule, - outputs: Any, - batch: Any, - batch_idx: int, - dataloader_idx: int = 0, - ) -> None: - """Wait for on_validation_batch_end PL hook to call 'process'.""" - self.callback.on_test_batch_end( - trainer_state=get_trainer_state(trainer, pl_module, val=True), - model=get_model(pl_module), - outputs=outputs, - batch=batch, - batch_idx=batch_idx, - dataloader_idx=dataloader_idx, - ) - - def on_validation_epoch_end( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - """Wait for on_validation_epoch_end PL hook to call 'evaluate'.""" - log_dict = self.callback.on_test_epoch_end( - get_trainer_state(trainer, pl_module, val=True), - get_model(pl_module), - ) - - if log_dict is not None: - for k, v in log_dict.items(): - pl_module.log( - f"val/{k}", v, sync_dist=True, rank_zero_only=True - ) - - def on_test_epoch_start( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - """Hook to run at the start of a testing epoch.""" - self.callback.on_test_epoch_start( - get_trainer_state(trainer, pl_module), get_model(pl_module) - ) - - def on_test_batch_end( # type: ignore - self, - trainer: pl.Trainer, - pl_module: pl.LightningModule, - outputs: Any, - batch: Any, - batch_idx: int, - dataloader_idx: int = 0, - ) -> None: - """Wait for on_test_batch_end PL hook to call 'process'.""" - self.callback.on_test_batch_end( - trainer_state=get_trainer_state(trainer, pl_module), - model=get_model(pl_module), - outputs=outputs, - batch=batch, - batch_idx=batch_idx, - dataloader_idx=dataloader_idx, - ) - - def on_test_epoch_end( - self, trainer: pl.Trainer, pl_module: pl.LightningModule - ) -> None: - """Wait for on_test_epoch_end PL hook to call 'evaluate'.""" - log_dict = self.callback.on_test_epoch_end( - get_trainer_state(trainer, pl_module), get_model(pl_module) - ) - - if log_dict is not None: - for k, v in log_dict.items(): - pl_module.log( - f"test/{k}", v, sync_dist=True, rank_zero_only=True - ) diff --git a/vis4d/pl/run.py b/vis4d/pl/run.py deleted file mode 100644 index e2f3f1b36..000000000 --- a/vis4d/pl/run.py +++ /dev/null @@ -1,157 +0,0 @@ -"""CLI interface using PyTorch Lightning.""" - -from __future__ import annotations - -import logging -import os.path as osp - -import torch -from absl import app # pylint: disable=no-name-in-module -from lightning.fabric.utilities.exceptions import MisconfigurationException -from lightning.pytorch import Callback -from torch.utils.collect_env import get_pretty_env_info - -from vis4d.common import ArgsType -from vis4d.common.logging import dump_config, rank_zero_info, setup_logger -from vis4d.common.util import set_tf32 -from vis4d.config import instantiate_classes -from vis4d.config.typing import ExperimentConfig -from vis4d.engine.callbacks import CheckpointCallback -from vis4d.engine.flag import _CKPT, _CONFIG, _GPUS, _RESUME, _SHOW_CONFIG -from vis4d.engine.parser import pprints_config -from vis4d.pl.callbacks import CallbackWrapper, LRSchedulerCallback -from vis4d.pl.data_module import DataModule -from vis4d.pl.trainer import PLTrainer -from vis4d.pl.training_module import TrainingModule - - -def main(argv: ArgsType) -> None: - """Main entry point for the CLI. - - Example to run this script: - >>> python -m vis4d.pl.run fit --config configs/faster_rcnn/faster_rcnn_coco.py - """ - # Get config - mode = argv[1] - assert mode in {"fit", "test"}, f"Invalid mode: {mode}" - config: ExperimentConfig = _CONFIG.value - num_gpus = _GPUS.value - - # Setup logging - logger_vis4d = logging.getLogger("vis4d") - logger_pl = logging.getLogger("pytorch_lightning") - log_file = osp.join(config.output_dir, f"log_{config.timestamp}.txt") - setup_logger(logger_vis4d, log_file) - setup_logger(logger_pl, log_file) - - # Dump config - config_file = osp.join( - config.output_dir, f"config_{config.timestamp}.yaml" - ) - dump_config(config, config_file) - - rank_zero_info("Environment info: %s", get_pretty_env_info()) - - # PyTorch Setting - set_tf32(config.use_tf32, config.tf32_matmul_precision) - torch.hub.set_dir(f"{config.work_dir}/.cache/torch/hub") - - # Setup device - if num_gpus > 0: - config.pl_trainer.accelerator = "gpu" - config.pl_trainer.devices = num_gpus - else: - config.pl_trainer.accelerator = "cpu" - config.pl_trainer.devices = 1 - - trainer_args = instantiate_classes(config.pl_trainer).to_dict() - - if _SHOW_CONFIG.value: - rank_zero_info(pprints_config(config)) - - # Instantiate classes - if mode == "fit": - train_data_connector = instantiate_classes(config.train_data_connector) - loss = instantiate_classes(config.loss) - else: - train_data_connector = None - loss = None - - if config.test_data_connector is not None: - test_data_connector = instantiate_classes(config.test_data_connector) - else: - test_data_connector = None - - # Callbacks - callbacks: list[Callback] = [] - for cb in config.callbacks: - callback = instantiate_classes(cb) - # Skip checkpoint callback to use PL ModelCheckpoint - if not isinstance(callback, CheckpointCallback): - callbacks.append(CallbackWrapper(callback)) - - if "pl_callbacks" in config: - pl_callbacks = [instantiate_classes(cb) for cb in config.pl_callbacks] - else: - pl_callbacks = [] - - for cb in pl_callbacks: - if not isinstance(cb, Callback): - raise MisconfigurationException( - "Callback must be a subclass of pytorch_lightning Callback. " - f"Provided callback: {cb} is not!" - ) - callbacks.append(cb) - - # Add needed callbacks - callbacks.append(LRSchedulerCallback()) - - # Checkpoint path - ckpt_path = _CKPT.value - - # Resume training - resume = _RESUME.value - if resume: - if ckpt_path is None: - resume_ckpt_path = osp.join( - config.output_dir, "checkpoints/last.ckpt" - ) - else: - resume_ckpt_path = ckpt_path - else: - resume_ckpt_path = None - - trainer = PLTrainer(callbacks=callbacks, **trainer_args) - - hyper_params = trainer_args - - if config.get("params", None) is not None: - hyper_params.update(config.params.to_dict()) - - training_module = TrainingModule( - config.model, - config.optimizers, - loss, - train_data_connector, - test_data_connector, - hyper_params, - config.seed, - ckpt_path if not resume else None, - ) - data_module = DataModule(config.data) - - if mode == "fit": - trainer.fit( - training_module, datamodule=data_module, ckpt_path=resume_ckpt_path - ) - elif mode == "test": - trainer.test(training_module, datamodule=data_module, verbose=False) - - -def entrypoint() -> None: - """Entry point for the CLI.""" - app.run(main) - - -if __name__ == "__main__": - entrypoint() diff --git a/vis4d/pl/trainer.py b/vis4d/pl/trainer.py deleted file mode 100644 index 53fc8115c..000000000 --- a/vis4d/pl/trainer.py +++ /dev/null @@ -1,133 +0,0 @@ -"""Trainer for PyTorch Lightning.""" - -from __future__ import annotations - -import os.path as osp - -from lightning.pytorch import Callback, Trainer -from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint -from lightning.pytorch.loggers import Logger, TensorBoardLogger -from lightning.pytorch.loggers.wandb import WandbLogger -from lightning.pytorch.strategies.ddp import DDPStrategy - -from vis4d.common import ArgsType -from vis4d.common.imports import TENSORBOARD_AVAILABLE -from vis4d.common.logging import rank_zero_info - - -class PLTrainer(Trainer): - """Trainer for PyTorch Lightning.""" - - def __init__( - self, - *args: ArgsType, - work_dir: str, - exp_name: str, - version: str, - epoch_based: bool = True, - find_unused_parameters: bool = False, - save_top_k: int = 1, - checkpoint_period: int = 1, - checkpoint_callback: ModelCheckpoint | None = None, - wandb: bool = False, - seed: int = -1, - **kwargs: ArgsType, - ) -> None: - """Perform some basic common setups at the beginning of a job. - - Args: - work_dir: Specific directory to save checkpoints, logs, etc. - Integrates with exp_name and version to get output_dir. - exp_name: Name of current experiment. - version: Version of current experiment. - epoch_based: Use epoch-based / iteration-based training. Default is - True. - find_unused_parameters: Activates PyTorch checking for unused - parameters in DDP setting. Default: False, for better - performance. - save_top_k: Save top k checkpoints. Default: 1 (save last). - checkpoint_period: After N epochs / stpes, save out checkpoints. - Default: 1. - checkpoint_callback: Custom PL checkpoint callback. Default: None. - wandb: Use weights and biases logging instead of tensorboard. - Default: False. - seed (int, optional): The integer value seed for global random - state. Defaults to -1. If -1, a random seed will be generated. - This will be set by TrainingModule. - """ - self.work_dir = work_dir - self.exp_name = exp_name - self.version = version - self.seed = seed - - self.output_dir = osp.join(work_dir, exp_name, version) - - # setup experiment logging - if "logger" not in kwargs or ( - isinstance(kwargs["logger"], bool) and kwargs["logger"] - ): - exp_logger: Logger | None = None - if wandb: # pragma: no cover - exp_logger = WandbLogger( - save_dir=work_dir, - project=exp_name, - name=version, - ) - elif TENSORBOARD_AVAILABLE: - exp_logger = TensorBoardLogger( - save_dir=work_dir, - name=exp_name, - version=version, - default_hp_metric=False, - ) - else: - rank_zero_info( - "Neither `tensorboard` nor `tensorboardX` is " - "available. Running without experiment logger. To log " - "your experiments, try `pip install`ing either." - ) - kwargs["logger"] = exp_logger - - callbacks: list[Callback] = [] - - # add learning rate / GPU stats monitor (logs to tensorboard) - if TENSORBOARD_AVAILABLE or wandb: - callbacks += [LearningRateMonitor(logging_interval="step")] - - # Model checkpointer - if checkpoint_callback is None: - if epoch_based: - checkpoint_cb = ModelCheckpoint( - dirpath=osp.join(self.output_dir, "checkpoints"), - verbose=True, - save_last=True, - save_top_k=save_top_k, - every_n_epochs=checkpoint_period, - save_on_train_epoch_end=True, - ) - else: - checkpoint_cb = ModelCheckpoint( - dirpath=osp.join(self.output_dir, "checkpoints"), - verbose=True, - save_last=True, - save_top_k=save_top_k, - every_n_train_steps=checkpoint_period, - ) - else: - checkpoint_cb = checkpoint_callback - callbacks += [checkpoint_cb] - - kwargs["callbacks"] += callbacks - - # add distributed strategy - if kwargs["devices"] == 0: - kwargs["accelerator"] = "cpu" - kwargs["devices"] = "auto" - elif kwargs["devices"] > 1: # pragma: no cover - if kwargs["accelerator"] == "gpu": - ddp_plugin = DDPStrategy( - find_unused_parameters=find_unused_parameters - ) - kwargs["strategy"] = ddp_plugin - - super().__init__(*args, **kwargs) diff --git a/vis4d/vis/functional/__init__.py b/vis4d/vis/functional/__init__.py deleted file mode 100644 index 5a46ff0ee..000000000 --- a/vis4d/vis/functional/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Function interface for visualization functions.""" - -from .image import ( - draw_bboxes, - draw_masks, - imshow, - imshow_bboxes, - imshow_masks, - imshow_topk_bboxes, - imshow_track_matches, -) -from .pointcloud import draw_points, show_3d, show_points - -__all__ = [ - "imshow", - "draw_masks", - "draw_bboxes", - "imshow_bboxes", - "imshow_masks", - "imshow_topk_bboxes", - "imshow_track_matches", - "show_3d", - "draw_points", - "show_points", -] diff --git a/vis4d/vis/image/bbox3d_visualizer.py b/vis4d/vis/image/bbox3d_visualizer.py index 74f4a7f35..1f336f48b 100644 --- a/vis4d/vis/image/bbox3d_visualizer.py +++ b/vis4d/vis/image/bbox3d_visualizer.py @@ -64,9 +64,11 @@ def __init__( image_mode: str = "RGB", width: int = 2, camera_near_clip: float = 0.15, + plot_heading: bool = True, axis_mode: AxisMode = AxisMode.ROS, trajectory_length: int = 10, plot_trajectory: bool = True, + save_boxes3d: bool = False, canvas: CanvasBackend | None = None, viewer: ImageViewerBackend | None = None, **kwargs: ArgsType, @@ -84,12 +86,16 @@ def __init__( width (int): Width of the drawn bounding boxes. Defaults to 2. camera_near_clip (float): Near clipping plane of the camera. Defaults to 0.15. + plot_heading (bool): If the heading should be plotted. Defaults to + True. axis_mode (AxisMode): Axis mode for the input bboxes. Defaults to AxisMode.ROS (i.e. global coordinate). trajectory_length (int): How many past frames should be used to draw the trajectory. Defaults to 10. plot_trajectory (bool): If the trajectory should be plotted. Defaults to True. + save_boxes3d (bool): If the corners of 3D boxes should be saved to + disk in the format of npy. Defaults to False. canvas (CanvasBackend): Backend that is used to draw on images. If None a PillowCanvasBackend is used. viewer (ImageViewerBackend): Backend that is used show images. If @@ -117,6 +123,9 @@ def __init__( self.width = width self.camera_near_clip = camera_near_clip + self.plot_heading = plot_heading + self.save_boxes3d = save_boxes3d + self.canvas = canvas if canvas is not None else PillowCanvasBackend() self.viewer = viewer if viewer is not None else MatplotlibImageViewer() @@ -124,7 +133,11 @@ def reset(self) -> None: """Reset visualizer.""" self._samples.clear() - def process( # type: ignore # pylint: disable=arguments-differ + def __repr__(self) -> str: + """Return string representation.""" + return "BoundingBox3DVisualizer" + + def process( # pylint: disable=arguments-differ self, cur_iter: int, images: list[ArrayLike], @@ -136,6 +149,7 @@ def process( # type: ignore # pylint: disable=arguments-differ class_ids: None | list[ArrayLikeInt] = None, track_ids: None | list[ArrayLikeInt] = None, sequence_names: None | list[str] = None, + categories: None | list[list[str]] = None, ) -> None: """Processes a batch of data. @@ -157,6 +171,9 @@ class ids each of shape [B, N]. Defaults to None. track ids each of shape [B, N]. Defaults to None. sequence_names (None | list[str], optional): List of sequence names of shape [B,]. Defaults to None. + categories (None | list[list[str]], optional): List of categories + for each image. Instead of class ids, the categories will be + used to label the boxes. Defaults to None. """ if self._run_on_batch(cur_iter): for batch, image in enumerate(images): @@ -172,6 +189,7 @@ class ids each of shape [B, N]. Defaults to None. None if class_ids is None else class_ids[batch], None if track_ids is None else track_ids[batch], None if sequence_names is None else sequence_names[batch], + None if categories is None else categories[batch], ) for tid in self.trajectories: @@ -189,6 +207,7 @@ def process_single_image( class_ids: None | ArrayLikeInt = None, track_ids: None | ArrayLikeInt = None, sequence_name: None | str = None, + categories: None | list[str] = None, camera_name: None | str = None, ) -> None: """Processes a single image entry. @@ -209,6 +228,9 @@ def process_single_image( shape [N]. Defaults to None. sequence_name (None | str, optional): Sequence name. Defaults to None. + categories (None | list[str], optional): List of categories for + each box. Instead of class ids, the categories will be used to + label the boxes. Defaults to None. camera_name (None | str, optional): Camera name. Defaults to None. """ img_normalized = preprocess_image(image, mode=self.image_mode) @@ -223,37 +245,40 @@ def process_single_image( data_sample = DataSample( img_normalized, image_name, - intrinsics_np, # type: ignore - extrinsics_np, # type: ignore + intrinsics_np, + extrinsics_np, sequence_name, camera_name, [], ) - for center, corners, label, color, track_id in zip( - *preprocess_boxes3d( - image_hw, - boxes3d, - intrinsics, - extrinsics, - scores, - class_ids, - track_ids, - self.color_palette, - self.class_id_mapping, - axis_mode=self.axis_mode, - ) - ): - data_sample.boxes.append( - DetectionBox3D( - corners=corners, - label=label, - color=color, - track_id=track_id, + if len(boxes3d) != 0: # type: ignore + for center, corners, label, color, track_id in zip( + *preprocess_boxes3d( + image_hw, + boxes3d, + intrinsics, + extrinsics, + scores, + class_ids, + track_ids, + self.color_palette, + self.class_id_mapping, + axis_mode=self.axis_mode, + categories=categories, ) - ) - if track_id is not None: - self.trajectories[track_id].append(center) + ): + data_sample.boxes.append( + DetectionBox3D( + corners=corners, + label=label, + color=color, + track_id=track_id, + ) + ) + if track_id is not None: + self.trajectories[track_id].append(center) + self._samples.append(data_sample) def show(self, cur_iter: int, blocking: bool = True) -> None: @@ -279,9 +304,13 @@ def _draw_image(self, sample: DataSample) -> NDArrayUI8: """ self.canvas.create_canvas(sample.image) - global_to_cam = inverse_rigid_transform( - torch.from_numpy(sample.extrinsics) - ).numpy() + if self.plot_trajectory: + assert ( + sample.extrinsics is not None + ), "Extrinsics is needed to plot trajectory." + global_to_cam = inverse_rigid_transform( + torch.from_numpy(sample.extrinsics) + ).numpy() for box in sample.boxes: self.canvas.draw_box_3d( @@ -290,18 +319,18 @@ def _draw_image(self, sample: DataSample) -> NDArrayUI8: sample.intrinsics, self.width, self.camera_near_clip, + self.plot_heading, ) selected_corner = project_point(box.corners[0], sample.intrinsics) self.canvas.draw_text( - (selected_corner[0], selected_corner[1]), - box.label, + (selected_corner[0], selected_corner[1]), box.label, box.color ) if self.plot_trajectory: assert ( - sample.extrinsics is not None and box.track_id is not None - ), "Extrinsics and track id must be set to plot trajectory." + box.track_id is not None + ), "track id must be set to plot trajectory." trajectory = self.trajectories[box.track_id] for center in trajectory: @@ -344,6 +373,14 @@ def save_to_disk(self, cur_iter: int, output_folder: str) -> None: os.makedirs(output_dir, exist_ok=True) self.canvas.save_to_disk(os.path.join(output_dir, image_name)) + if self.save_boxes3d: + corners = np.array([box.corners for box in sample.boxes]) + + np.save( + os.path.join(output_dir, f"{sample.image_name}.npy"), + corners, + ) + class MultiCameraBBox3DVisualizer(BoundingBox3DVisualizer): """Bounding box 3D visualizer class for multi-camera datasets.""" @@ -360,6 +397,10 @@ def __init__( self.cameras = cameras + def __repr__(self) -> str: + """Return string representation.""" + return "MultiCameraBBox3DVisualizer" + def process( # type: ignore # pylint: disable=arguments-differ self, cur_iter: int, @@ -372,6 +413,7 @@ def process( # type: ignore # pylint: disable=arguments-differ class_ids: list[ArrayLikeInt] | None = None, track_ids: list[ArrayLikeInt] | None = None, sequence_names: list[str] | None = None, + categories: None | list[list[str]] = None, ) -> None: """Processes a batch of data. @@ -393,6 +435,9 @@ class ids each of shape [B, N]. Defaults to None. track ids each of shape [B, N]. Defaults to None. sequence_names (None | list[str], optional): List of sequence names of shape [B,]. Defaults to None. + categories (None | list[list[str]], optional): List of categories + for each image. Instead of class ids, the categories will be + used to label the boxes. Defaults to None. """ if self._run_on_batch(cur_iter): for idx, batch_images in enumerate(images): @@ -415,5 +460,6 @@ class ids each of shape [B, N]. Defaults to None. if sequence_names is None else sequence_names[batch] ), + None if categories is None else categories[batch], self.cameras[idx], ) diff --git a/vis4d/vis/image/bev_visualizer.py b/vis4d/vis/image/bev_visualizer.py index 9ba2d4c7b..d6c68c901 100644 --- a/vis4d/vis/image/bev_visualizer.py +++ b/vis4d/vis/image/bev_visualizer.py @@ -115,11 +115,15 @@ def __init__( self.canvas = canvas if canvas is not None else PillowCanvasBackend() self.viewer = viewer if viewer is not None else MatplotlibImageViewer() + def __repr__(self) -> str: + """Return string representation.""" + return "BEVBBox3DVisualizer" + def reset(self) -> None: """Reset visualizer.""" self._samples.clear() - def process( # type: ignore # pylint: disable=arguments-differ + def process( # pylint: disable=arguments-differ self, cur_iter: int, sample_names: list[list[str]] | list[str], @@ -168,13 +172,13 @@ def process_single( extrinsics_np = array_to_numpy(extrinsics, n_dims=2, dtype=np.float32) data_sample = DataSample( sample_name, - extrinsics_np, # type: ignore + extrinsics_np, sequence_name, [], ) boxes3d_lidar, boxes3d = self._get_lidar_and_global_boxes3d( - boxes3d, extrinsics_np # type: ignore + boxes3d, extrinsics_np ) corners = boxes3d_to_corners( @@ -285,6 +289,7 @@ def _draw_image(self, sample: DataSample) -> NDArrayUI8: self.canvas.draw_text( (img_center[0] + distance - 25, img_center[1]), f"{10 * i} m", + color=(0, 0, 0), ) # Draw ego car diff --git a/vis4d/vis/image/bounding_box_visualizer.py b/vis4d/vis/image/bounding_box_visualizer.py index 72ece8d8a..5a436c89d 100644 --- a/vis4d/vis/image/bounding_box_visualizer.py +++ b/vis4d/vis/image/bounding_box_visualizer.py @@ -45,8 +45,9 @@ def __init__( self, *args: ArgsType, n_colors: int = 50, - class_id_mapping: dict[int, str] | None = None, + cat_mapping: dict[str, int] | None = None, file_type: str = "png", + width: int = 2, canvas: CanvasBackend = PillowCanvasBackend(), viewer: ImageViewerBackend = MatplotlibImageViewer(), **kwargs: ArgsType, @@ -55,28 +56,36 @@ def __init__( Args: n_colors (int): How many colors should be used for the internal - color map - class_id_mapping (dict[int, str]): Mapping from class id to - human readable name - file_type (str): Desired file type - canvas (CanvasBackend): Backend that is used to draw on images - viewer (ImageViewerBackend): Backend that is used show images + color map + cat_mapping (dict[str, int]): Mapping from class names to class + ids. Defaults to None. + file_type (str): Desired file type. Defaults to "png". + width (int): Width of the bounding box lines. Defaults to 2. + canvas (CanvasBackend): Backend that is used to draw on images. + viewer (ImageViewerBackend): Backend that is used show images. """ super().__init__(*args, **kwargs) self._samples: list[DataSample] = [] self.color_palette = generate_color_map(n_colors) self.class_id_mapping = ( - class_id_mapping if class_id_mapping is not None else {} + {v: k for k, v in cat_mapping.items()} + if cat_mapping is not None + else {} ) self.file_type = file_type + self.width = width self.canvas = canvas self.viewer = viewer + def __repr__(self) -> str: + """Return string representation of the visualizer.""" + return "BoundingBoxVisualizer" + def reset(self) -> None: """Reset visualizer.""" self._samples.clear() - def process( # type: ignore # pylint: disable=arguments-differ + def process( # pylint: disable=arguments-differ self, cur_iter: int, images: list[ArrayLike], @@ -85,6 +94,7 @@ def process( # type: ignore # pylint: disable=arguments-differ scores: None | list[ArrayLikeFloat] = None, class_ids: None | list[ArrayLikeInt] = None, track_ids: None | list[ArrayLikeInt] = None, + categories: None | list[list[str]] = None, ) -> None: """Processes a batch of data. @@ -100,6 +110,9 @@ def process( # type: ignore # pylint: disable=arguments-differ class ids each of shape [N]. Defaults to None. track_ids (None | list[ArrayLikeInt], optional): List of predicted track ids each of shape [N]. Defaults to None. + categories (None | list[list[str]], optional): List of categories + for each image. Instead of class ids, the categories will be + used to label the boxes. Defaults to None. """ if self._run_on_batch(cur_iter): for idx, image in enumerate(images): @@ -110,6 +123,7 @@ class ids each of shape [N]. Defaults to None. None if scores is None else scores[idx], None if class_ids is None else class_ids[idx], None if track_ids is None else track_ids[idx], + None if categories is None else categories[idx], ) def process_single_image( @@ -120,6 +134,7 @@ def process_single_image( scores: None | ArrayLikeFloat = None, class_ids: None | ArrayLikeInt = None, track_ids: None | ArrayLikeInt = None, + categories: None | list[str] = None, ) -> None: """Processes a single image entry. @@ -134,6 +149,9 @@ def process_single_image( shape [N]. Defaults to None. track_ids (None | ArrayLikeInt, optional): Predicted track ids of shape [N]. Defaults to None. + categories (None | list[str], optional): List of categories for + each box. Instead of class ids, the categories will be used to + label the boxes. Defaults to None. """ img_normalized = preprocess_image(image, mode=self.image_mode) data_sample = DataSample(img_normalized, image_name, []) @@ -146,6 +164,7 @@ def process_single_image( track_ids, self.color_palette, self.class_id_mapping, + categories=categories, ) ): data_sample.boxes.append( @@ -181,8 +200,8 @@ def _draw_image(self, sample: DataSample) -> NDArrayUI8: """ self.canvas.create_canvas(sample.image) for box in sample.boxes: - self.canvas.draw_box(box.corners, box.color) - self.canvas.draw_text(box.corners[:2], box.label) + self.canvas.draw_box(box.corners, box.color, width=self.width) + self.canvas.draw_text(box.corners[:2], box.label, box.color) return self.canvas.as_numpy_image() @@ -200,11 +219,7 @@ def save_to_disk(self, cur_iter: int, output_folder: str) -> None: for sample in self._samples: image_name = f"{sample.image_name}.{self.file_type}" - self.canvas.create_canvas(sample.image) - - for box in sample.boxes: - self.canvas.draw_box(box.corners, box.color) - self.canvas.draw_text(box.corners[:2], box.label) + _ = self._draw_image(sample) self.canvas.save_to_disk( os.path.join(output_folder, image_name) diff --git a/vis4d/vis/image/canvas/base.py b/vis4d/vis/image/canvas/base.py index 9025ff0e1..6eb31eda3 100644 --- a/vis4d/vis/image/canvas/base.py +++ b/vis4d/vis/image/canvas/base.py @@ -145,6 +145,7 @@ def draw_box_3d( intrinsics: NDArrayF32, width: int = 0, camera_near_clip: float = 0.15, + plot_heading: bool = True, ) -> None: """Draws a line between two points. @@ -156,6 +157,8 @@ def draw_box_3d( width (int, optional): The width of the line. Defaults to 0. camera_near_clip (float, optional): The near clipping plane of the camera. Defaults to 0.15. + plot_heading (bool, optional): If True, the heading of the box will + be plotted as a line. Defaults to True. """ raise NotImplementedError diff --git a/vis4d/vis/image/canvas/pillow_backend.py b/vis4d/vis/image/canvas/pillow_backend.py index 0dbb7c704..6bfd01040 100644 --- a/vis4d/vis/image/canvas/pillow_backend.py +++ b/vis4d/vis/image/canvas/pillow_backend.py @@ -2,12 +2,9 @@ from __future__ import annotations -import base64 -from io import BytesIO - import numpy as np from PIL import Image, ImageDraw -from PIL.ImageFont import ImageFont +from PIL.ImageFont import ImageFont, load_default from vis4d.common.typing import NDArrayBool, NDArrayF32, NDArrayF64, NDArrayUI8 @@ -18,14 +15,17 @@ class PillowCanvasBackend(CanvasBackend): """Canvas backend using Pillow.""" - def __init__(self, font: ImageFont | None = None) -> None: + def __init__( + self, font: ImageFont | None = None, font_size: int | None = None + ) -> None: """Creates a new canvas backend. Args: font (ImageFont): Pillow font to use for the label. + font_size (int): Font size to use for the label. """ self._image_draw: ImageDraw.ImageDraw | None = None - self._font = font if font is not None else load_default_font() + self._font = font if font is not None else load_default(font_size) self._image: Image.Image | None = None def create_canvas( @@ -45,11 +45,14 @@ def create_canvas( Raises: ValueError: If the canvas is not initialized. """ - if image is None and image_hw is None: - raise ValueError("Image or Image Shapes required to create canvas") if image_hw is not None: white_image = np.ones([*image_hw, 3]) * 255 image = white_image.astype(np.uint8) + else: + assert ( + image is not None + ), "Image or Image Shapes required to create canvas" + self._image = Image.fromarray(image) self._image_draw = ImageDraw.Draw(self._image) @@ -112,7 +115,13 @@ def draw_text( raise ValueError( "No Image Draw initialized! Did you call 'create_canvas'?" ) - self._image_draw.text(position, text, color, font=self._font) + left, top, right, bottom = self._image_draw.textbbox( + position, text, font=self._font + ) + self._image_draw.rectangle( + (left - 2, top - 2, right + 2, bottom + 2), fill=color + ) + self._image_draw.text(position, text, (255, 255, 255), font=self._font) def draw_box( self, @@ -279,6 +288,7 @@ def draw_box_3d( intrinsics: NDArrayF32, width: int = 0, camera_near_clip: float = 0.15, + plot_heading: bool = True, ) -> None: """Draws a 3D box onto the given canvas.""" # Draw Front @@ -324,16 +334,19 @@ def draw_box_3d( ) # Draw line indicating the front - center_bottom_forward = np.mean(corners[:2], axis=0, dtype=np.float32) - center_bottom = np.mean(corners[:4], axis=0, dtype=np.float32) - self._draw_box_3d_line( - tuple(center_bottom.tolist()), - tuple(center_bottom_forward.tolist()), - color, - intrinsics, - width, - camera_near_clip, - ) + if plot_heading: + center_bottom_forward = np.mean( + corners[:2], axis=0, dtype=np.float32 + ) + center_bottom = np.mean(corners[:4], axis=0, dtype=np.float32) + self._draw_box_3d_line( + tuple(center_bottom.tolist()), + tuple(center_bottom_forward.tolist()), + color, + intrinsics, + width, + camera_near_clip, + ) def as_numpy_image(self) -> NDArrayUI8: """Returns the current canvas as numpy image. @@ -361,140 +374,3 @@ def save_to_disk(self, image_path: str) -> None: "No Image initialized! Did you call 'create_canvas'?" ) self._image.save(image_path) - - -def load_default_font() -> ImageFont: - """Load a "better than nothing" default font.""" - f = ImageFont() - f._load_pilfont_data( # pylint: disable=protected-access - # courB08 - BytesIO( - base64.b64decode( - b""" -UElMZm9udAo7Ozs7OzsxMDsKREFUQQoAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAYAAAAA//8AAQAAAAAAAAABAAEA -BgAAAAH/+gADAAAAAQAAAAMABgAGAAAAAf/6AAT//QADAAAABgADAAYAAAAA//kABQABAAYAAAAL -AAgABgAAAAD/+AAFAAEACwAAABAACQAGAAAAAP/5AAUAAAAQAAAAFQAHAAYAAP////oABQAAABUA -AAAbAAYABgAAAAH/+QAE//wAGwAAAB4AAwAGAAAAAf/5AAQAAQAeAAAAIQAIAAYAAAAB//kABAAB -ACEAAAAkAAgABgAAAAD/+QAE//0AJAAAACgABAAGAAAAAP/6AAX//wAoAAAALQAFAAYAAAAB//8A -BAACAC0AAAAwAAMABgAAAAD//AAF//0AMAAAADUAAQAGAAAAAf//AAMAAAA1AAAANwABAAYAAAAB -//kABQABADcAAAA7AAgABgAAAAD/+QAFAAAAOwAAAEAABwAGAAAAAP/5AAYAAABAAAAARgAHAAYA -AAAA//kABQAAAEYAAABLAAcABgAAAAD/+QAFAAAASwAAAFAABwAGAAAAAP/5AAYAAABQAAAAVgAH -AAYAAAAA//kABQAAAFYAAABbAAcABgAAAAD/+QAFAAAAWwAAAGAABwAGAAAAAP/5AAUAAABgAAAA -ZQAHAAYAAAAA//kABQAAAGUAAABqAAcABgAAAAD/+QAFAAAAagAAAG8ABwAGAAAAAf/8AAMAAABv -AAAAcQAEAAYAAAAA//wAAwACAHEAAAB0AAYABgAAAAD/+gAE//8AdAAAAHgABQAGAAAAAP/7AAT/ -/gB4AAAAfAADAAYAAAAB//oABf//AHwAAACAAAUABgAAAAD/+gAFAAAAgAAAAIUABgAGAAAAAP/5 -AAYAAQCFAAAAiwAIAAYAAP////oABgAAAIsAAACSAAYABgAA////+gAFAAAAkgAAAJgABgAGAAAA -AP/6AAUAAACYAAAAnQAGAAYAAP////oABQAAAJ0AAACjAAYABgAA////+gAFAAAAowAAAKkABgAG -AAD////6AAUAAACpAAAArwAGAAYAAAAA//oABQAAAK8AAAC0AAYABgAA////+gAGAAAAtAAAALsA -BgAGAAAAAP/6AAQAAAC7AAAAvwAGAAYAAP////oABQAAAL8AAADFAAYABgAA////+gAGAAAAxQAA -AMwABgAGAAD////6AAUAAADMAAAA0gAGAAYAAP////oABQAAANIAAADYAAYABgAA////+gAGAAAA -2AAAAN8ABgAGAAAAAP/6AAUAAADfAAAA5AAGAAYAAP////oABQAAAOQAAADqAAYABgAAAAD/+gAF -AAEA6gAAAO8ABwAGAAD////6AAYAAADvAAAA9gAGAAYAAAAA//oABQAAAPYAAAD7AAYABgAA//// -+gAFAAAA+wAAAQEABgAGAAD////6AAYAAAEBAAABCAAGAAYAAP////oABgAAAQgAAAEPAAYABgAA -////+gAGAAABDwAAARYABgAGAAAAAP/6AAYAAAEWAAABHAAGAAYAAP////oABgAAARwAAAEjAAYA -BgAAAAD/+gAFAAABIwAAASgABgAGAAAAAf/5AAQAAQEoAAABKwAIAAYAAAAA//kABAABASsAAAEv -AAgABgAAAAH/+QAEAAEBLwAAATIACAAGAAAAAP/5AAX//AEyAAABNwADAAYAAAAAAAEABgACATcA -AAE9AAEABgAAAAH/+QAE//wBPQAAAUAAAwAGAAAAAP/7AAYAAAFAAAABRgAFAAYAAP////kABQAA -AUYAAAFMAAcABgAAAAD/+wAFAAABTAAAAVEABQAGAAAAAP/5AAYAAAFRAAABVwAHAAYAAAAA//sA -BQAAAVcAAAFcAAUABgAAAAD/+QAFAAABXAAAAWEABwAGAAAAAP/7AAYAAgFhAAABZwAHAAYAAP// -//kABQAAAWcAAAFtAAcABgAAAAD/+QAGAAABbQAAAXMABwAGAAAAAP/5AAQAAgFzAAABdwAJAAYA -AP////kABgAAAXcAAAF+AAcABgAAAAD/+QAGAAABfgAAAYQABwAGAAD////7AAUAAAGEAAABigAF -AAYAAP////sABQAAAYoAAAGQAAUABgAAAAD/+wAFAAABkAAAAZUABQAGAAD////7AAUAAgGVAAAB -mwAHAAYAAAAA//sABgACAZsAAAGhAAcABgAAAAD/+wAGAAABoQAAAacABQAGAAAAAP/7AAYAAAGn -AAABrQAFAAYAAAAA//kABgAAAa0AAAGzAAcABgAA////+wAGAAABswAAAboABQAGAAD////7AAUA -AAG6AAABwAAFAAYAAP////sABgAAAcAAAAHHAAUABgAAAAD/+wAGAAABxwAAAc0ABQAGAAD////7 -AAYAAgHNAAAB1AAHAAYAAAAA//sABQAAAdQAAAHZAAUABgAAAAH/+QAFAAEB2QAAAd0ACAAGAAAA -Av/6AAMAAQHdAAAB3gAHAAYAAAAA//kABAABAd4AAAHiAAgABgAAAAD/+wAF//0B4gAAAecAAgAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA -AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAYAAAAB -//sAAwACAecAAAHpAAcABgAAAAD/+QAFAAEB6QAAAe4ACAAGAAAAAP/5AAYAAAHuAAAB9AAHAAYA -AAAA//oABf//AfQAAAH5AAUABgAAAAD/+QAGAAAB+QAAAf8ABwAGAAAAAv/5AAMAAgH/AAACAAAJ -AAYAAAAA//kABQABAgAAAAIFAAgABgAAAAH/+gAE//sCBQAAAggAAQAGAAAAAP/5AAYAAAIIAAAC -DgAHAAYAAAAB//kABf/+Ag4AAAISAAUABgAA////+wAGAAACEgAAAhkABQAGAAAAAP/7AAX//gIZ -AAACHgADAAYAAAAA//wABf/9Ah4AAAIjAAEABgAAAAD/+QAHAAACIwAAAioABwAGAAAAAP/6AAT/ -+wIqAAACLgABAAYAAAAA//kABP/8Ai4AAAIyAAMABgAAAAD/+gAFAAACMgAAAjcABgAGAAAAAf/5 -AAT//QI3AAACOgAEAAYAAAAB//kABP/9AjoAAAI9AAQABgAAAAL/+QAE//sCPQAAAj8AAgAGAAD/ -///7AAYAAgI/AAACRgAHAAYAAAAA//kABgABAkYAAAJMAAgABgAAAAH//AAD//0CTAAAAk4AAQAG -AAAAAf//AAQAAgJOAAACUQADAAYAAAAB//kABP/9AlEAAAJUAAQABgAAAAH/+QAF//4CVAAAAlgA -BQAGAAD////7AAYAAAJYAAACXwAFAAYAAP////kABgAAAl8AAAJmAAcABgAA////+QAGAAACZgAA -Am0ABwAGAAD////5AAYAAAJtAAACdAAHAAYAAAAA//sABQACAnQAAAJ5AAcABgAA////9wAGAAAC -eQAAAoAACQAGAAD////3AAYAAAKAAAAChwAJAAYAAP////cABgAAAocAAAKOAAkABgAA////9wAG -AAACjgAAApUACQAGAAD////4AAYAAAKVAAACnAAIAAYAAP////cABgAAApwAAAKjAAkABgAA//// -+gAGAAACowAAAqoABgAGAAAAAP/6AAUAAgKqAAACrwAIAAYAAP////cABQAAAq8AAAK1AAkABgAA -////9wAFAAACtQAAArsACQAGAAD////3AAUAAAK7AAACwQAJAAYAAP////gABQAAAsEAAALHAAgA -BgAAAAD/9wAEAAACxwAAAssACQAGAAAAAP/3AAQAAALLAAACzwAJAAYAAAAA//cABAAAAs8AAALT -AAkABgAAAAD/+AAEAAAC0wAAAtcACAAGAAD////6AAUAAALXAAAC3QAGAAYAAP////cABgAAAt0A -AALkAAkABgAAAAD/9wAFAAAC5AAAAukACQAGAAAAAP/3AAUAAALpAAAC7gAJAAYAAAAA//cABQAA -Au4AAALzAAkABgAAAAD/9wAFAAAC8wAAAvgACQAGAAAAAP/4AAUAAAL4AAAC/QAIAAYAAAAA//oA -Bf//Av0AAAMCAAUABgAA////+gAGAAADAgAAAwkABgAGAAD////3AAYAAAMJAAADEAAJAAYAAP// -//cABgAAAxAAAAMXAAkABgAA////9wAGAAADFwAAAx4ACQAGAAD////4AAYAAAAAAAoABwASAAYA -AP////cABgAAAAcACgAOABMABgAA////+gAFAAAADgAKABQAEAAGAAD////6AAYAAAAUAAoAGwAQ -AAYAAAAA//gABgAAABsACgAhABIABgAAAAD/+AAGAAAAIQAKACcAEgAGAAAAAP/4AAYAAAAnAAoA -LQASAAYAAAAA//gABgAAAC0ACgAzABIABgAAAAD/+QAGAAAAMwAKADkAEQAGAAAAAP/3AAYAAAA5 -AAoAPwATAAYAAP////sABQAAAD8ACgBFAA8ABgAAAAD/+wAFAAIARQAKAEoAEQAGAAAAAP/4AAUA -AABKAAoATwASAAYAAAAA//gABQAAAE8ACgBUABIABgAAAAD/+AAFAAAAVAAKAFkAEgAGAAAAAP/5 -AAUAAABZAAoAXgARAAYAAAAA//gABgAAAF4ACgBkABIABgAAAAD/+AAGAAAAZAAKAGoAEgAGAAAA -AP/4AAYAAABqAAoAcAASAAYAAAAA//kABgAAAHAACgB2ABEABgAAAAD/+AAFAAAAdgAKAHsAEgAG -AAD////4AAYAAAB7AAoAggASAAYAAAAA//gABQAAAIIACgCHABIABgAAAAD/+AAFAAAAhwAKAIwA -EgAGAAAAAP/4AAUAAACMAAoAkQASAAYAAAAA//gABQAAAJEACgCWABIABgAAAAD/+QAFAAAAlgAK -AJsAEQAGAAAAAP/6AAX//wCbAAoAoAAPAAYAAAAA//oABQABAKAACgClABEABgAA////+AAGAAAA -pQAKAKwAEgAGAAD////4AAYAAACsAAoAswASAAYAAP////gABgAAALMACgC6ABIABgAA////+QAG -AAAAugAKAMEAEQAGAAD////4AAYAAgDBAAoAyAAUAAYAAP////kABQACAMgACgDOABMABgAA//// -+QAGAAIAzgAKANUAEw== -""" - ) - ), - Image.open( - BytesIO( - base64.b64decode( - b""" -iVBORw0KGgoAAAANSUhEUgAAAx4AAAAUAQAAAAArMtZoAAAEwElEQVR4nABlAJr/AHVE4czCI/4u -Mc4b7vuds/xzjz5/3/7u/n9vMe7vnfH/9++vPn/xyf5zhxzjt8GHw8+2d83u8x27199/nxuQ6Od9 -M43/5z2I+9n9ZtmDBwMQECDRQw/eQIQohJXxpBCNVE6QCCAAAAD//wBlAJr/AgALyj1t/wINwq0g -LeNZUworuN1cjTPIzrTX6ofHWeo3v336qPzfEwRmBnHTtf95/fglZK5N0PDgfRTslpGBvz7LFc4F -IUXBWQGjQ5MGCx34EDFPwXiY4YbYxavpnhHFrk14CDAAAAD//wBlAJr/AgKqRooH2gAgPeggvUAA -Bu2WfgPoAwzRAABAAAAAAACQgLz/3Uv4Gv+gX7BJgDeeGP6AAAD1NMDzKHD7ANWr3loYbxsAD791 -NAADfcoIDyP44K/jv4Y63/Z+t98Ovt+ub4T48LAAAAD//wBlAJr/AuplMlADJAAAAGuAphWpqhMx -in0A/fRvAYBABPgBwBUgABBQ/sYAyv9g0bCHgOLoGAAAAAAAREAAwI7nr0ArYpow7aX8//9LaP/9 -SjdavWA8ePHeBIKB//81/83ndznOaXx379wAAAD//wBlAJr/AqDxW+D3AABAAbUh/QMnbQag/gAY -AYDAAACgtgD/gOqAAAB5IA/8AAAk+n9w0AAA8AAAmFRJuPo27ciC0cD5oeW4E7KA/wD3ECMAn2tt -y8PgwH8AfAxFzC0JzeAMtratAsC/ffwAAAD//wBlAJr/BGKAyCAA4AAAAvgeYTAwHd1kmQF5chkG -ABoMIHcL5xVpTfQbUqzlAAAErwAQBgAAEOClA5D9il08AEh/tUzdCBsXkbgACED+woQg8Si9VeqY -lODCn7lmF6NhnAEYgAAA/NMIAAAAAAD//2JgjLZgVGBg5Pv/Tvpc8hwGBjYGJADjHDrAwPzAjv/H -/Wf3PzCwtzcwHmBgYGcwbZz8wHaCAQMDOwMDQ8MCBgYOC3W7mp+f0w+wHOYxO3OG+e376hsMZjk3 -AAAAAP//YmCMY2A4wMAIN5e5gQETPD6AZisDAwMDgzSDAAPjByiHcQMDAwMDg1nOze1lByRu5/47 -c4859311AYNZzg0AAAAA//9iYGDBYihOIIMuwIjGL39/fwffA8b//xv/P2BPtzzHwCBjUQAAAAD/ -/yLFBrIBAAAA//9i1HhcwdhizX7u8NZNzyLbvT97bfrMf/QHI8evOwcSqGUJAAAA//9iYBB81iSw -pEE170Qrg5MIYydHqwdDQRMrAwcVrQAAAAD//2J4x7j9AAMDn8Q/BgYLBoaiAwwMjPdvMDBYM1Tv -oJodAAAAAP//Yqo/83+dxePWlxl3npsel9lvLfPcqlE9725C+acfVLMEAAAA//9i+s9gwCoaaGMR -evta/58PTEWzr21hufPjA8N+qlnBwAAAAAD//2JiWLci5v1+HmFXDqcnULE/MxgYGBj+f6CaJQAA -AAD//2Ji2FrkY3iYpYC5qDeGgeEMAwPDvwQBBoYvcTwOVLMEAAAA//9isDBgkP///0EOg9z35v// -Gc/eeW7BwPj5+QGZhANUswMAAAD//2JgqGBgYGBgqEMXlvhMPUsAAAAA//8iYDd1AAAAAP//AwDR -w7IkEbzhVQAAAABJRU5ErkJggg== -""" - ) - ) - ), - ) - return f diff --git a/vis4d/vis/functional/image.py b/vis4d/vis/image/functional.py similarity index 88% rename from vis4d/vis/functional/image.py rename to vis4d/vis/image/functional.py index baf516438..b9355c6f2 100644 --- a/vis4d/vis/functional/image.py +++ b/vis4d/vis/image/functional.py @@ -4,7 +4,7 @@ import numpy as np -from vis4d.common.array import array_to_numpy, arrays_to_numpy +from vis4d.common.array import array_to_numpy from vis4d.common.typing import ( ArrayLike, ArrayLikeBool, @@ -13,22 +13,24 @@ NDArrayF32, NDArrayUI8, ) -from vis4d.vis.image.canvas import CanvasBackend, PillowCanvasBackend -from vis4d.vis.image.util import ( + +from ..util import generate_color_map +from .canvas import CanvasBackend, PillowCanvasBackend +from .util import ( preprocess_boxes, preprocess_boxes3d, preprocess_image, preprocess_masks, project_point, ) -from vis4d.vis.image.viewer import ImageViewerBackend, MatplotlibImageViewer -from vis4d.vis.util import generate_color_map +from .viewer import ImageViewerBackend, MatplotlibImageViewer def imshow( image: ArrayLike, image_mode: str = "RGB", image_viewer: ImageViewerBackend = MatplotlibImageViewer(), + file_path: str | None = None, ) -> None: """Shows a single image. @@ -37,28 +39,13 @@ def imshow( image_mode (str, optional): Image Mode. Defaults to "RGB". image_viewer (ImageViewerBackend, optional): The Image viewer backend to use. Defaults to MatplotlibImageViewer(). + file_path (str): The path to save the image to. Defaults to None. """ image = preprocess_image(image, image_mode) image_viewer.show_images([image]) - -def imsave( - image: ArrayLike, - file_path: str, - image_mode: str = "RGB", - image_viewer: ImageViewerBackend = MatplotlibImageViewer(), -) -> None: - """Shows a single image. - - Args: - image (NDArrayNumber): The image to show. - file_path (str): The path to save the image to. - image_mode (str, optional): Image Mode. Defaults to "RGB". - image_viewer (ImageViewerBackend, optional): The Image viewer backend - to use. Defaults to MatplotlibImageViewer(). - """ - image = preprocess_image(image, image_mode) - image_viewer.save_images([image], [file_path]) + if file_path is not None: + image_viewer.save_images([image], [file_path]) def draw_masks( @@ -158,6 +145,7 @@ def imshow_bboxes( image_mode: str = "RGB", box_width: int = 1, image_viewer: ImageViewerBackend = MatplotlibImageViewer(), + file_path: str | None = None, ) -> None: """Shows the bounding boxes overlayed on the given image. @@ -176,6 +164,7 @@ class id to class name box_width (int, optional): Width of the box border. Defaults to 1. image_viewer (ImageViewerBackend, optional): The Image viewer backend to use. Defaults to MatplotlibImageViewer(). + file_path (str): The path to save the image to. Defaults to None. """ image = preprocess_image(image, mode=image_mode) img = draw_bboxes( @@ -189,7 +178,7 @@ class id to class name image_mode, box_width, ) - imshow(img, image_mode, image_viewer) + imshow(img, image_mode, image_viewer, file_path) def draw_bbox3d( @@ -244,6 +233,7 @@ def imshow_bboxes3d( n_colors: int = 50, image_mode: str = "RGB", image_viewer: ImageViewerBackend = MatplotlibImageViewer(), + file_path: str | None = None, ) -> None: """Show image with bounding boxes.""" image = preprocess_image(image, mode=image_mode) @@ -259,7 +249,7 @@ def imshow_bboxes3d( n_colors=n_colors, image_mode=image_mode, ) - imshow(img, image_mode, image_viewer) + imshow(img, image_mode, image_viewer, file_path) def imshow_masks( @@ -270,6 +260,7 @@ def imshow_masks( image_mode: str = "RGB", canvas: CanvasBackend = PillowCanvasBackend(), image_viewer: ImageViewerBackend = MatplotlibImageViewer(), + file_path: str | None = None, ) -> None: """Shows semantic masks overlayed over the given image. @@ -286,11 +277,13 @@ def imshow_masks( Defaults to PillowCanvasBackend(). image_viewer (ImageViewerBackend, optional): The Image viewer backend to use. Defaults to MatplotlibImageViewer(). + file_path (str): The path to save the image to. Defaults to None. """ imshow( draw_masks(image, masks, class_ids, n_colors, image_mode, canvas), image_mode, image_viewer, + file_path, ) @@ -306,6 +299,7 @@ def imshow_topk_bboxes( image_mode: str = "RGB", box_width: int = 1, image_viewer: ImageViewerBackend = MatplotlibImageViewer(), + file_path: str | None = None, ) -> None: """Visualize the 'topk' bounding boxes with highest score. @@ -325,6 +319,7 @@ class id to class name box_width (int, optional): Width of the box border. Defaults to 1. image_viewer (ImageViewerBackend, optional): The Image viewer backend to use. Defaults to MatplotlibImageViewer(). + file_path (str): The path to save the image to. Defaults to None. """ scores = array_to_numpy(scores, n_dims=1, dtype=np.float32) @@ -344,6 +339,7 @@ class id to class name image_mode, box_width, image_viewer, + file_path, ) @@ -356,6 +352,7 @@ def imshow_track_matches( ref_track_ids: list[ArrayLikeInt], image_mode: str = "RGB", image_viewer: ImageViewerBackend = MatplotlibImageViewer(), + file_path: str | None = None, ) -> None: """Visualize paired bounding boxes successively for batched frame pairs. @@ -372,16 +369,25 @@ def imshow_track_matches( image_mode (str, optional): Color mode if the image. Defaults to "RGB". image_viewer (ImageViewerBackend, optional): The Image viewer backend to use. Defaults to MatplotlibImageViewer(). + file_path (str): The path to save the image to. Defaults to None. """ - key_imgs_np = arrays_to_numpy(*key_imgs, n_dims=3, dtype=np.float32) - ref_imgs_np = arrays_to_numpy(*ref_imgs, n_dims=3, dtype=np.float32) - key_boxes_np = arrays_to_numpy(*key_boxes, n_dims=2, dtype=np.float32) - ref_boxes_np = arrays_to_numpy(*ref_boxes, n_dims=2, dtype=np.float32) - key_track_ids_np = arrays_to_numpy( - *key_track_ids, n_dims=1, dtype=np.int32 + key_imgs_np = tuple( + array_to_numpy(img, n_dims=3, dtype=np.float32) for img in key_imgs + ) + ref_imgs_np = tuple( + array_to_numpy(img, n_dims=3, dtype=np.float32) for img in ref_imgs + ) + key_boxes_np = tuple( + array_to_numpy(b, n_dims=2, dtype=np.float32) for b in key_boxes + ) + ref_boxes_np = tuple( + array_to_numpy(b, n_dims=2, dtype=np.float32) for b in ref_boxes + ) + key_track_ids_np = tuple( + array_to_numpy(t, n_dims=1, dtype=np.int32) for t in key_track_ids ) - ref_track_ids_np = arrays_to_numpy( - *ref_track_ids, n_dims=1, dtype=np.int32 + ref_track_ids_np = tuple( + array_to_numpy(t, n_dims=1, dtype=np.int32) for t in ref_track_ids ) for batch_i, (key_box, ref_box) in enumerate( @@ -404,12 +410,14 @@ def imshow_track_matches( key_box[key_i], image_mode=image_mode, image_viewer=image_viewer, + file_path=file_path, ) imshow_bboxes( ref_image, ref_box[ref_i], image_mode=image_mode, image_viewer=image_viewer, + file_path=file_path, ) else: # stack imgs horizontal @@ -420,4 +428,4 @@ def imshow_track_matches( ref_image, ref_box[batch_i], image_mode=image_mode ) stacked_img = np.vstack([k_img, r_img]) - imshow(stacked_img, image_mode, image_viewer) + imshow(stacked_img, image_mode, image_viewer, file_path) diff --git a/vis4d/vis/image/seg_mask_visualizer.py b/vis4d/vis/image/seg_mask_visualizer.py index faaddb2d4..7b95451a0 100644 --- a/vis4d/vis/image/seg_mask_visualizer.py +++ b/vis4d/vis/image/seg_mask_visualizer.py @@ -120,7 +120,7 @@ def _draw_image(self, sample: ImageWithSegMask) -> NDArrayUI8: self.canvas.draw_bitmap(mask.mask, mask.color) return self.canvas.as_numpy_image() - def process( # type: ignore # pylint: disable=arguments-differ + def process( # pylint: disable=arguments-differ self, cur_iter: int, images: list[ArrayLikeFloat], diff --git a/vis4d/vis/image/util.py b/vis4d/vis/image/util.py index c2a20afbd..d32f223ad 100644 --- a/vis4d/vis/image/util.py +++ b/vis4d/vis/image/util.py @@ -27,29 +27,25 @@ def _get_box_label( - class_id: int | None, + category: str | None, score: float | None, track_id: int | None, - class_id_mapping: dict[int, str] | None = None, ) -> str: """Gets a unique string representation for a box definition. Args: - class_id (int): The class id for this box + category (str): The category name score (float): The confidence score track_id (int): The track id - class_id_mapping (dict[int,str]): Mapping of class_id to class name Returns: str: Label for this box of format 'class_name, track_id, score%' """ labels = [] - if class_id_mapping is None: - class_id_mapping = {} - if class_id is not None: - labels.append(class_id_mapping.get(class_id, str(class_id))) + if category is not None: + labels.append(category) if track_id is not None: labels.append(str(track_id)) if score is not None: @@ -88,6 +84,7 @@ def preprocess_boxes( color_palette: list[tuple[int, int, int]] = DEFAULT_COLOR_MAPPING, class_id_mapping: dict[int, str] | None = None, default_color: tuple[int, int, int] = (255, 0, 0), + categories: None | list[str] = None, ) -> tuple[ list[tuple[float, float, float, float]], list[str], @@ -111,6 +108,8 @@ def preprocess_boxes( to color tuple (0-255). default_color (tuple[int, int, int]): fallback color for boxes of no class or track id is given. + categories (None | list[str], optional): List of categories for each + box. Returns: boxes_proc (list[tuple[float, float, float, float]]): List of box @@ -157,9 +156,15 @@ class or track id is given. ) ) colors_proc.append(color) - labels_proc.append( - _get_box_label(class_id, score, track_id, class_id_mapping) - ) + + if categories is not None: + category = categories[idx] + elif class_id is not None: + category = class_id_mapping.get(class_id, str(class_id)) + else: + category = None + + labels_proc.append(_get_box_label(category, score, track_id)) return boxes_proc, labels_proc, colors_proc @@ -175,6 +180,7 @@ def preprocess_boxes3d( class_id_mapping: dict[int, str] | None = None, default_color: tuple[int, int, int] = (255, 0, 0), axis_mode: AxisMode = AxisMode.OPENCV, + categories: None | list[str] = None, ) -> tuple[ list[tuple[float, float, float]], list[list[tuple[float, float, float]]], @@ -223,18 +229,28 @@ def preprocess_boxes3d( class_ids_np = array_to_numpy(class_ids, n_dims=1, dtype=np.int32) track_ids_np = array_to_numpy(track_ids, n_dims=1, dtype=np.int32) - boxes3d_np = boxes3d_np[mask] - corners_np = corners_np[mask] - scores_np = scores_np[mask] if scores_np is not None else None - class_ids_np = class_ids_np[mask] if class_ids_np is not None else None - track_ids_np = track_ids_np[mask] if track_ids_np is not None else None - centers_proc: list[tuple[float, float, float]] = [] corners_proc: list[list[tuple[float, float, float]]] = [] colors_proc: list[tuple[int, int, int]] = [] labels_proc: list[str] = [] track_ids_proc: list[int] = [] + if len(mask) == 1: + if not mask[0]: + return ( + centers_proc, + corners_proc, + labels_proc, + colors_proc, + track_ids_proc, + ) + else: + boxes3d_np = boxes3d_np[mask] + corners_np = corners_np[mask] + scores_np = scores_np[mask] if scores_np is not None else None + class_ids_np = class_ids_np[mask] if class_ids_np is not None else None + track_ids_np = track_ids_np[mask] if track_ids_np is not None else None + for idx in range(corners_np.shape[0]): class_id = None if class_ids_np is None else class_ids_np[idx].item() score = None if scores_np is None else scores_np[idx].item() @@ -256,10 +272,17 @@ def preprocess_boxes3d( ) corners_proc.append([tuple(pts) for pts in corners_np[idx].tolist()]) colors_proc.append(color) - labels_proc.append( - _get_box_label(class_id, score, track_id, class_id_mapping) - ) - track_ids_proc.append(track_id) + + if categories is not None: + category = categories[idx] + elif class_id is not None: + category = class_id_mapping.get(class_id, str(class_id)) + else: + category = None + + labels_proc.append(_get_box_label(category, score, track_id)) + if track_id is not None: + track_ids_proc.append(track_id) return centers_proc, corners_proc, labels_proc, colors_proc, track_ids_proc @@ -289,9 +312,7 @@ def preprocess_masks( Raises: ValueError: If the masks have an invalid shape. """ - masks_np: NDArrayUI8 = array_to_numpy( # type: ignore - masks, n_dims=None, dtype=np.uint8 - ) + masks_np = array_to_numpy(masks, n_dims=None, dtype=np.uint8) if len(masks_np.shape) == 2: masks_np, class_ids = _to_binary_mask(masks_np) @@ -338,7 +359,7 @@ def preprocess_image(image: ArrayLike, mode: str = "RGB") -> NDArrayUI8: # Convert torch to numpy convention if not image_np.shape[-1] == 3: - image_np = np.transpose(image_np, (1, 2, 0)) # type: ignore + image_np = np.transpose(image_np, (1, 2, 0)) # Convert image_np to [0, 255] min_val, max_val = ( diff --git a/vis4d/vis/functional/pointcloud.py b/vis4d/vis/pointcloud/functional.py similarity index 93% rename from vis4d/vis/functional/pointcloud.py rename to vis4d/vis/pointcloud/functional.py index ed4335a6d..fec7fd83a 100644 --- a/vis4d/vis/functional/pointcloud.py +++ b/vis4d/vis/pointcloud/functional.py @@ -3,12 +3,10 @@ from __future__ import annotations from vis4d.common.typing import ArrayLikeFloat, ArrayLikeInt -from vis4d.vis.pointcloud.scene import Scene3D -from vis4d.vis.pointcloud.viewer import ( - Open3DVisualizationBackend, - PointCloudVisualizerBackend, -) -from vis4d.vis.util import DEFAULT_COLOR_MAPPING + +from ..util import DEFAULT_COLOR_MAPPING +from .scene import Scene3D +from .viewer import Open3DVisualizationBackend, PointCloudVisualizerBackend def show_3d( diff --git a/vis4d/vis/pointcloud/pointcloud_visualizer.py b/vis4d/vis/pointcloud/pointcloud_visualizer.py index e0dbc315c..ae2721134 100644 --- a/vis4d/vis/pointcloud/pointcloud_visualizer.py +++ b/vis4d/vis/pointcloud/pointcloud_visualizer.py @@ -110,7 +110,7 @@ def process_single( points_xyz, colors=colors, classes=semantics, instances=instances ) - def process( # type: ignore # pylint: disable=arguments-differ + def process( # pylint: disable=arguments-differ self, cur_iter: int, points_xyz: NDArrayF64, diff --git a/vis4d/vis/pointcloud/scene.py b/vis4d/vis/pointcloud/scene.py index 37f941d3e..963a99f7c 100644 --- a/vis4d/vis/pointcloud/scene.py +++ b/vis4d/vis/pointcloud/scene.py @@ -112,7 +112,7 @@ def __init__( else: self.num_instances = len( np.unique( - self.classes * np.max(self.instances) + self.instances # type: ignore # pylint: disable=line-too-long + self.classes * np.max(self.instances) + self.instances ) ) diff --git a/vis4d/zoo/base/dataloader.py b/vis4d/zoo/base/dataloader.py index 6dcb70bd8..91a32695a 100644 --- a/vis4d/zoo/base/dataloader.py +++ b/vis4d/zoo/base/dataloader.py @@ -32,6 +32,7 @@ def get_train_dataloader_cfg( sensors: Sequence[str] | None = None, pin_memory: bool | FieldReference = True, shuffle: bool | FieldReference = True, + aspect_ratio_grouping: bool | FieldReference = False, ) -> ConfigDict: """Creates dataloader configuration given dataset and preprocessing. @@ -59,6 +60,8 @@ def get_train_dataloader_cfg( Defaults to True. shuffle (bool | FieldReference, optional): Whether to shuffle the dataset. Defaults to True. + aspect_ratio_grouping (bool | FieldReference, optional): Whether to + group the samples by aspect ratio. Defaults to False. Returns: ConfigDict: Configuration that can be instantiate as a dataloader. @@ -84,6 +87,7 @@ def get_train_dataloader_cfg( sensors=sensors, pin_memory=pin_memory, shuffle=shuffle, + aspect_ratio_grouping=aspect_ratio_grouping, ) diff --git a/vis4d/zoo/base/datasets/shift/common.py b/vis4d/zoo/base/datasets/shift/common.py index 6d38d3031..add1e8397 100644 --- a/vis4d/zoo/base/datasets/shift/common.py +++ b/vis4d/zoo/base/datasets/shift/common.py @@ -346,7 +346,7 @@ def get_shift_dataloader_config( return data -def get_shift_config( # pylint: disable=too-many-arguments +def get_shift_config( # pylint: disable=too-many-arguments, too-many-positional-arguments, line-too-long data_root: str = "data/shift/images", train_split: str = "train", train_framerate: str = "images", diff --git a/vis4d/zoo/base/models/yolox.py b/vis4d/zoo/base/models/yolox.py index 044650deb..eadc3bc19 100644 --- a/vis4d/zoo/base/models/yolox.py +++ b/vis4d/zoo/base/models/yolox.py @@ -3,8 +3,8 @@ from __future__ import annotations from ml_collections import ConfigDict, FieldReference -from torch.optim import SGD from torch.optim.lr_scheduler import CosineAnnealingLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import OptimizerConfig diff --git a/vis4d/zoo/base/runtime.py b/vis4d/zoo/base/runtime.py index 65e842f97..d127ace78 100644 --- a/vis4d/zoo/base/runtime.py +++ b/vis4d/zoo/base/runtime.py @@ -5,11 +5,11 @@ import platform from datetime import datetime -from ml_collections import ConfigDict, FieldReference +from ml_collections import ConfigDict from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig -from vis4d.engine.callbacks import CheckpointCallback, LoggingCallback +from vis4d.engine.callbacks import LoggingCallback def get_default_cfg( @@ -58,13 +58,13 @@ def get_default_cfg( config.use_tf32 = False config.tf32_matmul_precision = "highest" config.benchmark = False + config.compute_flops = False + config.check_unused_parameters = False return config def get_default_callbacks_cfg( - output_dir: str | FieldReference, - checkpoint_period: int = 1, epoch_based: bool = True, refresh_rate: int = 50, ) -> list[ConfigDict]: @@ -72,11 +72,8 @@ def get_default_callbacks_cfg( It will return a list of callbacks config including: - LoggingCallback - - CheckpointCallback Args: - output_dir (str | FieldReference): Output directory. - checkpoint_period (int, optional): Checkpoint period. Defaults to 1. epoch_based (bool, optional): Whether to use epoch based logging. refresh_rate (int, optional): Refresh rate for the logging. Defaults to 50. @@ -93,14 +90,4 @@ def get_default_callbacks_cfg( ) ) - # Checkpoint - callbacks.append( - class_config( - CheckpointCallback, - epoch_based=epoch_based, - save_prefix=output_dir, - checkpoint_period=checkpoint_period, - ) - ) - return callbacks diff --git a/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_1x_bdd100k.py b/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_1x_bdd100k.py index d92196693..cfb13fcd6 100644 --- a/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_1x_bdd100k.py +++ b/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_1x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -120,7 +120,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_3x_bdd100k.py b/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_3x_bdd100k.py index 7abc3a7c1..b9e12b39a 100644 --- a/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_3x_bdd100k.py +++ b/vis4d/zoo/bdd100k/faster_rcnn/faster_rcnn_r50_3x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -121,7 +121,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_1x_bdd100k.py b/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_1x_bdd100k.py index 8fc5385d9..a9e50b762 100644 --- a/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_1x_bdd100k.py +++ b/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_1x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -126,7 +126,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_3x_bdd100k.py b/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_3x_bdd100k.py index 9a709d213..25e93e1cf 100644 --- a/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_3x_bdd100k.py +++ b/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_3x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -126,7 +126,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_5x_bdd100k.py b/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_5x_bdd100k.py index 9c6dec315..c4fcec939 100644 --- a/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_5x_bdd100k.py +++ b/vis4d/zoo/bdd100k/mask_rcnn/mask_rcnn_r50_5x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -126,7 +126,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/bdd100k/qdtrack/qdtrack_frcnn_r50_fpn_1x_bdd100k.py b/vis4d/zoo/bdd100k/qdtrack/qdtrack_frcnn_r50_fpn_1x_bdd100k.py index 9b047d2d3..0e3c878a7 100644 --- a/vis4d/zoo/bdd100k/qdtrack/qdtrack_frcnn_r50_fpn_1x_bdd100k.py +++ b/vis4d/zoo/bdd100k/qdtrack/qdtrack_frcnn_r50_fpn_1x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -110,7 +110,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( diff --git a/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r101_80k_bdd100k.py b/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r101_80k_bdd100k.py index db45092d3..c4c3701f1 100644 --- a/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r101_80k_bdd100k.py +++ b/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r101_80k_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -150,11 +150,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_40k_bdd100k.py b/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_40k_bdd100k.py index deb8f78fe..6301745b0 100644 --- a/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_40k_bdd100k.py +++ b/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_40k_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -140,11 +140,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_80k_bdd100k.py b/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_80k_bdd100k.py index 06b95cd37..a3682d244 100644 --- a/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_80k_bdd100k.py +++ b/vis4d/zoo/bdd100k/semantic_fpn/semantic_fpn_r50_80k_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -140,11 +140,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/bevformer/bevformer_base.py b/vis4d/zoo/bevformer/bevformer_base.py index c31202022..5c386e5bc 100644 --- a/vis4d/zoo/bevformer/bevformer_base.py +++ b/vis4d/zoo/bevformer/bevformer_base.py @@ -2,8 +2,8 @@ """BEVFormer base with ResNet-101-DCN backbone.""" from __future__ import annotations -import pytorch_lightning as pl -from torch.optim import AdamW +import lightning.pytorch as pl +from torch.optim.adamw import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR from vis4d.config import class_config @@ -122,7 +122,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( diff --git a/vis4d/zoo/bevformer/bevformer_tiny.py b/vis4d/zoo/bevformer/bevformer_tiny.py index 942455183..db82f2a51 100644 --- a/vis4d/zoo/bevformer/bevformer_tiny.py +++ b/vis4d/zoo/bevformer/bevformer_tiny.py @@ -2,8 +2,8 @@ """BEVFormer tiny with ResNet-50 backbone.""" from __future__ import annotations -import pytorch_lightning as pl -from torch.optim import AdamW +import lightning.pytorch as pl +from torch.optim.adamw import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR from vis4d.config import class_config @@ -160,7 +160,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( diff --git a/vis4d/zoo/bevformer/bevformer_vis.py b/vis4d/zoo/bevformer/bevformer_vis.py index fee098fea..97b411da3 100644 --- a/vis4d/zoo/bevformer/bevformer_vis.py +++ b/vis4d/zoo/bevformer/bevformer_vis.py @@ -35,7 +35,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_kf3d_24e_nusc.py b/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_kf3d_24e_nusc.py index 6f685615e..5d2c928bc 100644 --- a/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_kf3d_24e_nusc.py +++ b/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_kf3d_24e_nusc.py @@ -2,9 +2,9 @@ """CC-3DT with Faster-RCNN ResNet-101 detector using KF3D motion model.""" from __future__ import annotations -import pytorch_lightning as pl -from torch.optim import SGD +import lightning.pytorch as pl from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -155,7 +155,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( diff --git a/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_pure_det_nusc.py b/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_pure_det_nusc.py index ae44762cb..3a9afd1be 100644 --- a/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_pure_det_nusc.py +++ b/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r101_fpn_pure_det_nusc.py @@ -60,7 +60,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( diff --git a/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r50_fpn_kf3d_12e_nusc.py b/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r50_fpn_kf3d_12e_nusc.py index 1354b44a0..61bfc0f46 100644 --- a/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r50_fpn_kf3d_12e_nusc.py +++ b/vis4d/zoo/cc_3dt/cc_3dt_frcnn_r50_fpn_kf3d_12e_nusc.py @@ -2,9 +2,9 @@ """CC-3DT with Faster-RCNN ResNet-50 detector using KF3D motion model.""" from __future__ import annotations -import pytorch_lightning as pl -from torch.optim import SGD +import lightning.pytorch as pl from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -155,7 +155,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( @@ -168,7 +168,7 @@ def get_config() -> ExperimentConfig: split=test_split, ), save_predictions=True, - save_prefix=config.output_dir, + output_dir=config.output_dir, test_connector=class_config( CallbackConnector, key_mapping=CONN_NUSC_DET3D_EVAL ), @@ -180,7 +180,7 @@ def get_config() -> ExperimentConfig: EvaluatorCallback, evaluator=class_config(NuScenesTrack3DEvaluator), save_predictions=True, - save_prefix=config.output_dir, + output_dir=config.output_dir, test_connector=class_config( CallbackConnector, key_mapping=CONN_NUSC_TRACK3D_EVAL ), diff --git a/vis4d/zoo/cc_3dt/cc_3dt_nusc_test.py b/vis4d/zoo/cc_3dt/cc_3dt_nusc_test.py index 27a2bd047..192b09b26 100644 --- a/vis4d/zoo/cc_3dt/cc_3dt_nusc_test.py +++ b/vis4d/zoo/cc_3dt/cc_3dt_nusc_test.py @@ -68,7 +68,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Evaluator callbacks.append( diff --git a/vis4d/zoo/cc_3dt/cc_3dt_nusc_vis.py b/vis4d/zoo/cc_3dt/cc_3dt_nusc_vis.py index af086411e..b23011647 100644 --- a/vis4d/zoo/cc_3dt/cc_3dt_nusc_vis.py +++ b/vis4d/zoo/cc_3dt/cc_3dt_nusc_vis.py @@ -36,7 +36,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( @@ -50,7 +50,8 @@ def get_config() -> ExperimentConfig: cameras=NuScenes.CAMERAS, vis_freq=1, ), - save_prefix=config.output_dir, + output_dir=config.output_dir, + save_prefix="boxes3d", test_connector=class_config( MultiSensorCallbackConnector, key_mapping=CONN_NUSC_BBOX_3D_VIS, @@ -62,7 +63,8 @@ def get_config() -> ExperimentConfig: class_config( VisualizerCallback, visualizer=class_config(BEVBBox3DVisualizer, width=2, vis_freq=1), - save_prefix=config.output_dir, + output_dir=config.output_dir, + save_prefix="bev", test_connector=class_config( MultiSensorCallbackConnector, key_mapping=CONN_NUSC_BEV_BBOX_3D_VIS, diff --git a/vis4d/zoo/cc_3dt/velo_lstm_bevformer_base_100e_nusc.py b/vis4d/zoo/cc_3dt/velo_lstm_bevformer_base_100e_nusc.py index 55d6eac90..eb5fc2796 100644 --- a/vis4d/zoo/cc_3dt/velo_lstm_bevformer_base_100e_nusc.py +++ b/vis4d/zoo/cc_3dt/velo_lstm_bevformer_base_100e_nusc.py @@ -2,8 +2,8 @@ """CC-3DT VeloLSTM for BEVFormer on nuScenes.""" from __future__ import annotations -import pytorch_lightning as pl -from torch.optim import Adam +import lightning.pytorch as pl +from torch.optim.adam import Adam from torch.optim.lr_scheduler import MultiStepLR from vis4d.config import class_config @@ -134,7 +134,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() config.callbacks = callbacks diff --git a/vis4d/zoo/cc_3dt/velo_lstm_frcnn_r101_fpn_100e_nusc.py b/vis4d/zoo/cc_3dt/velo_lstm_frcnn_r101_fpn_100e_nusc.py index 01a8dbe08..4c3b3406a 100644 --- a/vis4d/zoo/cc_3dt/velo_lstm_frcnn_r101_fpn_100e_nusc.py +++ b/vis4d/zoo/cc_3dt/velo_lstm_frcnn_r101_fpn_100e_nusc.py @@ -2,8 +2,8 @@ """CC-3DT VeloLSTM on nuScenes.""" from __future__ import annotations -import pytorch_lightning as pl -from torch.optim import Adam +import lightning.pytorch as pl +from torch.optim.adam import Adam from torch.optim.lr_scheduler import MultiStepLR from vis4d.config import class_config @@ -133,7 +133,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() config.callbacks = callbacks diff --git a/vis4d/zoo/faster_rcnn/faster_rcnn_coco.py b/vis4d/zoo/faster_rcnn/faster_rcnn_coco.py index 005bedb81..8fccda47b 100644 --- a/vis4d/zoo/faster_rcnn/faster_rcnn_coco.py +++ b/vis4d/zoo/faster_rcnn/faster_rcnn_coco.py @@ -2,18 +2,11 @@ """Faster RCNN COCO training example.""" from __future__ import annotations -import lightning.pytorch as pl -import numpy as np -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config -from vis4d.config.sweep import grid_search -from vis4d.config.typing import ( - ExperimentConfig, - ExperimentParameters, - ParameterSweepConfig, -) +from vis4d.config.typing import ExperimentConfig, ExperimentParameters from vis4d.data.io.hdf5 import HDF5Backend from vis4d.engine.callbacks import EvaluatorCallback, VisualizerCallback from vis4d.engine.connectors import CallbackConnector, DataConnector @@ -136,14 +129,14 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( class_config( VisualizerCallback, visualizer=class_config(BoundingBoxVisualizer, vis_freq=100), - save_prefix=config.output_dir, + output_dir=config.output_dir, test_connector=class_config( CallbackConnector, key_mapping=CONN_BBOX_2D_VIS ), @@ -174,29 +167,4 @@ def get_config() -> ExperimentConfig: pl_trainer.max_epochs = params.num_epochs config.pl_trainer = pl_trainer - # PL Callbacks - pl_callbacks: list[pl.callbacks.Callback] = [] - config.pl_callbacks = pl_callbacks - return config.value_mode() - - -def get_sweep() -> ParameterSweepConfig: # pragma: no cover - """Returns the config dict for a grid search over learning rate. - - The name of the experiments will also be updated to include the learning - rate in the format "lr_{params.lr:.3f}_". - - Returns: - ParameterSweepConfig: The configuration that can be used to run a grid - search. It can be passed to replicate_config to create a list of - configs that can be used to run a grid search. - """ - # Here we define the parameters that we want to sweep over. - # In order to sweep over multiple parameters, we can pass a list of - # parameters to the grid_search function. - sweep_config = grid_search("params.lr", list(np.linspace(0.001, 0.01, 3))) - - # Here we update the name of the experiment to include the learning rate. - sweep_config.suffix = "lr_{params.lr:.3f}_" - return sweep_config diff --git a/vis4d/zoo/fcn_resnet/fcn_resnet_coco.py b/vis4d/zoo/fcn_resnet/fcn_resnet_coco.py index 1f177c3fc..8a77b45f5 100644 --- a/vis4d/zoo/fcn_resnet/fcn_resnet_coco.py +++ b/vis4d/zoo/fcn_resnet/fcn_resnet_coco.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -140,11 +140,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) config.callbacks = callbacks diff --git a/vis4d/zoo/mask_rcnn/mask_rcnn_coco.py b/vis4d/zoo/mask_rcnn/mask_rcnn_coco.py index 58816730a..260bd8c08 100644 --- a/vis4d/zoo/mask_rcnn/mask_rcnn_coco.py +++ b/vis4d/zoo/mask_rcnn/mask_rcnn_coco.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -148,7 +148,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/qdtrack/qdtrack_frcnn_r50_fpn_augs_1x_bdd100k.py b/vis4d/zoo/qdtrack/qdtrack_frcnn_r50_fpn_augs_1x_bdd100k.py index a49ca5462..497186f75 100644 --- a/vis4d/zoo/qdtrack/qdtrack_frcnn_r50_fpn_augs_1x_bdd100k.py +++ b/vis4d/zoo/qdtrack/qdtrack_frcnn_r50_fpn_augs_1x_bdd100k.py @@ -3,8 +3,8 @@ from __future__ import annotations from lightning.pytorch.callbacks import ModelCheckpoint -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -116,7 +116,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Mode switch for strong augmentations callbacks += [class_config(YOLOXModeSwitchCallback, switch_epoch=9)] diff --git a/vis4d/zoo/qdtrack/qdtrack_yolox_x_25e_bdd100k.py b/vis4d/zoo/qdtrack/qdtrack_yolox_x_25e_bdd100k.py index 0cfcbabd6..160a62bb6 100644 --- a/vis4d/zoo/qdtrack/qdtrack_yolox_x_25e_bdd100k.py +++ b/vis4d/zoo/qdtrack/qdtrack_yolox_x_25e_bdd100k.py @@ -2,7 +2,7 @@ """QDTrack with YOLOX-x on BDD100K.""" from __future__ import annotations -import pytorch_lightning as pl +import lightning.pytorch as pl from lightning.pytorch.callbacks import ModelCheckpoint from vis4d.config import class_config @@ -101,7 +101,7 @@ def get_config() -> ExperimentConfig: ###################################################### # Logger and Checkpoint callbacks = get_default_callbacks_cfg( - config.output_dir, refresh_rate=config.log_every_n_steps + refresh_rate=config.log_every_n_steps ) # YOLOX callbacks diff --git a/vis4d/zoo/retinanet/retinanet_coco.py b/vis4d/zoo/retinanet/retinanet_coco.py index 5ab9d42d5..e0b8df8b8 100644 --- a/vis4d/zoo/retinanet/retinanet_coco.py +++ b/vis4d/zoo/retinanet/retinanet_coco.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -162,7 +162,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_12e_shift.py b/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_12e_shift.py index 1e452102d..90df01636 100644 --- a/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_12e_shift.py +++ b/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_12e_shift.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -126,7 +126,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_36e_shift.py b/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_36e_shift.py index 946e1fb61..afe377388 100644 --- a/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_36e_shift.py +++ b/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_36e_shift.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -126,7 +126,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_6e_shift_all_domains.py b/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_6e_shift_all_domains.py index 7d84cae75..76bcc8208 100644 --- a/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_6e_shift_all_domains.py +++ b/vis4d/zoo/shift/faster_rcnn/faster_rcnn_r50_6e_shift_all_domains.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -126,7 +126,7 @@ def get_config() -> ExperimentConfig: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_12e_shift.py b/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_12e_shift.py index d9da7bb43..53ce37b7e 100644 --- a/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_12e_shift.py +++ b/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_12e_shift.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import FieldConfigDict, class_config from vis4d.data.io.hdf5 import HDF5Backend @@ -127,7 +127,7 @@ def get_config() -> FieldConfigDict: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_36e_shift.py b/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_36e_shift.py index 40963714e..2a3b3eb4e 100644 --- a/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_36e_shift.py +++ b/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_36e_shift.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import FieldConfigDict, class_config from vis4d.data.io.hdf5 import HDF5Backend @@ -127,7 +127,7 @@ def get_config() -> FieldConfigDict: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_6e_shift_all_domains.py b/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_6e_shift_all_domains.py index 1f915339d..d03e1c132 100644 --- a/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_6e_shift_all_domains.py +++ b/vis4d/zoo/shift/mask_rcnn/mask_rcnn_r50_6e_shift_all_domains.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch.optim import SGD from torch.optim.lr_scheduler import LinearLR, MultiStepLR +from torch.optim.sgd import SGD from vis4d.config import FieldConfigDict, class_config from vis4d.data.io.hdf5 import HDF5Backend @@ -127,7 +127,7 @@ def get_config() -> FieldConfigDict: ## CALLBACKS ## ###################################################### # Logger and Checkpoint - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # Visualizer callbacks.append( diff --git a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift.py b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift.py index 983d8fef3..b41cede54 100644 --- a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift.py +++ b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch import optim from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -106,7 +106,7 @@ def get_config() -> ExperimentConfig: config.optimizers = [ get_optimizer_cfg( optimizer=class_config( - optim.SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 + SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 ), lr_schedulers=[ get_lr_scheduler_cfg( @@ -143,11 +143,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift_all_domains.py b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift_all_domains.py index 6fa1b785f..94ed983fc 100644 --- a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift_all_domains.py +++ b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_160k_shift_all_domains.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch import optim from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -108,7 +108,7 @@ def get_config() -> ExperimentConfig: config.optimizers = [ get_optimizer_cfg( optimizer=class_config( - optim.SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 + SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 ), lr_schedulers=[ get_lr_scheduler_cfg( @@ -145,11 +145,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift.py b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift.py index bc86ea101..70bc89188 100644 --- a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift.py +++ b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch import optim from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -106,7 +106,7 @@ def get_config() -> ExperimentConfig: config.optimizers = [ get_optimizer_cfg( optimizer=class_config( - optim.SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 + SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 ), lr_schedulers=[ get_lr_scheduler_cfg( @@ -143,11 +143,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift_all_domains.py b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift_all_domains.py index 8a6f40959..789ab1757 100644 --- a/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift_all_domains.py +++ b/vis4d/zoo/shift/semantic_fpn/semantic_fpn_r50_40k_shift_all_domains.py @@ -3,8 +3,8 @@ from __future__ import annotations import lightning.pytorch as pl -from torch import optim from torch.optim.lr_scheduler import LinearLR +from torch.optim.sgd import SGD from vis4d.config import class_config from vis4d.config.typing import ExperimentConfig, ExperimentParameters @@ -106,7 +106,7 @@ def get_config() -> ExperimentConfig: config.optimizers = [ get_optimizer_cfg( optimizer=class_config( - optim.SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 + SGD, lr=params.lr, momentum=0.9, weight_decay=0.0005 ), lr_schedulers=[ get_lr_scheduler_cfg( @@ -143,11 +143,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg( - config.output_dir, - epoch_based=False, - checkpoint_period=config.val_check_interval, - ) + callbacks = get_default_callbacks_cfg(epoch_based=False) # Evaluator callbacks.append( diff --git a/vis4d/zoo/util.py b/vis4d/zoo/util.py new file mode 100644 index 000000000..b58c5f870 --- /dev/null +++ b/vis4d/zoo/util.py @@ -0,0 +1,14 @@ +"""Utility functions for the zoo module.""" + +from __future__ import annotations + +import importlib + +from vis4d.config.typing import ExperimentConfig + + +def get_config_for_name(config_name: str) -> ExperimentConfig: + """Get config for name.""" + module = importlib.import_module("vis4d.zoo." + config_name) + + return module.get_config() diff --git a/vis4d/zoo/vit/vit_small_imagenet.py b/vis4d/zoo/vit/vit_small_imagenet.py index 85ae78f02..88f219064 100644 --- a/vis4d/zoo/vit/vit_small_imagenet.py +++ b/vis4d/zoo/vit/vit_small_imagenet.py @@ -2,9 +2,9 @@ """VIT ImageNet-1k training example.""" from __future__ import annotations -import pytorch_lightning as pl +import lightning.pytorch as pl from torch import nn -from torch.optim import AdamW +from torch.optim.adamw import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR from vis4d.config import class_config @@ -150,7 +150,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## GENERIC CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # EMA callback callbacks.append(class_config(EMACallback)) diff --git a/vis4d/zoo/vit/vit_tiny_imagenet.py b/vis4d/zoo/vit/vit_tiny_imagenet.py index 0d7f8c588..a109da931 100644 --- a/vis4d/zoo/vit/vit_tiny_imagenet.py +++ b/vis4d/zoo/vit/vit_tiny_imagenet.py @@ -2,9 +2,9 @@ """VIT ImageNet-1k training example.""" from __future__ import annotations -import pytorch_lightning as pl +import lightning.pytorch as pl from torch import nn -from torch.optim import AdamW +from torch.optim.adamw import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR from vis4d.config import class_config @@ -150,7 +150,7 @@ def get_config() -> ExperimentConfig: ###################################################### ## GENERIC CALLBACKS ## ###################################################### - callbacks = get_default_callbacks_cfg(config.output_dir) + callbacks = get_default_callbacks_cfg() # EMA callback callbacks.append(class_config(EMACallback)) diff --git a/vis4d/zoo/yolox/yolox_s_300e_coco.py b/vis4d/zoo/yolox/yolox_s_300e_coco.py index 7403d3625..8b654980f 100644 --- a/vis4d/zoo/yolox/yolox_s_300e_coco.py +++ b/vis4d/zoo/yolox/yolox_s_300e_coco.py @@ -98,7 +98,7 @@ def get_config() -> ExperimentConfig: ###################################################### # Logger and Checkpoint callbacks = get_default_callbacks_cfg( - config.output_dir, refresh_rate=config.log_every_n_steps + refresh_rate=config.log_every_n_steps ) # YOLOX callbacks diff --git a/vis4d/zoo/yolox/yolox_tiny_300e_coco.py b/vis4d/zoo/yolox/yolox_tiny_300e_coco.py index e1c8d50f3..7d63b07cd 100644 --- a/vis4d/zoo/yolox/yolox_tiny_300e_coco.py +++ b/vis4d/zoo/yolox/yolox_tiny_300e_coco.py @@ -101,7 +101,7 @@ def get_config() -> ExperimentConfig: ###################################################### # Logger and Checkpoint callbacks = get_default_callbacks_cfg( - config.output_dir, refresh_rate=config.log_every_n_steps + refresh_rate=config.log_every_n_steps ) # YOLOX callbacks