diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 504d1291..94adef06 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -257,6 +257,14 @@ repos: pass_filenames: false files: ^visdet/engine/hub/ + - id: zuban-models + name: Zuban type checker (visdet/models) + entry: uv run zuban check visdet/models --exclude '/dense_heads/' --exclude '/backbones/hrnet.py' --exclude '/backbones/swin.py' --exclude '/backbones/regnet.py' --exclude '/backbones/res2net.py' --exclude '/backbones/resnest.py' --exclude '/backbones/cspnext.py' + language: system + types: [python] + pass_filenames: false + files: ^visdet/models/ + - repo: https://github.com/kynan/nbstripout rev: 0.8.1 hooks: diff --git a/docs/about/contributing.md b/docs/about/contributing.md index 9dac21c3..4dbf0dca 100644 --- a/docs/about/contributing.md +++ b/docs/about/contributing.md @@ -28,7 +28,7 @@ visdet-worktrees/ We use the following tools to maintain code quality: - `ruff` for linting and formatting -- `pyright` for type checking +- `zuban` for type checking - `prek` hooks for automated checks (faster Rust-based alternative to pre-commit) To set up the development environment: diff --git a/docs/development/contributing.md b/docs/development/contributing.md index f557f8cc..1e658587 100644 --- a/docs/development/contributing.md +++ b/docs/development/contributing.md @@ -15,7 +15,7 @@ We appreciate all contributions to improve VisDet. Please follow the guidelines We use the following tools to maintain code quality: - `ruff` for linting and formatting -- `pyright` for type checking +- `zuban` for type checking - `prek` hooks for automated checks (faster Rust-based alternative to pre-commit) To set up the development environment: diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index 4710a9e9..a99f8027 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -1,234 +1,143 @@ -# Prerequisites +# Installation -In this section we demonstrate how to prepare an environment with PyTorch. +VisDet is published as a normal Python package. In most cases you **do not** need to clone this repository to use it. -This framework works on Linux, Windows and macOS. It requires Python 3.7+, CUDA 9.2+ and PyTorch 1.5+. +Cloning the repo is only needed if you want to: -```{note} -If you are experienced with PyTorch and have already installed it, just skip this part and jump to the [next section](#installation). Otherwise, you can follow these steps for the preparation. -``` +- Develop VisDet itself (editable install, tests, docs) +- Use the repo’s training scripts under `tools/` +- Use the repo’s example assets/configs as-is + +## Requirements -**Step 0.** Install [uv](https://docs.astral.sh/uv/) - a fast Python package manager. +- Python `>=3.10,<3.13` +- PyTorch + torchvision + - For CUDA / GPU support, install PyTorch following https://pytorch.org first (so you get the correct CUDA build). + +## Install with uv (recommended) + +**Step 0.** Install [uv](https://docs.astral.sh/uv/) (if you don’t already have it). ```shell -# On macOS and Linux +# macOS / Linux curl -LsSf https://astral.sh/uv/install.sh | sh -# On Windows +# Windows powershell -c "irm https://astral.sh/uv/install.ps1 | iex" ``` -**Step 1.** uv will automatically manage Python versions and virtual environments for you. No separate setup needed! - -# Installation - -We recommend that users follow our best practices to install VisDet using uv. However, the whole process is highly customizable. See [Customize Installation](#customize-installation) section for more information. - -## Best Practices - -**Step 0.** Clone the repository and navigate to it. +**Step 1.** Create a virtual environment. ```shell -git clone -cd visdet +uv venv --python 3.12 ``` -**Step 1.** Install dependencies using uv. - -For development (includes all dependencies): +**Step 2.** Activate it. ```shell -uv sync +# macOS / Linux +source .venv/bin/activate + +# Windows (PowerShell) +.venv\Scripts\Activate.ps1 ``` -For specific extras (e.g., just documentation): +**Step 3.** Install VisDet. ```shell -uv sync --extra mkdocs +uv pip install visdet ``` -**Step 2.** That's it! uv has: -- Created a virtual environment -- Installed Python 3.12 (or the version specified in `.python-version`) -- Installed all dependencies from `pyproject.toml` -- Installed the package in editable mode - -## Verify the installation - -To verify whether This framework installed correctly, we provide some sample codes to run an inference demo. - -**Step 1.** We need to download config and checkpoint files. +### Optional extras ```shell -uv run mim download mmdet --config yolov3_mobilenetv2_320_300e_coco --dest . +# Extra (optional) dependencies used by some features +uv pip install "visdet[optional]" + +# Everything in the optional group +uv pip install "visdet[all]" ``` -The downloading will take several seconds or more, depending on your network environment. When it is done, you will find two files `yolov3_mobilenetv2_320_300e_coco.py` and `yolov3_mobilenetv2_320_300e_coco_20210719_215349-d18dff72.pth` in your current folder. +## Install with pip -**Step 2.** Verify the inference demo. +If you prefer standard tooling: ```shell -uv run python demo/image_demo.py demo/demo.jpg yolov3_mobilenetv2_320_300e_coco.py yolov3_mobilenetv2_320_300e_coco_20210719_215349-d18dff72.pth --device cpu --out-file result.jpg +python -m venv .venv +source .venv/bin/activate # or .venv\Scripts\activate on Windows +pip install -U pip +pip install visdet ``` -You will see a new image `result.jpg` on your current folder, where bounding boxes are plotted on cars, benches, etc. - -Alternatively, you can run Python code directly: +## Verify the installation -```python -# Run with: uv run python -from mmdet.apis import init_detector, inference_detector +A minimal smoke-check (no repo clone required): -config_file = 'yolov3_mobilenetv2_320_300e_coco.py' -checkpoint_file = 'yolov3_mobilenetv2_320_300e_coco_20210719_215349-d18dff72.pth' -model = init_detector(config_file, checkpoint_file, device='cpu') # or device='cuda:0' -inference_detector(model, 'demo/demo.jpg') +```shell +python -c "import visdet; print(visdet.__version__)" ``` -You will see a list of arrays printed, indicating the detected bounding boxes. - -## Customize Installation - -### CUDA versions - -When installing PyTorch, you need to specify the version of CUDA. If you are not clear on which to choose, follow our recommendations: +Optional: run a quick inference using a built-in YAML preset (this may download model weights on first use): -- For Ampere-based NVIDIA GPUs, such as GeForce 30 series and NVIDIA A100, CUDA 11 is a must. -- For older NVIDIA GPUs, CUDA 11 is backward compatible, but CUDA 10.2 offers better compatibility and is more lightweight. - -Please make sure the GPU driver satisfies the minimum version requirements. See [this table](https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#cuda-major-component-versions__table-cuda-toolkit-driver-versions) for more information. +```python +from visdet.apis import DetInferencer -```{note} -uv handles PyTorch installation automatically based on your `pyproject.toml` configuration. For specific CUDA versions, you may need to configure the PyTorch index URL in your project settings. +inferencer = DetInferencer(model="rtmdet-s", device="cpu") +results = inferencer("path/to/your_image.jpg") +print(results) ``` -### Installing additional packages +## Install from source (development) -To add additional packages to your project: +Only needed if you’re contributing to VisDet. ```shell -# Add a new dependency -uv add package-name - -# Add a development dependency -uv add --dev package-name - -# Add an optional dependency to a specific group -uv add --optional group-name package-name +git clone +cd visdet +uv sync ``` -### Install on CPU-only platforms - -This framework can be built for CPU only environment. In CPU mode you can train, test or inference a model. - -However some functionalities are gone in this mode: - -- Deformable Convolution -- Modulated Deformable Convolution -- ROI pooling -- Deformable ROI pooling -- CARAFE -- SyncBatchNorm -- CrissCrossAttention -- MaskedConv2d -- Temporal Interlace Shift -- nms_cuda -- sigmoid_focal_loss_cuda -- bbox_overlaps +That will: -If you try to train/test/inference a model containing above ops, an error will be raised. -The following table lists affected algorithms. +- Create a virtual environment +- Install dependencies from `pyproject.toml` +- Install VisDet in editable mode -| Operator | Model | -| :-----------------------------------------------------: | :--------------------------------------------------------------------------------------: | -| Deformable Convolution/Modulated Deformable Convolution | DCN、Guided Anchoring、RepPoints、CentripetalNet、VFNet、CascadeRPN、NAS-FCOS、DetectoRS | -| MaskedConv2d | Guided Anchoring | -| CARAFE | CARAFE | -| SyncBatchNorm | ResNeSt | +## Install on Google Colab -### Install on Google Colab - -[Google Colab](https://research.google.com/) usually has PyTorch installed. -Here's how to install visdet with uv on Colab: - -**Step 1.** Install uv in Colab. +Colab usually already has PyTorch installed. ```shell !curl -LsSf https://astral.sh/uv/install.sh | sh +!uv venv --python 3.12 +!uv pip install visdet ``` -**Step 2.** Clone and install visdet. - -```shell -!git clone -%cd visdet -!uv sync -``` - -**Step 3.** Verification. - ```python -!uv run python -c "import mmdet; print(mmdet.__version__)" +import visdet +print(visdet.__version__) ``` ```{note} -Within Jupyter, the exclamation mark `!` is used to call external executables and `%cd` is a [magic command](https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-cd) to change the current working directory of Python. +Within Jupyter, the exclamation mark `!` runs shell commands. ``` -### Using VisDet with Docker +## Using VisDet with Docker -We provide a Dockerfile in the `docker/` directory to build an image. Ensure that your [docker version](https://docs.docker.com/engine/install/) >=19.03. +The repo contains a `docker/` folder with a Dockerfile. This path **does** require cloning the repository. ```shell -# build an image with PyTorch 1.6, CUDA 10.1 -# If you prefer other versions, just modified the Dockerfile docker build -t visdet docker/ ``` -Run it with +Run it with: ```shell docker run --gpus all --shm-size=8g -it -v {DATA_DIR}:/visdet/data visdet ``` -### Running commands with uv - -All commands should be prefixed with `uv run` to ensure they use the project's virtual environment: - -```shell -# Running Python scripts -uv run python your_script.py - -# Running installed CLI tools -uv run pytest tests/ - -# Running prek hooks -uv run prek run --all-files - -# Building documentation -uv run mkdocs build -``` - -Alternatively, you can activate the virtual environment manually: - -```shell -source .venv/bin/activate # On Linux/macOS -# or -.venv\Scripts\activate # On Windows - -# Now you can run commands without 'uv run' prefix -python your_script.py -pytest tests/ -``` - -## Trouble shooting - -If you have some issues during the installation, please check the documentation carefully. - -Common issues and solutions: - -- **Import errors**: Make sure you have installed the package with `uv sync` -- **CUDA issues**: Verify your CUDA version matches your PyTorch installation -- **Version conflicts**: Try creating a fresh virtual environment +## Troubleshooting -If you encounter problems not covered in the documentation, please refer to the [Contributing Guide](../development/contributing.md) for how to report issues. +- **Import errors**: ensure you installed into the environment you’re running (`which python` / `python -V`). +- **CUDA issues**: install PyTorch for your CUDA version (then reinstall `visdet` if needed). +- **Version conflicts**: try a fresh env: `rm -rf .venv && uv venv && uv pip install visdet`. diff --git a/docs/getting-started/quick-start.md b/docs/getting-started/quick-start.md index 1da4b4c4..72f0f504 100644 --- a/docs/getting-started/quick-start.md +++ b/docs/getting-started/quick-start.md @@ -22,31 +22,40 @@ runner.train() ## Inference with Pre-trained Models -### Using High-level APIs +### Using YAML model presets (no repo clone) -You can use high-level APIs to perform inference on images: +VisDet ships with YAML model presets, so you can run inference without needing this repository’s Python config files. ```python -from visdet.apis import init_detector, inference_detector +from visdet.apis import DetInferencer -# Specify the config file and checkpoint file -config_file = 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py' -checkpoint_file = 'checkpoints/faster_rcnn_r50_fpn_1x_coco.pth' +# Uses a built-in preset name/alias; may download weights on first use +inferencer = DetInferencer(model="rtmdet-s", device="cuda:0") +result = inferencer("path/to/image.jpg") +print(result) +``` -# Build the model from a config file and a checkpoint file -model = init_detector(config_file, checkpoint_file, device='cuda:0') +### Using explicit config + checkpoint (repo-style) -# Test a single image -img = 'demo/demo.jpg' -result = inference_detector(model, img) +If you *are* working from a cloned repo (or you have your own configs/checkpoints), you can still use the classic APIs: + +```python +from visdet.apis import init_detector, inference_detector, show_result_pyplot -# Show the results -from visdet.apis import show_result_pyplot +config_file = "path/to/config.py" +checkpoint_file = "path/to/checkpoint.pth" + +model = init_detector(config_file, checkpoint_file, device="cuda:0") +img = "path/to/image.jpg" +result = inference_detector(model, img) show_result_pyplot(model, img, result) ``` ## Training a Model +The training / testing entrypoints under `tools/` are part of this repository. +If you want to use them, clone the repo and run from the repo root. + ### Train with a Single GPU ```bash diff --git a/pyproject.toml b/pyproject.toml index c8281cb7..fb091b57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -127,13 +127,6 @@ known-first-party = ["visdet"] split-on-trailing-comma = true force-single-line = false -[tool.pyright] -typeCheckingMode = "basic" -pythonVersion = "3.10" -reportMissingImports = true -reportMissingTypeStubs = false -include = ["visdet"] -exclude = ["visdet/tests", "tests", "archive", "configs"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/visdet/models/backbones/resnet.py b/visdet/models/backbones/resnet.py index 8de2bff3..8eaede99 100644 --- a/visdet/models/backbones/resnet.py +++ b/visdet/models/backbones/resnet.py @@ -1,9 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. import logging import warnings +from typing import Any, cast import torch.nn as nn import torch.utils.checkpoint as cp +from torch import Tensor from torch.nn.modules.batchnorm import _BatchNorm from visdet.cv.cnn import build_conv_layer, build_norm_layer @@ -36,19 +38,19 @@ class BasicBlock(BaseModule): def __init__( self, - inplanes, - planes, - stride=1, - dilation=1, - downsample=None, - style="pytorch", - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type="BN"), - dcn=None, - plugins=None, - init_cfg=None, - ): + inplanes: int, + planes: int, + stride: int = 1, + dilation: int = 1, + downsample: nn.Module | None = None, + style: str = "pytorch", + with_cp: bool = False, + conv_cfg: dict | None = None, + norm_cfg: dict = dict(type="BN"), + dcn: dict | None = None, + plugins: list[dict] | None = None, + init_cfg: dict | list[dict] | None = None, + ) -> None: super(BasicBlock, self).__init__(init_cfg=init_cfg) assert dcn is None, "DCN is not supported in BasicBlock" assert plugins is None, "Plugins are not supported yet" @@ -86,12 +88,11 @@ def norm2(self): """nn.Module: normalization layer after the second convolution layer""" return getattr(self, self.norm2_name) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: """Forward function.""" def _inner_forward(x): identity = x - out = self.conv1(x) out = self.norm1(out) out = self.relu(out) @@ -138,19 +139,19 @@ class Bottleneck(BaseModule): def __init__( self, - inplanes, - planes, - stride=1, - dilation=1, - downsample=None, - style="pytorch", - with_cp=False, - conv_cfg=None, - norm_cfg=dict(type="BN"), - dcn=None, - plugins=None, - init_cfg=None, - ): + inplanes: int, + planes: int, + stride: int = 1, + dilation: int = 1, + downsample: nn.Module | None = None, + style: str = "pytorch", + with_cp: bool = False, + conv_cfg: dict | None = None, + norm_cfg: dict = dict(type="BN"), + dcn: dict | None = None, + plugins: list[dict] | None = None, + init_cfg: dict | list[dict] | None = None, + ) -> None: """Bottleneck block for ResNet. If style is "pytorch", the stride-two layer is the 3x3 conv layer, if @@ -232,21 +233,28 @@ def norm3(self): """nn.Module: normalization layer after the third convolution layer""" return getattr(self, self.norm3_name) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: """Forward function.""" - def _inner_forward(x): + def _inner_forward(x: Tensor) -> Tensor: + conv1 = cast(nn.Module, self.conv1) + conv2 = cast(nn.Module, self.conv2) + conv3 = cast(nn.Module, self.conv3) + norm1 = cast(nn.Module, self.norm1) + norm2 = cast(nn.Module, self.norm2) + norm3 = cast(nn.Module, self.norm3) + identity = x - out = self.conv1(x) - out = self.norm1(out) + out = conv1(x) + out = norm1(out) out = self.relu(out) - out = self.conv2(out) - out = self.norm2(out) + out = conv2(out) + out = norm2(out) out = self.relu(out) - out = self.conv3(out) - out = self.norm3(out) + out = conv3(out) + out = norm3(out) if self.downsample is not None: identity = self.downsample(x) @@ -512,7 +520,7 @@ def _freeze_stages(self): for param in m.parameters(): param.requires_grad = False - def forward(self, x): + def forward(self, x: Tensor) -> tuple[Tensor, ...]: """Forward function.""" if self.deep_stem: x = self.stem(x) @@ -521,7 +529,7 @@ def forward(self, x): x = self.norm1(x) x = self.relu(x) x = self.maxpool(x) - outs = [] + outs: list[Tensor] = [] for i, layer_name in enumerate(self.res_layers): res_layer = getattr(self, layer_name) x = res_layer(x) diff --git a/visdet/models/backbones/resnext.py b/visdet/models/backbones/resnext.py index 87110912..47d8f087 100644 --- a/visdet/models/backbones/resnext.py +++ b/visdet/models/backbones/resnext.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. import math +from typing import Any from visdet.cv.cnn import build_conv_layer, build_norm_layer from visdet.models.backbones.resnet import Bottleneck as _Bottleneck @@ -17,7 +18,15 @@ class Bottleneck(_Bottleneck): expansion = 4 - def __init__(self, inplanes, planes, groups=1, base_width=4, base_channels=64, **kwargs): + def __init__( + self, + inplanes: int, + planes: int, + groups: int = 1, + base_width: int = 4, + base_channels: int = 64, + **kwargs: Any, + ) -> None: # Extract groups and base_width before calling parent self.groups = groups self.base_width = base_width diff --git a/visdet/models/detectors/single_stage.py b/visdet/models/detectors/single_stage.py index 9f08c805..b8eed213 100644 --- a/visdet/models/detectors/single_stage.py +++ b/visdet/models/detectors/single_stage.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any + from torch import Tensor from visdet.registry import MODELS @@ -26,12 +28,34 @@ def __init__( init_cfg: OptMultiConfig = None, ) -> None: super().__init__(data_preprocessor=data_preprocessor, init_cfg=init_cfg) - self.backbone = MODELS.build(backbone) + + backbone_cfg: dict[Any, Any] + if isinstance(backbone, dict): + backbone_cfg = backbone + else: + backbone_cfg = {"type": backbone} + self.backbone = MODELS.build(backbone_cfg) + if neck is not None: - self.neck = MODELS.build(neck) - bbox_head.update(train_cfg=train_cfg) - bbox_head.update(test_cfg=test_cfg) - self.bbox_head = MODELS.build(bbox_head) + neck_cfg: dict[Any, Any] + if isinstance(neck, dict): + neck_cfg = neck + else: + neck_cfg = {"type": neck} + self.neck = MODELS.build(neck_cfg) + + if bbox_head is None: + raise ValueError("SingleStageDetector requires a bbox_head") + + bbox_head_cfg: dict[Any, Any] + if isinstance(bbox_head, dict): + bbox_head_cfg = bbox_head.copy() + else: + bbox_head_cfg = {"type": bbox_head} + + bbox_head_cfg.update(train_cfg=train_cfg) + bbox_head_cfg.update(test_cfg=test_cfg) + self.bbox_head = MODELS.build(bbox_head_cfg) self.train_cfg = train_cfg self.test_cfg = test_cfg diff --git a/visdet/models/layers/__init__.py b/visdet/models/layers/__init__.py index 2eca53b9..f0475ae2 100644 --- a/visdet/models/layers/__init__.py +++ b/visdet/models/layers/__init__.py @@ -1,8 +1,14 @@ # ruff: noqa +from __future__ import annotations + import math +from typing import Any, cast + import torch import torch.nn as nn import torch.nn.functional as F +from torch import Tensor + from visdet.engine.model import BaseModule from visdet.engine.utils import to_2tuple from visdet.models.layers.bbox_nms import multiclass_nms @@ -28,20 +34,26 @@ class AdaptivePadding(nn.Module): pad zero around input. Default: "corner". """ - def __init__(self, kernel_size=1, stride=1, dilation=1, padding="corner"): + def __init__( + self, + kernel_size: int | tuple[int, int] = 1, + stride: int | tuple[int, int] = 1, + dilation: int | tuple[int, int] = 1, + padding: str = "corner", + ) -> None: super().__init__() assert padding in ("same", "corner") - kernel_size = to_2tuple(kernel_size) - stride = to_2tuple(stride) - dilation = to_2tuple(dilation) + kernel_size = cast(tuple[int, int], to_2tuple(kernel_size)) + stride = cast(tuple[int, int], to_2tuple(stride)) + dilation = cast(tuple[int, int], to_2tuple(dilation)) self.padding = padding self.kernel_size = kernel_size self.stride = stride self.dilation = dilation - def get_pad_shape(self, input_shape): + def get_pad_shape(self, input_shape: tuple[int, int]) -> tuple[int, int]: input_h, input_w = input_shape kernel_h, kernel_w = self.kernel_size stride_h, stride_w = self.stride @@ -57,8 +69,8 @@ def get_pad_shape(self, input_shape): ) return pad_h, pad_w - def forward(self, x): - pad_h, pad_w = self.get_pad_shape(x.size()[-2:]) + def forward(self, x: Tensor) -> Tensor: + pad_h, pad_w = self.get_pad_shape((int(x.size(-2)), int(x.size(-1)))) if pad_h > 0 or pad_w > 0: if self.padding == "corner": x = F.pad(x, [0, pad_w, 0, pad_h]) @@ -72,20 +84,20 @@ class PatchEmbed(BaseModule): def __init__( self, - img_size=224, - patch_size=4, - in_channels=3, - embed_dims=96, - norm_cfg=None, - conv_type="Conv2d", - kernel_size=None, - stride=None, - padding="corner", - dilation=1, - bias=True, - input_size=None, - init_cfg=None, - ): + img_size: int = 224, + patch_size: int = 4, + in_channels: int = 3, + embed_dims: int = 96, + norm_cfg: dict | None = None, + conv_type: str = "Conv2d", + kernel_size: int | None = None, + stride: int | None = None, + padding: str | int | tuple[int, int] = "corner", + dilation: int | tuple[int, int] = 1, + bias: bool = True, + input_size: tuple[int, int] | None = None, + init_cfg: dict | list[dict] | None = None, + ) -> None: super().__init__(init_cfg=init_cfg) # Handle different parameter names if kernel_size is not None: @@ -96,15 +108,16 @@ def __init__( self.img_size = img_size self.patch_size = patch_size - kernel_size = to_2tuple(patch_size) - stride = to_2tuple(stride) - dilation = to_2tuple(dilation) + kernel_size_tuple = cast(tuple[int, int], to_2tuple(patch_size)) + stride_tuple = cast(tuple[int, int], to_2tuple(stride)) + dilation_tuple = cast(tuple[int, int], to_2tuple(dilation)) + self.adap_padding: AdaptivePadding | None if isinstance(padding, str): self.adap_padding = AdaptivePadding( - kernel_size=kernel_size, - stride=stride, - dilation=dilation, + kernel_size=kernel_size_tuple, + stride=stride_tuple, + dilation=dilation_tuple, padding=padding, ) # disable the padding of conv @@ -112,31 +125,32 @@ def __init__( else: self.adap_padding = None - padding = to_2tuple(padding) + padding_tuple = cast(tuple[int, int], to_2tuple(padding)) self.proj = nn.Conv2d( in_channels, embed_dims, - kernel_size=kernel_size, - stride=stride, - padding=padding, - dilation=dilation, + kernel_size=kernel_size_tuple, + stride=stride_tuple, + padding=padding_tuple, + dilation=dilation_tuple, bias=bias, ) + self.norm: nn.LayerNorm | None if norm_cfg is not None: self.norm = nn.LayerNorm(embed_dims) else: self.norm = None - def forward(self, x): + def forward(self, x: Tensor) -> tuple[Tensor, tuple[int, int]]: B, C, H, W = x.shape if self.adap_padding: x = self.adap_padding(x) x = self.proj(x) # B, embed_dims, H/patch_size, W/patch_size - Hp, Wp = x.shape[2], x.shape[3] + Hp, Wp = int(x.shape[2]), int(x.shape[3]) x = x.flatten(2).transpose(1, 2) # B, H*W/patch_size^2, embed_dims if self.norm is not None: x = self.norm(x) diff --git a/visdet/models/layers/bbox_nms.py b/visdet/models/layers/bbox_nms.py index 09ec323f..dd358fc8 100644 --- a/visdet/models/layers/bbox_nms.py +++ b/visdet/models/layers/bbox_nms.py @@ -1,5 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any + import torch from torch import Tensor @@ -8,6 +10,12 @@ from visdet.utils import ConfigType +def _cfg_to_dict(cfg: ConfigType) -> dict[str, Any]: + if isinstance(cfg, dict): + return cfg + return {"type": cfg} + + def multiclass_nms( multi_bboxes: Tensor, multi_scores: Tensor, @@ -89,7 +97,7 @@ def multiclass_nms( else: return dets, labels - dets, keep = batched_nms(bboxes, scores, labels, nms_cfg) + dets, keep = batched_nms(bboxes, scores, labels, _cfg_to_dict(nms_cfg)) if max_num > 0: dets = dets[:max_num] diff --git a/visdet/models/layers/csp_layer.py b/visdet/models/layers/csp_layer.py index 82fbc326..af29101e 100644 --- a/visdet/models/layers/csp_layer.py +++ b/visdet/models/layers/csp_layer.py @@ -23,10 +23,10 @@ def __init__( expansion: float = 0.5, add_identity: bool = True, use_depthwise: bool = False, - conv_cfg: OptConfigType = None, - norm_cfg: ConfigType = dict(type="BN", momentum=0.03, eps=0.001), - act_cfg: ConfigType = dict(type="Swish"), - init_cfg: OptMultiConfig = None, + conv_cfg: dict | None = None, + norm_cfg: dict = dict(type="BN", momentum=0.03, eps=0.001), + act_cfg: dict = dict(type="Swish"), + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__(init_cfg=init_cfg) hidden_channels = int(out_channels * expansion) @@ -69,10 +69,10 @@ def __init__( add_identity: bool = True, use_depthwise: bool = False, kernel_size: int = 5, - conv_cfg: OptConfigType = None, - norm_cfg: ConfigType = dict(type="BN", momentum=0.03, eps=0.001), - act_cfg: ConfigType = dict(type="SiLU"), - init_cfg: OptMultiConfig = None, + conv_cfg: dict | None = None, + norm_cfg: dict = dict(type="BN", momentum=0.03, eps=0.001), + act_cfg: dict = dict(type="SiLU"), + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__(init_cfg=init_cfg) hidden_channels = int(out_channels * expansion) @@ -118,10 +118,10 @@ def __init__( use_depthwise: bool = False, use_cspnext_block: bool = False, channel_attention: bool = False, - conv_cfg: OptConfigType = None, - norm_cfg: ConfigType = dict(type="BN", momentum=0.03, eps=0.001), - act_cfg: ConfigType = dict(type="Swish"), - init_cfg: OptMultiConfig = None, + conv_cfg: dict | None = None, + norm_cfg: dict = dict(type="BN", momentum=0.03, eps=0.001), + act_cfg: dict = dict(type="Swish"), + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__(init_cfg=init_cfg) block = CSPNeXtBlock if use_cspnext_block else DarknetBottleneck diff --git a/visdet/models/layers/normed_predictor.py b/visdet/models/layers/normed_predictor.py index 70f50a1f..0e7cd78e 100644 --- a/visdet/models/layers/normed_predictor.py +++ b/visdet/models/layers/normed_predictor.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any, Callable, cast + import torch import torch.nn as nn import torch.nn.functional as F @@ -25,7 +27,7 @@ def __init__( self, *args, tempearture: float = 20, - power: int = 1.0, + power: float = 1.0, eps: float = 1e-6, **kwargs, ) -> None: @@ -67,7 +69,7 @@ def __init__( self, *args, tempearture: float = 20, - power: int = 1.0, + power: float = 1.0, eps: float = 1e-6, norm_over_kernel: bool = False, **kwargs, @@ -90,11 +92,10 @@ def forward(self, x: Tensor) -> Tensor: x_ = x / (x.norm(dim=1, keepdim=True).pow(self.power) + self.eps) x_ = x_ * self.tempearture - if hasattr(self, "conv2d_forward"): - x_ = self.conv2d_forward(x_, weight_) + conv2d_forward = getattr(self, "conv2d_forward", None) + if callable(conv2d_forward): + conv2d_forward_fn = cast(Callable[[Tensor, Tensor], Tensor], conv2d_forward) + x_ = conv2d_forward_fn(x_, weight_) else: - if digit_version(torch.__version__) >= digit_version("1.8"): - x_ = self._conv_forward(x_, weight_, self.bias) - else: - x_ = self._conv_forward(x_, weight_) + x_ = self._conv_forward(x_, weight_, self.bias) return x_ diff --git a/visdet/models/losses/cross_entropy_loss.py b/visdet/models/losses/cross_entropy_loss.py index 9a6c0878..b5796cb1 100644 --- a/visdet/models/losses/cross_entropy_loss.py +++ b/visdet/models/losses/cross_entropy_loss.py @@ -1,9 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. +from __future__ import annotations + import warnings +from typing import Any, cast import torch import torch.nn as nn import torch.nn.functional as F +from torch import Tensor from visdet.models.losses.accuracy import accuracy from visdet.models.losses.utils import weight_reduce_loss @@ -11,15 +15,15 @@ def cross_entropy( - pred, - label, - weight=None, - reduction="mean", - avg_factor=None, - class_weight=None, - ignore_index=-100, - avg_non_ignore=False, -): + pred: Tensor, + label: Tensor, + weight: Tensor | None = None, + reduction: str = "mean", + avg_factor: float | None = None, + class_weight: Tensor | None = None, + ignore_index: int | None = -100, + avg_non_ignore: bool = False, +) -> Tensor: """Calculate the CrossEntropy loss. Args: @@ -48,7 +52,7 @@ def cross_entropy( # pytorch's official cross_entropy average loss over non-ignored elements # refer to https://github.com/pytorch/pytorch/blob/56b43f4fec1f76953f15a627694d4bba34588969/torch/nn/functional.py#L2660 if (avg_factor is None) and avg_non_ignore and reduction == "mean": - avg_factor = label.numel() - (label == ignore_index).sum().item() + avg_factor = float(label.numel() - (label == ignore_index).sum().item()) # apply weights and do the reduction if weight is not None: @@ -58,7 +62,12 @@ def cross_entropy( return loss -def _expand_onehot_labels(labels, label_weights, label_channels, ignore_index): +def _expand_onehot_labels( + labels: Tensor, + label_weights: Tensor | None, + label_channels: int, + ignore_index: int, +) -> tuple[Tensor, Tensor, Tensor]: """Expand onehot labels to match the size of prediction.""" bin_labels = labels.new_full((labels.size(0), label_channels), 0) valid_mask = (labels >= 0) & (labels != ignore_index) @@ -78,15 +87,15 @@ def _expand_onehot_labels(labels, label_weights, label_channels, ignore_index): def binary_cross_entropy( - pred, - label, - weight=None, - reduction="mean", - avg_factor=None, - class_weight=None, - ignore_index=-100, - avg_non_ignore=False, -): + pred: Tensor, + label: Tensor, + weight: Tensor | None = None, + reduction: str = "mean", + avg_factor: float | None = None, + class_weight: Tensor | None = None, + ignore_index: int | None = -100, + avg_non_ignore: bool = False, +) -> Tensor: """Calculate the binary CrossEntropy loss. Args: @@ -128,7 +137,7 @@ def binary_cross_entropy( # average loss over non-ignored elements if (avg_factor is None) and avg_non_ignore and reduction == "mean": - avg_factor = valid_mask.sum().item() + avg_factor = float(valid_mask.sum().item()) # weighted element-wise losses weight = weight.float() @@ -140,15 +149,15 @@ def binary_cross_entropy( def mask_cross_entropy( - pred, - target, - label, - reduction="mean", - avg_factor=None, - class_weight=None, - ignore_index=None, - **kwargs, -): + pred: Tensor, + target: Tensor, + label: Tensor, + reduction: str = "mean", + avg_factor: float | None = None, + class_weight: Tensor | None = None, + ignore_index: int | None = None, + **kwargs: Any, +) -> Tensor: """Calculate the CrossEntropy loss for masks. Args: @@ -196,14 +205,14 @@ def mask_cross_entropy( class CrossEntropyLoss(nn.Module): def __init__( self, - use_sigmoid=False, - use_mask=False, - reduction="mean", - class_weight=None, - ignore_index=None, - loss_weight=1.0, - avg_non_ignore=False, - ): + use_sigmoid: bool = False, + use_mask: bool = False, + reduction: str = "mean", + class_weight: list[float] | None = None, + ignore_index: int | None = None, + loss_weight: float = 1.0, + avg_non_ignore: bool = False, + ) -> None: """CrossEntropyLoss. Args: @@ -223,6 +232,7 @@ def __init__( """ super(CrossEntropyLoss, self).__init__() assert (use_sigmoid is False) or (use_mask is False) + self.cls_criterion: Any self.use_sigmoid = use_sigmoid self.use_mask = use_mask self.reduction = reduction @@ -252,14 +262,14 @@ def extra_repr(self): def forward( self, - cls_score, - label, - weight=None, - avg_factor=None, - reduction_override=None, - ignore_index=None, - **kwargs, - ): + cls_score: Tensor, + label: Tensor, + weight: Tensor | None = None, + avg_factor: float | None = None, + reduction_override: str | None = None, + ignore_index: int | None = None, + **kwargs: Any, + ) -> Tensor: """Forward function. Args: @@ -302,15 +312,15 @@ def forward( class CrossEntropyCustomLoss(CrossEntropyLoss): def __init__( self, - use_sigmoid=False, - use_mask=False, - reduction="mean", - num_classes=-1, - class_weight=None, - ignore_index=None, - loss_weight=1.0, - avg_non_ignore=False, - ): + use_sigmoid: bool = False, + use_mask: bool = False, + reduction: str = "mean", + num_classes: int = -1, + class_weight: list[float] | None = None, + ignore_index: int | None = None, + loss_weight: float = 1.0, + avg_non_ignore: bool = False, + ) -> None: """CrossEntropyCustomLoss. Args: @@ -331,6 +341,7 @@ def __init__( """ super(CrossEntropyCustomLoss, self).__init__() assert (use_sigmoid is False) or (use_mask is False) + self.cls_criterion: Any self.use_sigmoid = use_sigmoid self.use_mask = use_mask self.reduction = reduction @@ -364,14 +375,14 @@ def __init__( # custom accuracy of the classsifier self.custom_accuracy = True - def get_cls_channels(self, num_classes): + def get_cls_channels(self, num_classes: int) -> int: assert num_classes == self.num_classes if not self.use_sigmoid: return num_classes + 1 else: return num_classes - def get_activation(self, cls_score): + def get_activation(self, cls_score: Tensor) -> Tensor: fine_cls_score = cls_score[:, : self.num_classes] if not self.use_sigmoid: @@ -386,11 +397,10 @@ def get_activation(self, cls_score): return scores - def get_accuracy(self, cls_score, labels): + def get_accuracy(self, cls_score: Tensor, labels: Tensor) -> dict[str, Tensor]: fine_cls_score = cls_score[:, : self.num_classes] pos_inds = labels < self.num_classes - acc_classes = accuracy(fine_cls_score[pos_inds], labels[pos_inds]) - acc = dict() - acc["acc_classes"] = acc_classes + acc_classes = cast(Tensor, accuracy(fine_cls_score[pos_inds], labels[pos_inds])) + acc: dict[str, Tensor] = {"acc_classes": acc_classes} return acc diff --git a/visdet/models/necks/bfp.py b/visdet/models/necks/bfp.py index c96ac064..e9a5409c 100644 --- a/visdet/models/necks/bfp.py +++ b/visdet/models/necks/bfp.py @@ -1,4 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. +import torch import torch.nn.functional as F from torch import Tensor, nn @@ -49,12 +50,12 @@ def __init__( norm_cfg=self.norm_cfg, ) - def forward(self, inputs: tuple[Tensor]) -> tuple[Tensor]: + def forward(self, inputs: tuple[Tensor, ...]) -> tuple[Tensor, ...]: """Forward function.""" assert len(inputs) == self.num_levels # step 1: gather multi-level features by resize and average - feats = [] + feats: list[Tensor] = [] gather_size = inputs[self.refine_level].size()[2:] for i in range(self.num_levels): if i < self.refine_level: @@ -63,14 +64,14 @@ def forward(self, inputs: tuple[Tensor]) -> tuple[Tensor]: gathered = F.interpolate(inputs[i], size=gather_size, mode="nearest") feats.append(gathered) - bsf = sum(feats) / len(feats) + bsf = torch.stack(feats, dim=0).mean(dim=0) # step 2: refine gathered features if self.refine_type is not None: bsf = self.refine(bsf) # step 3: scatter refined features to multi-levels by a residual path - outs = [] + outs: list[Tensor] = [] for i in range(self.num_levels): out_size = inputs[i].size()[2:] if i < self.refine_level: diff --git a/visdet/models/necks/cspnext_pafpn.py b/visdet/models/necks/cspnext_pafpn.py index 5f496c87..f5393830 100644 --- a/visdet/models/necks/cspnext_pafpn.py +++ b/visdet/models/necks/cspnext_pafpn.py @@ -14,7 +14,7 @@ from visdet.engine.model import BaseModule from visdet.models.layers import CSPLayer from visdet.registry import MODELS -from visdet.utils.typing_utils import ConfigType, OptMultiConfig +from visdet.utils.typing_utils import OptMultiConfig @MODELS.register_module() @@ -28,11 +28,11 @@ def __init__( num_csp_blocks: int = 3, use_depthwise: bool = False, expand_ratio: float = 0.5, - upsample_cfg: ConfigType = dict(scale_factor=2, mode="nearest"), + upsample_cfg: dict = dict(scale_factor=2, mode="nearest"), conv_cfg: dict | None = None, - norm_cfg: ConfigType = dict(type="BN", momentum=0.03, eps=0.001), - act_cfg: ConfigType = dict(type="Swish"), - init_cfg: OptMultiConfig = dict( + norm_cfg: dict = dict(type="BN", momentum=0.03, eps=0.001), + act_cfg: dict = dict(type="Swish"), + init_cfg: dict | list[dict] | None = dict( type="Kaiming", layer="Conv2d", a=math.sqrt(5), diff --git a/visdet/models/roi_heads/base_roi_head.py b/visdet/models/roi_heads/base_roi_head.py index 01d1a568..2223a5ae 100644 --- a/visdet/models/roi_heads/base_roi_head.py +++ b/visdet/models/roi_heads/base_roi_head.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. from abc import ABCMeta, abstractmethod +from typing import Any from torch import Tensor @@ -21,13 +22,17 @@ def __init__( shared_head: OptConfigType = None, train_cfg: OptConfigType = None, test_cfg: OptConfigType = None, - init_cfg: OptMultiConfig = None, + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__(init_cfg=init_cfg) self.train_cfg = train_cfg self.test_cfg = test_cfg if shared_head is not None: - self.shared_head = MODELS.build(shared_head) + if isinstance(shared_head, dict): + shared_head_cfg: dict[Any, Any] = shared_head + else: + shared_head_cfg = {"type": shared_head} + self.shared_head = MODELS.build(shared_head_cfg) if bbox_head is not None: self.init_bbox_head(bbox_roi_extractor, bbox_head) @@ -67,10 +72,31 @@ def init_assigner_sampler(self, *args, **kwargs): """Initialize assigner and sampler.""" pass + @abstractmethod + def predict_bbox( + self, + x: tuple[Tensor, ...], + batch_img_metas: list[dict], + rpn_results_list: InstanceList, + rcnn_test_cfg: OptConfigType, + rescale: bool = False, + ) -> InstanceList: + """Predict bbox results from features.""" + + @abstractmethod + def predict_mask( + self, + x: tuple[Tensor, ...], + batch_img_metas: list[dict], + results_list: InstanceList, + rescale: bool = False, + ) -> InstanceList: + """Predict mask results from features.""" + @abstractmethod def loss( self, - x: tuple[Tensor], + x: tuple[Tensor, ...], rpn_results_list: InstanceList, batch_data_samples: SampleList, **kwargs, @@ -80,7 +106,7 @@ def loss( def predict( self, - x: tuple[Tensor], + x: tuple[Tensor, ...], rpn_results_list: InstanceList, batch_data_samples: SampleList, rescale: bool = False, diff --git a/visdet/models/roi_heads/bbox_heads/bbox_head.py b/visdet/models/roi_heads/bbox_heads/bbox_head.py index eb26c7a4..c4017dc6 100644 --- a/visdet/models/roi_heads/bbox_heads/bbox_head.py +++ b/visdet/models/roi_heads/bbox_heads/bbox_head.py @@ -1,5 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any, cast + import torch import torch.nn as nn import torch.nn.functional as F @@ -18,6 +20,12 @@ from visdet.utils import ConfigType, InstanceList, OptMultiConfig +def _cfg_to_dict(cfg: ConfigType) -> dict[str, Any]: + if isinstance(cfg, dict): + return cfg + return {"type": cfg} + + @MODELS.register_module() class BBoxHead(BaseModule): """Simplest RoI head, with only two fc layers for classification and @@ -44,14 +52,14 @@ def __init__( cls_predictor_cfg: ConfigType = dict(type="Linear"), loss_cls: ConfigType = dict(type="CrossEntropyLoss", use_sigmoid=False, loss_weight=1.0), loss_bbox: ConfigType = dict(type="SmoothL1Loss", beta=1.0, loss_weight=1.0), - init_cfg: OptMultiConfig = None, + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__(init_cfg=init_cfg) assert with_cls or with_reg self.with_avg_pool = with_avg_pool self.with_cls = with_cls self.with_reg = with_reg - self.roi_feat_size = _pair(roi_feat_size) + self.roi_feat_size = cast(tuple[int, int], _pair(roi_feat_size)) self.roi_feat_area = self.roi_feat_size[0] * self.roi_feat_size[1] self.in_channels = in_channels self.num_classes = num_classes @@ -61,9 +69,9 @@ def __init__( self.reg_predictor_cfg = reg_predictor_cfg self.cls_predictor_cfg = cls_predictor_cfg - self.bbox_coder = TASK_UTILS.build(bbox_coder) - self.loss_cls = MODELS.build(loss_cls) - self.loss_bbox = MODELS.build(loss_bbox) + self.bbox_coder = TASK_UTILS.build(_cfg_to_dict(bbox_coder)) + self.loss_cls = MODELS.build(_cfg_to_dict(loss_cls)) + self.loss_bbox = MODELS.build(_cfg_to_dict(loss_bbox)) in_channels = self.in_channels if self.with_avg_pool: @@ -76,15 +84,14 @@ def __init__( cls_channels = self.loss_cls.get_cls_channels(self.num_classes) else: cls_channels = num_classes + 1 - cls_predictor_cfg_ = self.cls_predictor_cfg.copy() + cls_predictor_cfg_ = _cfg_to_dict(self.cls_predictor_cfg).copy() cls_predictor_cfg_.update(in_features=in_channels, out_features=cls_channels) self.fc_cls = MODELS.build(cls_predictor_cfg_) if self.with_reg: box_dim = self.bbox_coder.encode_size out_dim_reg = box_dim if reg_class_agnostic else box_dim * num_classes - reg_predictor_cfg_ = self.reg_predictor_cfg.copy() - if isinstance(reg_predictor_cfg_, (dict, ConfigDict)): - reg_predictor_cfg_.update(in_features=in_channels, out_features=out_dim_reg) + reg_predictor_cfg_ = _cfg_to_dict(self.reg_predictor_cfg).copy() + reg_predictor_cfg_.update(in_features=in_channels, out_features=out_dim_reg) self.fc_reg = MODELS.build(reg_predictor_cfg_) self.debug_imgs = None if init_cfg is None: @@ -112,7 +119,7 @@ def custom_accuracy(self) -> bool: """get custom_accuracy from loss_cls.""" return getattr(self.loss_cls, "custom_accuracy", False) - def forward(self, x: tuple[Tensor]) -> tuple: + def forward(self, x: Tensor) -> tuple[Tensor | None, Tensor | None]: """Forward features from the upstream network. Args: @@ -144,8 +151,8 @@ def forward(self, x: tuple[Tensor]) -> tuple: def get_bboxes( self, rois: Tensor, - cls_score: Tensor, - bbox_pred: Tensor, + cls_score: Tensor | None, + bbox_pred: Tensor | None, img_shape: tuple | None = None, scale_factor: Tensor | None = None, rescale: bool = False, @@ -199,7 +206,11 @@ def get_bboxes( if cfg is None: return bboxes, scores else: - det_bboxes, det_labels = multiclass_nms(bboxes, scores, cfg.score_thr, cfg.nms, cfg.max_per_img) + assert scores is not None + det_bboxes, det_labels = cast( + tuple[Tensor, Tensor], + multiclass_nms(bboxes, scores, cfg.score_thr, cfg.nms, cfg.max_per_img, return_inds=False), + ) return det_bboxes, det_labels @@ -375,17 +386,25 @@ def loss_and_target( The targets are only used for cascade rcnn. """ - cls_reg_targets = self.get_targets(sampling_results, rcnn_train_cfg, concat=concat) + labels, label_weights, bbox_targets, bbox_weights = self.get_targets( + sampling_results, + rcnn_train_cfg, + concat=concat, + ) losses = self.loss( cls_score, bbox_pred, rois, - *cls_reg_targets, + labels, + label_weights, + bbox_targets, + bbox_weights, reduction_override=reduction_override, ) # cls_reg_targets is only for cascade rcnn - return dict(loss_bbox=losses, bbox_targets=cls_reg_targets) + # bbox_targets is only for cascade rcnn + return dict(loss_bbox=losses, bbox_targets=(labels, label_weights, bbox_targets, bbox_weights)) def loss( self, @@ -483,9 +502,9 @@ def loss( def predict_by_feat( self, - rois: tuple[Tensor], - cls_scores: tuple[Tensor], - bbox_preds: tuple[Tensor], + rois: tuple[Tensor, ...], + cls_scores: tuple[Tensor, ...], + bbox_preds: tuple[Tensor | None, ...], batch_img_metas: list[dict], rcnn_test_cfg: ConfigDict | None = None, rescale: bool = False, @@ -538,8 +557,8 @@ def predict_by_feat( def _predict_by_feat_single( self, roi: Tensor, - cls_score: Tensor, - bbox_pred: Tensor, + cls_score: Tensor | None, + bbox_pred: Tensor | None, img_meta: dict, rescale: bool = False, rcnn_test_cfg: ConfigDict | None = None, @@ -607,8 +626,9 @@ def _predict_by_feat_single( if rescale and bboxes.size(0) > 0: assert img_meta.get("scale_factor") is not None - scale_factor = [1 / s for s in img_meta["scale_factor"]] - bboxes = scale_boxes(bboxes, scale_factor) + w_scale = float(img_meta["scale_factor"][0]) + h_scale = float(img_meta["scale_factor"][1]) + bboxes = scale_boxes(bboxes, (1.0 / w_scale, 1.0 / h_scale)) # Get the inside tensor when `bboxes` is a box type bboxes = get_box_tensor(bboxes) @@ -621,13 +641,18 @@ def _predict_by_feat_single( results.bboxes = bboxes results.scores = scores else: - det_bboxes, det_labels = multiclass_nms( - bboxes, - scores, - rcnn_test_cfg["score_thr"], - rcnn_test_cfg["nms"], - rcnn_test_cfg["max_per_img"], - box_dim=box_dim, + assert scores is not None + det_bboxes, det_labels = cast( + tuple[Tensor, Tensor], + multiclass_nms( + bboxes, + scores, + rcnn_test_cfg["score_thr"], + rcnn_test_cfg["nms"], + rcnn_test_cfg["max_per_img"], + return_inds=False, + box_dim=box_dim, + ), ) results.bboxes = det_bboxes[:, :-1] results.scores = det_bboxes[:, -1] diff --git a/visdet/models/roi_heads/bbox_heads/convfc_bbox_head.py b/visdet/models/roi_heads/bbox_heads/convfc_bbox_head.py index 968b6326..2d08b49f 100644 --- a/visdet/models/roi_heads/bbox_heads/convfc_bbox_head.py +++ b/visdet/models/roi_heads/bbox_heads/convfc_bbox_head.py @@ -1,11 +1,13 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any + import torch.nn as nn from torch import Tensor from visdet.cv.cnn import ConvModule from visdet.engine.config import ConfigDict -from visdet.models.roi_heads.bbox_heads.bbox_head import BBoxHead +from visdet.models.roi_heads.bbox_heads.bbox_head import BBoxHead, _cfg_to_dict from visdet.registry import MODELS @@ -31,13 +33,13 @@ def __init__( num_reg_fcs: int = 0, conv_out_channels: int = 256, fc_out_channels: int = 1024, - conv_cfg: dict | ConfigDict | None = None, - norm_cfg: dict | ConfigDict | None = None, - init_cfg: dict | ConfigDict | None = None, - *args, - **kwargs, + conv_cfg: dict | None = None, + norm_cfg: dict | None = None, + init_cfg: dict | list[dict] | None = None, + **kwargs: Any, ) -> None: - super().__init__(*args, init_cfg=init_cfg, **kwargs) + kwargs.pop("init_cfg", None) + super().__init__(init_cfg=init_cfg, **kwargs) assert num_shared_convs + num_shared_fcs + num_cls_convs + num_cls_fcs + num_reg_convs + num_reg_fcs > 0 if num_cls_convs > 0 or num_reg_convs > 0: assert num_shared_fcs == 0 @@ -85,18 +87,21 @@ def __init__( cls_channels = self.loss_cls.get_cls_channels(self.num_classes) else: cls_channels = self.num_classes + 1 - cls_predictor_cfg_ = self.cls_predictor_cfg.copy() + cls_predictor_cfg_ = _cfg_to_dict(self.cls_predictor_cfg).copy() cls_predictor_cfg_.update(in_features=self.cls_last_dim, out_features=cls_channels) self.fc_cls = MODELS.build(cls_predictor_cfg_) if self.with_reg: box_dim = self.bbox_coder.encode_size out_dim_reg = box_dim if self.reg_class_agnostic else box_dim * self.num_classes - reg_predictor_cfg_ = self.reg_predictor_cfg.copy() - if isinstance(reg_predictor_cfg_, (dict, ConfigDict)): - reg_predictor_cfg_.update(in_features=self.reg_last_dim, out_features=out_dim_reg) + reg_predictor_cfg_ = _cfg_to_dict(self.reg_predictor_cfg).copy() + reg_predictor_cfg_.update(in_features=self.reg_last_dim, out_features=out_dim_reg) self.fc_reg = MODELS.build(reg_predictor_cfg_) if init_cfg is None: + if self.init_cfg is None: + self.init_cfg = [] + elif isinstance(self.init_cfg, dict): + self.init_cfg = [self.init_cfg] self.init_cfg += [ dict( type="Xavier", @@ -150,7 +155,7 @@ def _add_conv_fc_branch( last_layer_dim = self.fc_out_channels return branch_convs, branch_fcs, last_layer_dim - def forward(self, x: tuple[Tensor]) -> tuple: + def forward(self, x: Tensor) -> tuple[Tensor | None, Tensor | None]: """Forward features from the upstream network. Args: @@ -211,7 +216,7 @@ def forward(self, x: tuple[Tensor]) -> tuple: # reduce the dumb classifications errors @MODELS.register_module() class Shared2FCBBoxHead(ConvFCBBoxHead): - def __init__(self, fc_out_channels: int = 1024, *args, **kwargs) -> None: + def __init__(self, fc_out_channels: int = 1024, **kwargs: Any) -> None: super().__init__( num_shared_convs=0, num_shared_fcs=2, @@ -220,14 +225,13 @@ def __init__(self, fc_out_channels: int = 1024, *args, **kwargs) -> None: num_reg_convs=0, num_reg_fcs=0, fc_out_channels=fc_out_channels, - *args, **kwargs, ) @MODELS.register_module() class Shared4Conv1FCBBoxHead(ConvFCBBoxHead): - def __init__(self, fc_out_channels: int = 1024, *args, **kwargs) -> None: + def __init__(self, fc_out_channels: int = 1024, **kwargs: Any) -> None: super().__init__( num_shared_convs=4, num_shared_fcs=1, @@ -236,6 +240,5 @@ def __init__(self, fc_out_channels: int = 1024, *args, **kwargs) -> None: num_reg_convs=0, num_reg_fcs=0, fc_out_channels=fc_out_channels, - *args, **kwargs, ) diff --git a/visdet/models/roi_heads/bbox_heads/double_bbox_head.py b/visdet/models/roi_heads/bbox_heads/double_bbox_head.py index d3be4901..2392dc54 100644 --- a/visdet/models/roi_heads/bbox_heads/double_bbox_head.py +++ b/visdet/models/roi_heads/bbox_heads/double_bbox_head.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch.nn as nn +from torch import Tensor from visdet.cv.cnn import ConvModule from visdet.models.backbones.resnet import Bottleneck @@ -117,9 +118,11 @@ def __init__( self.fc_cls = nn.Linear(self.fc_out_channels, self.num_classes + 1) self.relu = nn.ReLU(inplace=True) - def forward(self, x_cls, x_reg): + def forward(self, x_cls: Tensor, x_reg: Tensor | None = None): + x_reg_tensor = x_cls if x_reg is None else x_reg + # conv head - x_conv = self.res_block(x_reg) + x_conv = self.res_block(x_reg_tensor) for conv in self.conv_branch: x_conv = conv(x_conv) diff --git a/visdet/models/roi_heads/dynamic_roi_head.py b/visdet/models/roi_heads/dynamic_roi_head.py index faf602d2..e90e2e4d 100644 --- a/visdet/models/roi_heads/dynamic_roi_head.py +++ b/visdet/models/roi_heads/dynamic_roi_head.py @@ -1,10 +1,16 @@ # Copyright (c) OpenMMLab. All rights reserved. +from __future__ import annotations + +from typing import Any, cast + import numpy as np import torch +from torch import Tensor +from visdet.engine.structures import InstanceData +from visdet.models.roi_heads.standard_roi_head import StandardRoIHead from visdet.registry import MODELS - -from .standard_roi_head import StandardRoIHead +from visdet.structures import DetDataSample EPS = 1e-15 @@ -20,7 +26,13 @@ def __init__(self, **kwargs) -> None: # the beta history of the past `update_iter_interval` iterations self.beta_history = [] - def loss(self, x, rpn_results_list, batch_data_samples): + def loss( + self, + x: tuple[Tensor, ...], + rpn_results_list: list[InstanceData], + batch_data_samples: list[DetDataSample], + **kwargs: Any, + ) -> dict[Any, Any]: """Calculate losses from a batch of inputs and data samples.""" num_imgs = len(batch_data_samples) batch_gt_instances = [data_samples.gt_instances for data_samples in batch_data_samples] @@ -28,6 +40,14 @@ def loss(self, x, rpn_results_list, batch_data_samples): getattr(data_samples, "ignored_instances", None) for data_samples in batch_data_samples ] + assert self.train_cfg is not None + assert self.bbox_assigner is not None + assert self.bbox_sampler is not None + + train_cfg = cast(dict[str, Any], self.train_cfg) + bbox_assigner = self.bbox_assigner + bbox_sampler = self.bbox_sampler + sampling_results = [] cur_iou = [] for i in range(num_imgs): @@ -36,12 +56,12 @@ def loss(self, x, rpn_results_list, batch_data_samples): if "bboxes" in rpn_results: rpn_results.priors = rpn_results.pop("bboxes") - assign_result = self.bbox_assigner.assign(rpn_results, batch_gt_instances[i], batch_gt_instances_ignore[i]) - sampling_result = self.bbox_sampler.sample(assign_result, rpn_results, batch_gt_instances[i]) + assign_result = bbox_assigner.assign(rpn_results, batch_gt_instances[i], batch_gt_instances_ignore[i]) + sampling_result = bbox_sampler.sample(assign_result, rpn_results, batch_gt_instances[i]) sampling_results.append(sampling_result) # record the `iou_topk`-th largest IoU in an image - dynamic_cfg = self.train_cfg.get("dynamic_rcnn", self.train_cfg) + dynamic_cfg = train_cfg.get("dynamic_rcnn", train_cfg) iou_topk = min( dynamic_cfg.get("iou_topk", 75), len(assign_result.max_overlaps), @@ -60,7 +80,7 @@ def loss(self, x, rpn_results_list, batch_data_samples): bbox_results = self.bbox_loss(x, sampling_results) # update IoU threshold and SmoothL1 beta - dynamic_cfg = self.train_cfg.get("dynamic_rcnn", self.train_cfg) + dynamic_cfg = train_cfg.get("dynamic_rcnn", train_cfg) update_iter_interval = dynamic_cfg.get("update_iter_interval", 100) if len(self.iou_history) % update_iter_interval == 0: self.update_hyperparameters() @@ -68,6 +88,9 @@ def loss(self, x, rpn_results_list, batch_data_samples): return bbox_results def bbox_loss(self, x, sampling_results): + assert self.train_cfg is not None + train_cfg = cast(dict[str, Any], self.train_cfg) + from visdet.models.task_modules.assigners import get_box_tensor from visdet.structures.bbox import bbox2roi @@ -86,7 +109,7 @@ def bbox_loss(self, x, sampling_results): num_pos = len(pos_inds) if num_pos > 0: cur_target = bbox_targets_vals[pos_inds, :2].abs().mean(dim=1) - dynamic_cfg = self.train_cfg.get("dynamic_rcnn", self.train_cfg) + dynamic_cfg = train_cfg.get("dynamic_rcnn", train_cfg) beta_topk = min(dynamic_cfg.get("beta_topk", 10) * len(sampling_results), num_pos) cur_target = torch.kthvalue(cur_target, beta_topk)[0].item() self.beta_history.append(cur_target) @@ -94,14 +117,19 @@ def bbox_loss(self, x, sampling_results): loss_bbox = self.bbox_head.loss(bbox_results["cls_score"], bbox_results["bbox_pred"], rois, *bbox_targets) return loss_bbox - def update_hyperparameters(self): + def update_hyperparameters(self) -> None: """Update hyperparameters.""" - dynamic_cfg = self.train_cfg.get("dynamic_rcnn", self.train_cfg) + assert self.train_cfg is not None + train_cfg = cast(dict[str, Any], self.train_cfg) + dynamic_cfg = train_cfg.get("dynamic_rcnn", train_cfg) new_iou_thr = max(dynamic_cfg.get("initial_iou", 0.4), np.mean(self.iou_history)) self.iou_history = [] - self.bbox_assigner.pos_iou_thr = new_iou_thr - self.bbox_assigner.neg_iou_thr = new_iou_thr - self.bbox_assigner.min_pos_iou = new_iou_thr + + assert self.bbox_assigner is not None + bbox_assigner = self.bbox_assigner + bbox_assigner.pos_iou_thr = new_iou_thr + bbox_assigner.neg_iou_thr = new_iou_thr + bbox_assigner.min_pos_iou = new_iou_thr if len(self.beta_history) > 0: if np.median(self.beta_history) < EPS: diff --git a/visdet/models/roi_heads/mask_heads/fcn_mask_head.py b/visdet/models/roi_heads/mask_heads/fcn_mask_head.py index 8f80ffbd..b50f1947 100644 --- a/visdet/models/roi_heads/mask_heads/fcn_mask_head.py +++ b/visdet/models/roi_heads/mask_heads/fcn_mask_head.py @@ -1,5 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any, cast + import numpy as np import torch import torch.nn as nn @@ -17,6 +19,13 @@ from visdet.structures.mask import mask_target from visdet.utils import ConfigType, InstanceList, OptConfigType, OptMultiConfig + +def _cfg_to_dict(cfg: ConfigType) -> dict[str, Any]: + if isinstance(cfg, dict): + return cfg + return {"type": cfg} + + BYTES_PER_FLOAT = 4 # TODO: This memory limit may be too much or too little. It would be better to # determine it based on available resources. @@ -33,24 +42,24 @@ def __init__( conv_kernel_size: int = 3, conv_out_channels: int = 256, num_classes: int = 80, - class_agnostic: int = False, + class_agnostic: bool = False, upsample_cfg: ConfigType = dict(type="deconv", scale_factor=2), - conv_cfg: OptConfigType = None, - norm_cfg: OptConfigType = None, + conv_cfg: dict | None = None, + norm_cfg: dict | None = None, predictor_cfg: ConfigType = dict(type="Conv"), loss_mask: ConfigType = dict(type="CrossEntropyLoss", use_mask=True, loss_weight=1.0), - init_cfg: OptMultiConfig = None, + init_cfg: dict | list[dict] | None = None, ) -> None: assert init_cfg is None, "To prevent abnormal initialization behavior, init_cfg is not allowed to be set" super().__init__(init_cfg=init_cfg) - self.upsample_cfg = upsample_cfg.copy() + self.upsample_cfg = _cfg_to_dict(upsample_cfg).copy() if self.upsample_cfg["type"] not in [None, "deconv", "nearest", "bilinear"]: raise ValueError( f'Invalid upsample method {self.upsample_cfg["type"]}, accepted methods are "deconv", "nearest", "bilinear"' ) self.num_convs = num_convs # WARN: roi_feat_size is reserved and not used - self.roi_feat_size = _pair(roi_feat_size) + self.roi_feat_size = cast(tuple[int, int], _pair(roi_feat_size)) self.in_channels = in_channels self.conv_kernel_size = conv_kernel_size self.conv_out_channels = conv_out_channels @@ -60,8 +69,8 @@ def __init__( self.class_agnostic = class_agnostic self.conv_cfg = conv_cfg self.norm_cfg = norm_cfg - self.predictor_cfg = predictor_cfg - self.loss_mask = MODELS.build(loss_mask) + self.predictor_cfg = _cfg_to_dict(predictor_cfg) + self.loss_mask = MODELS.build(_cfg_to_dict(loss_mask)) self.convs = ModuleList() for i in range(self.num_convs): @@ -112,8 +121,8 @@ def init_weights(self) -> None: if m is None: continue elif hasattr(m, "weight") and hasattr(m, "bias"): - nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") - nn.init.constant_(m.bias, 0) + nn.init.kaiming_normal_(cast(Tensor, m.weight), mode="fan_out", nonlinearity="relu") + nn.init.constant_(cast(Tensor, m.bias), 0) def forward(self, x: Tensor) -> Tensor: """Forward features from the upstream network. @@ -157,7 +166,7 @@ def get_targets( pos_assigned_gt_inds = [res.pos_assigned_gt_inds for res in sampling_results] gt_masks = [getattr(res, "instance_masks", getattr(res, "masks", None)) for res in batch_gt_instances] mask_targets = mask_target(pos_proposals, pos_assigned_gt_inds, gt_masks, rcnn_train_cfg) - return mask_targets + return cast(Tensor, mask_targets) def loss_and_target( self, @@ -243,7 +252,7 @@ def predict_by_feat( for img_id in range(len(batch_img_metas)): img_meta = batch_img_metas[img_id] results = results_list[img_id] - bboxes = results.bboxes + bboxes = cast(Tensor, getattr(results, "bboxes")) if bboxes.shape[0] == 0: results_list[img_id] = empty_instances( [img_meta], @@ -256,7 +265,7 @@ def predict_by_feat( im_mask = self._predict_by_feat_single( mask_preds=mask_preds[img_id], bboxes=bboxes, - labels=results.labels, + labels=cast(Tensor, getattr(results, "labels")), img_meta=img_meta, rcnn_test_cfg=rcnn_test_cfg, rescale=rescale, @@ -435,8 +444,13 @@ def _do_paste_mask(masks: Tensor, boxes: Tensor, img_h: int, img_w: int, skip_em N = masks.shape[0] - img_y = torch.arange(y0_int, y1_int, device=device).to(torch.float32) + 0.5 - img_x = torch.arange(x0_int, x1_int, device=device).to(torch.float32) + 0.5 + y0i = int(y0_int.item()) if isinstance(y0_int, torch.Tensor) else int(y0_int) + y1i = int(y1_int.item()) if isinstance(y1_int, torch.Tensor) else int(y1_int) + x0i = int(x0_int.item()) if isinstance(x0_int, torch.Tensor) else int(x0_int) + x1i = int(x1_int.item()) if isinstance(x1_int, torch.Tensor) else int(x1_int) + + img_y = torch.arange(y0i, y1i, device=device).to(torch.float32) + 0.5 + img_x = torch.arange(x0i, x1i, device=device).to(torch.float32) + 0.5 img_y = (img_y - y0) / (y1 - y0) * 2 - 1 img_x = (img_x - x0) / (x1 - x0) * 2 - 1 # img_x, img_y have shapes (N, w), (N, h) diff --git a/visdet/models/roi_heads/mask_heads/mask_iou_head.py b/visdet/models/roi_heads/mask_heads/mask_iou_head.py index f5af59a4..8134c231 100644 --- a/visdet/models/roi_heads/mask_heads/mask_iou_head.py +++ b/visdet/models/roi_heads/mask_heads/mask_iou_head.py @@ -15,7 +15,7 @@ def __init__( self, num_convs: int = 4, num_fcs: int = 2, - roi_feat_size: int = 14, + roi_feat_size: int | tuple[int, int] = 14, in_channels: int = 256, conv_out_channels: int = 256, fc_out_channels: int = 1024, @@ -39,8 +39,10 @@ def __init__( self.convs.append(nn.Conv2d(ch, self.conv_out_channels, 3, stride=stride, padding=1)) if isinstance(roi_feat_size, int): - roi_feat_size = (roi_feat_size, roi_feat_size) - pooled_area = (roi_feat_size[0] // 2) * (roi_feat_size[1] // 2) + roi_feat_size_tuple = (roi_feat_size, roi_feat_size) + else: + roi_feat_size_tuple = roi_feat_size + pooled_area = (roi_feat_size_tuple[0] // 2) * (roi_feat_size_tuple[1] // 2) self.fcs = nn.ModuleList() for i in range(num_fcs): ch = self.conv_out_channels * pooled_area if i == 0 else self.fc_out_channels @@ -104,17 +106,17 @@ def _get_area_ratio(self, pos_proposals: Tensor, pos_assigned_gt_inds: Tensor, g if num_pos > 0: area_ratios = [] proposals_np = pos_proposals.cpu().numpy() - pos_assigned_gt_inds = pos_assigned_gt_inds.cpu().numpy() + pos_assigned_gt_inds_np = pos_assigned_gt_inds.cpu().numpy() # compute mask areas of gt instances gt_instance_mask_area = gt_masks.areas for i in range(num_pos): - gt_mask = gt_masks[pos_assigned_gt_inds[i]] + gt_mask = gt_masks[pos_assigned_gt_inds_np[i]] # crop the gt mask inside the proposal bbox = proposals_np[i, :].astype(np.int32) gt_mask_in_proposal = gt_mask.crop(bbox) - ratio = gt_mask_in_proposal.areas[0] / (gt_instance_mask_area[pos_assigned_gt_inds[i]] + 1e-7) + ratio = gt_mask_in_proposal.areas[0] / (gt_instance_mask_area[pos_assigned_gt_inds_np[i]] + 1e-7) area_ratios.append(ratio) area_ratios = torch.from_numpy(np.stack(area_ratios)).float().to(pos_proposals.device) else: @@ -125,6 +127,6 @@ def get_mask_scores(self, mask_iou_pred: Tensor, det_bboxes: Tensor, det_labels: """Get the mask scores.""" inds = range(det_labels.size(0)) mask_scores = mask_iou_pred[inds, det_labels] * det_bboxes[inds, -1] - mask_scores = mask_scores.cpu().numpy() - det_labels = det_labels.cpu().numpy() - return [mask_scores[det_labels == i] for i in range(self.num_classes)] + mask_scores_np = mask_scores.cpu().numpy() + det_labels_np = det_labels.cpu().numpy() + return [mask_scores_np[det_labels_np == i] for i in range(self.num_classes)] diff --git a/visdet/models/roi_heads/mask_scoring_roi_head.py b/visdet/models/roi_heads/mask_scoring_roi_head.py index 78e9d309..f3a54009 100644 --- a/visdet/models/roi_heads/mask_scoring_roi_head.py +++ b/visdet/models/roi_heads/mask_scoring_roi_head.py @@ -1,10 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import cast + import torch +from visdet.models.roi_heads.standard_roi_head import StandardRoIHead from visdet.registry import MODELS -from .standard_roi_head import StandardRoIHead - @MODELS.register_module() class MaskScoringRoIHead(StandardRoIHead): @@ -31,14 +32,6 @@ def mask_loss(self, x, sampling_results, bbox_feats, batch_gt_instances): mask_iou_pred = self.mask_iou_head(mask_results["mask_feats"], pos_mask_pred) pos_mask_iou_pred = mask_iou_pred[range(mask_iou_pred.size(0)), pos_labels] - # Get mask targets for mask iou head - # We need mask targets from standard mask head - # Wait, standard mask head might not return mask targets in its mask_loss results. - # Let's check FCNMaskHead.loss_and_target - - # Actually, we can re-calculate them or ensure they are available. - # For now, assume we can get them. - # Re-get mask targets as they are needed for MaskIoUHead mask_targets = self.mask_head.get_targets(sampling_results, batch_gt_instances, self.train_cfg) gt_masks = [getattr(res, "instance_masks", getattr(res, "masks", None)) for res in batch_gt_instances] @@ -63,9 +56,11 @@ def predict_mask(self, x, batch_img_metas, results_list, rescale=False): return results_list # Get mask scores - from visdet.structures.bbox import bbox2roi + from visdet.structures.bbox import BaseBoxes, bbox2roi - bboxes = [res.bboxes for res in results_list] + bboxes: list[torch.Tensor | BaseBoxes] = [ + cast(torch.Tensor | BaseBoxes, getattr(res, "bboxes")) for res in results_list + ] mask_rois = bbox2roi(bboxes) # We need mask features and predictions to get mask IoU scores @@ -73,7 +68,7 @@ def predict_mask(self, x, batch_img_metas, results_list, rescale=False): mask_feats = mask_results["mask_feats"] mask_preds = mask_results["mask_preds"] - concat_det_labels = torch.cat([res.labels for res in results_list]) + concat_det_labels = torch.cat([cast(torch.Tensor, getattr(res, "labels")) for res in results_list]) mask_iou_pred = self.mask_iou_head( mask_feats, @@ -87,6 +82,10 @@ def predict_mask(self, x, batch_img_metas, results_list, rescale=False): for i in range(len(results_list)): if len(results_list[i]) > 0: # Get mask scores from mask IoU head - self.mask_iou_head.get_mask_scores(mask_iou_preds[i], results_list[i].bboxes, results_list[i].labels) + self.mask_iou_head.get_mask_scores( + mask_iou_preds[i], + cast(torch.Tensor, getattr(results_list[i], "bboxes")), + cast(torch.Tensor, getattr(results_list[i], "labels")), + ) return results_list diff --git a/visdet/models/roi_heads/roi_extractors/base_roi_extractor.py b/visdet/models/roi_heads/roi_extractors/base_roi_extractor.py index bb1c1f40..5467daf5 100644 --- a/visdet/models/roi_heads/roi_extractors/base_roi_extractor.py +++ b/visdet/models/roi_heads/roi_extractors/base_roi_extractor.py @@ -1,5 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. from abc import ABCMeta, abstractmethod +from typing import Any import torch import torch.nn as nn @@ -27,7 +28,7 @@ def __init__( roi_layer: ConfigType, out_channels: int, featmap_strides: list[int], - init_cfg: OptMultiConfig = None, + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__(init_cfg=init_cfg) self.roi_layers = self.build_roi_layers(roi_layer, featmap_strides) @@ -56,7 +57,11 @@ def build_roi_layers(self, layer_cfg: ConfigType, featmap_strides: list[int]) -> feature map. """ - cfg = layer_cfg.copy() + cfg: dict[str, Any] + if isinstance(layer_cfg, dict): + cfg = layer_cfg.copy() + else: + cfg = {"type": layer_cfg} layer_type = cfg.pop("type") if isinstance(layer_type, str): assert hasattr(ops, layer_type) diff --git a/visdet/models/roi_heads/roi_extractors/single_level_roi_extractor.py b/visdet/models/roi_heads/roi_extractors/single_level_roi_extractor.py index 2c588232..5547ee4d 100644 --- a/visdet/models/roi_heads/roi_extractors/single_level_roi_extractor.py +++ b/visdet/models/roi_heads/roi_extractors/single_level_roi_extractor.py @@ -1,5 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import cast + import torch from torch import Tensor @@ -33,7 +35,7 @@ def __init__( out_channels: int, featmap_strides: list[int], finest_scale: int = 56, - init_cfg: OptMultiConfig = None, + init_cfg: dict | list[dict] | None = None, ) -> None: super().__init__( roi_layer=roi_layer, @@ -78,9 +80,12 @@ def forward(self, feats: tuple[Tensor], rois: Tensor, roi_scale_factor: float | """ # convert fp32 to fp16 when amp is on rois = rois.type_as(feats[0]) - out_size = self.roi_layers[0].output_size + out_size = cast(tuple[int, int], self.roi_layers[0].output_size) + out_h = int(out_size[0]) + out_w = int(out_size[1]) + num_levels = len(feats) - roi_feats = feats[0].new_zeros(rois.size(0), self.out_channels, *out_size) + roi_feats = feats[0].new_zeros((rois.size(0), self.out_channels, out_h, out_w)) # TODO: remove this when parrots supports if torch.__version__ == "parrots": diff --git a/visdet/models/task_modules/assigners/assign_result.py b/visdet/models/task_modules/assigners/assign_result.py index afd6db1e..21e71063 100644 --- a/visdet/models/task_modules/assigners/assign_result.py +++ b/visdet/models/task_modules/assigners/assign_result.py @@ -38,7 +38,7 @@ class AssignResult(util_mixins.NiceRepr): labels.shape=(7,))> """ - def __init__(self, num_gts: int, gt_inds: Tensor, max_overlaps: Tensor, labels: Tensor) -> None: + def __init__(self, num_gts: int, gt_inds: Tensor, max_overlaps: Tensor, labels: Tensor | None = None) -> None: self.num_gts = num_gts self.gt_inds = gt_inds self.max_overlaps = max_overlaps @@ -77,14 +77,8 @@ def __nice__(self): """str: a "nice" summary string describing this assign result""" parts = [] parts.append(f"num_gts={self.num_gts!r}") - if self.gt_inds is None: - parts.append(f"gt_inds={self.gt_inds!r}") - else: - parts.append(f"gt_inds.shape={tuple(self.gt_inds.shape)!r}") - if self.max_overlaps is None: - parts.append(f"max_overlaps={self.max_overlaps!r}") - else: - parts.append(f"max_overlaps.shape={tuple(self.max_overlaps.shape)!r}") + parts.append(f"gt_inds.shape={tuple(self.gt_inds.shape)!r}") + parts.append(f"max_overlaps.shape={tuple(self.max_overlaps.shape)!r}") if self.labels is None: parts.append(f"labels={self.labels!r}") else: @@ -113,7 +107,7 @@ def random(cls, **kwargs): >>> self = AssignResult.random() >>> print(self.info) """ - from ..samplers.sampling_result import ensure_rng + from visdet.core.bbox.demodata import ensure_rng rng = ensure_rng(kwargs.get("rng", None)) @@ -186,9 +180,15 @@ def add_gt_(self, gt_labels): Args: gt_labels (torch.Tensor): Labels of gt boxes """ + old_num_preds = self.gt_inds.numel() + self_inds = torch.arange(1, len(gt_labels) + 1, dtype=torch.long, device=gt_labels.device) self.gt_inds = torch.cat([self_inds, self.gt_inds]) self.max_overlaps = torch.cat([self.max_overlaps.new_ones(len(gt_labels)), self.max_overlaps]) - self.labels = torch.cat([gt_labels, self.labels]) + if self.labels is None: + old_labels = gt_labels.new_zeros((old_num_preds,), dtype=torch.long) + else: + old_labels = self.labels + self.labels = torch.cat([gt_labels, old_labels]) diff --git a/visdet/models/task_modules/assigners/atss_assigner.py b/visdet/models/task_modules/assigners/atss_assigner.py index 6471d8ce..25631da1 100644 --- a/visdet/models/task_modules/assigners/atss_assigner.py +++ b/visdet/models/task_modules/assigners/atss_assigner.py @@ -1,7 +1,10 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any, cast + import torch from torch import Tensor +from visdet.engine.structures import InstanceData from visdet.registry import TASK_UTILS from .assign_result import AssignResult @@ -27,6 +30,33 @@ def __init__( self.ignore_iof_thr = ignore_iof_thr def assign( + self, + pred_instances: InstanceData, + gt_instances: InstanceData, + gt_instances_ignore: InstanceData | None = None, + **kwargs: Any, + ) -> AssignResult: + num_level_bboxes = cast(list[int], kwargs.get("num_level_bboxes")) + if num_level_bboxes is None: + raise ValueError("ATSSAssigner requires `num_level_bboxes` in kwargs") + + bboxes = cast(Tensor, getattr(pred_instances, "priors", getattr(pred_instances, "bboxes"))) + gt_bboxes = cast(Tensor, getattr(gt_instances, "bboxes")) + gt_labels = cast(Tensor | None, getattr(gt_instances, "labels", None)) + + gt_bboxes_ignore: Tensor | None = None + if gt_instances_ignore is not None: + gt_bboxes_ignore = cast(Tensor, getattr(gt_instances_ignore, "bboxes")) + + return self.assign_by_bboxes( + bboxes=bboxes, + num_level_bboxes=num_level_bboxes, + gt_bboxes=gt_bboxes, + gt_bboxes_ignore=gt_bboxes_ignore, + gt_labels=gt_labels, + ) + + def assign_by_bboxes( self, bboxes: Tensor, num_level_bboxes: list[int], diff --git a/visdet/models/task_modules/assigners/dynamic_soft_label_assigner.py b/visdet/models/task_modules/assigners/dynamic_soft_label_assigner.py index 64b120ca..2200d3b4 100644 --- a/visdet/models/task_modules/assigners/dynamic_soft_label_assigner.py +++ b/visdet/models/task_modules/assigners/dynamic_soft_label_assigner.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Optional +from typing import Any, Optional, cast import torch import torch.nn.functional as F @@ -56,14 +56,18 @@ def assign( gt_instances_ignore: Optional[InstanceData] = None, **kwargs, ) -> AssignResult: - gt_bboxes = gt_instances.bboxes - gt_labels = gt_instances.labels - num_gt = gt_bboxes.size(0) + gt_bboxes = cast(Tensor | BaseBoxes, getattr(gt_instances, "bboxes")) + gt_labels = cast(Tensor, getattr(gt_instances, "labels")) + if isinstance(gt_bboxes, BaseBoxes): + num_gt = len(gt_bboxes) + else: + num_gt = gt_bboxes.size(0) + num_gt = int(num_gt) - decoded_bboxes = pred_instances.bboxes - pred_scores = pred_instances.scores - priors = pred_instances.priors - num_bboxes = decoded_bboxes.size(0) + decoded_bboxes = cast(Tensor, getattr(pred_instances, "bboxes")) + pred_scores = cast(Tensor, getattr(pred_instances, "scores")) + priors = cast(Tensor, getattr(pred_instances, "priors")) + num_bboxes = int(decoded_bboxes.size(0)) assigned_gt_inds = decoded_bboxes.new_full((num_bboxes,), 0, dtype=torch.long) @@ -85,7 +89,7 @@ def assign( valid_decoded_bbox = decoded_bboxes[valid_mask] valid_pred_scores = pred_scores[valid_mask] - num_valid = valid_decoded_bbox.size(0) + num_valid = int(valid_decoded_bbox.size(0)) if num_valid == 0: max_overlaps = decoded_bboxes.new_zeros((num_bboxes,)) @@ -144,7 +148,7 @@ def dynamic_k_matching( topk_ious, _ = torch.topk(pairwise_ious, candidate_topk, dim=0) dynamic_ks = torch.clamp(topk_ious.sum(0).int(), min=1) for gt_idx in range(num_gt): - _, pos_idx = torch.topk(cost[:, gt_idx], k=dynamic_ks[gt_idx], largest=False) + _, pos_idx = torch.topk(cost[:, gt_idx], k=int(dynamic_ks[gt_idx].item()), largest=False) matching_matrix[:, gt_idx][pos_idx] = 1 prior_match_gt_mask = matching_matrix.sum(1) > 1 diff --git a/visdet/models/task_modules/assigners/max_iou_assigner.py b/visdet/models/task_modules/assigners/max_iou_assigner.py index 6c533252..700f2f29 100644 --- a/visdet/models/task_modules/assigners/max_iou_assigner.py +++ b/visdet/models/task_modules/assigners/max_iou_assigner.py @@ -1,6 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import copy import logging +from typing import cast import torch from torch import Tensor @@ -25,7 +26,7 @@ def _perm_box(bboxes, iou_calculator, iou_thr=0.97, perm_range=0.01, counter=0, Returns: Tensor: The permuted bboxes. """ - ori_bboxes = copy.deepcopy(bboxes) + ori_bboxes = bboxes.clone() is_valid = True N = bboxes.size(0) perm_factor = bboxes.new_empty(N, 4).uniform_(1 - perm_range, 1 + perm_range) @@ -61,7 +62,7 @@ def perm_repeat_bboxes(bboxes, iou_calculator=None, perm_repeat_cfg=None): import torchvision iou_calculator = torchvision.ops.box_iou - bboxes = copy.deepcopy(bboxes) + bboxes = bboxes.clone() unique_bboxes = bboxes.unique(dim=0) iou_thr = perm_repeat_cfg.get("iou_thr", 0.97) perm_range = perm_repeat_cfg.get("perm_range", 0.01) @@ -194,11 +195,11 @@ def assign( >>> expected_gt_inds = torch.LongTensor([1, 0]) >>> assert torch.all(assign_result.gt_inds == expected_gt_inds) """ - gt_bboxes = gt_instances.bboxes - priors = pred_instances.priors - gt_labels = gt_instances.labels + gt_bboxes = cast(Tensor, getattr(gt_instances, "bboxes")) + priors = cast(Tensor, getattr(pred_instances, "priors")) + gt_labels = cast(Tensor, getattr(gt_instances, "labels")) if gt_instances_ignore is not None: - gt_bboxes_ignore = gt_instances_ignore.bboxes + gt_bboxes_ignore = cast(Tensor, getattr(gt_instances_ignore, "bboxes")) else: gt_bboxes_ignore = None diff --git a/visdet/models/task_modules/coders/__init__.py b/visdet/models/task_modules/coders/__init__.py index 2719a5c2..19baca0c 100644 --- a/visdet/models/task_modules/coders/__init__.py +++ b/visdet/models/task_modules/coders/__init__.py @@ -10,10 +10,10 @@ class DeltaXYWHBBoxCoder: def __init__( self, - target_means=(0.0, 0.0, 0.0, 0.0), - target_stds=(1.0, 1.0, 1.0, 1.0), - clip_border=True, - ): + target_means: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0), + target_stds: tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0), + clip_border: bool = True, + ) -> None: self.means = target_means self.stds = target_stds self.clip_border = clip_border @@ -101,7 +101,7 @@ class DistancePointBBoxCoder: border of the image. Defaults to True. """ - def __init__(self, clip_border=True): + def __init__(self, clip_border: bool = True) -> None: self.clip_border = clip_border @property @@ -121,7 +121,7 @@ def decode(self, points, pred_bboxes, max_shape=None): assert points.size(0) == pred_bboxes.size(0) assert points.size(-1) == 2 assert pred_bboxes.size(-1) == 4 - if self.clip_border is False: + if not self.clip_border: max_shape = None return distance2bbox(points, pred_bboxes, max_shape) diff --git a/visdet/models/task_modules/prior_generators/point_generator.py b/visdet/models/task_modules/prior_generators/point_generator.py index 899ca0d9..e46a6a6e 100644 --- a/visdet/models/task_modules/prior_generators/point_generator.py +++ b/visdet/models/task_modules/prior_generators/point_generator.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import cast + import numpy as np import torch from torch import Tensor @@ -57,7 +59,7 @@ class MlvlPointGenerator: """ def __init__(self, strides: list[int] | list[tuple[int, int]], offset: float = 0.5) -> None: - self.strides = [_pair(stride) for stride in strides] + self.strides: list[tuple[int, int]] = [cast(tuple[int, int], _pair(stride)) for stride in strides] self.offset = offset @property @@ -87,7 +89,7 @@ def grid_priors( ) -> list[Tensor]: """Generate grid points of multiple feature levels.""" assert self.num_levels == len(featmap_sizes) - multi_level_priors = [] + multi_level_priors: list[Tensor] = [] for i in range(self.num_levels): priors = self.single_level_grid_priors( featmap_sizes[i], @@ -130,7 +132,7 @@ def valid_flags( ) -> list[Tensor]: """Generate valid flags of points of multiple feature levels.""" assert self.num_levels == len(featmap_sizes) - multi_level_flags = [] + multi_level_flags: list[Tensor] = [] for i in range(self.num_levels): point_stride = self.strides[i] feat_h, feat_w = featmap_sizes[i] diff --git a/visdet/models/task_modules/samplers/__init__.py b/visdet/models/task_modules/samplers/__init__.py index 33d4d317..00656d68 100644 --- a/visdet/models/task_modules/samplers/__init__.py +++ b/visdet/models/task_modules/samplers/__init__.py @@ -1,10 +1,13 @@ # ruff: noqa from abc import ABCMeta, abstractmethod -import torch +from typing import Any, cast + import numpy as np +import torch + +from visdet.engine.structures import InstanceData from visdet.registry import TASK_UTILS from visdet.utils import util_mixins -from visdet.engine.structures import InstanceData class SamplingResult(util_mixins.NiceRepr): @@ -66,7 +69,14 @@ def __nice__(self): class BaseSampler(metaclass=ABCMeta): """Base class of samplers.""" - def __init__(self, num, pos_fraction, neg_pos_ub=-1, add_gt_as_proposals=True, **kwargs): + def __init__( + self, + num: int, + pos_fraction: float, + neg_pos_ub: int = -1, + add_gt_as_proposals: bool = True, + **kwargs: Any, + ) -> None: self.num = num self.pos_fraction = pos_fraction self.neg_pos_ub = neg_pos_ub @@ -76,19 +86,21 @@ def __init__(self, num, pos_fraction, neg_pos_ub=-1, add_gt_as_proposals=True, * self.context = kwargs.pop("context", None) @abstractmethod - def _sample_pos(self, assign_result, num_expected, **kwargs): + def _sample_pos(self, assign_result: Any, num_expected: int, **kwargs: Any) -> torch.Tensor: """Sample positive samples.""" - pass + raise NotImplementedError @abstractmethod - def _sample_neg(self, assign_result, num_expected, **kwargs): + def _sample_neg(self, assign_result: Any, num_expected: int, **kwargs: Any) -> torch.Tensor: """Sample negative samples.""" - pass + raise NotImplementedError - def sample(self, assign_result, pred_instances: InstanceData, gt_instances: InstanceData, **kwargs): + def sample( + self, assign_result: Any, pred_instances: InstanceData, gt_instances: InstanceData, **kwargs: Any + ) -> SamplingResult: """Sample positive and negative bboxes.""" - bboxes = pred_instances.priors - gt_bboxes = gt_instances.bboxes + bboxes = cast(torch.Tensor, getattr(pred_instances, "priors")) + gt_bboxes = cast(torch.Tensor, getattr(gt_instances, "bboxes")) gt_labels = getattr(gt_instances, "labels", None) if len(bboxes.shape) < 2: @@ -135,10 +147,12 @@ def _sample_pos(self, assign_result, num_expected, **kwargs): def _sample_neg(self, assign_result, num_expected, **kwargs): return torch.nonzero(assign_result.gt_inds == 0, as_tuple=False).squeeze(-1) - def sample(self, assign_result, pred_instances: InstanceData, gt_instances: InstanceData, **kwargs): + def sample( + self, assign_result: Any, pred_instances: InstanceData, gt_instances: InstanceData, **kwargs: Any + ) -> SamplingResult: """Directly returns the positive and negative indices.""" - priors = pred_instances.priors - gt_bboxes = gt_instances.bboxes + priors = cast(torch.Tensor, getattr(pred_instances, "priors")) + gt_bboxes = cast(torch.Tensor, getattr(gt_instances, "bboxes")) pos_inds = self._sample_pos(assign_result, 0) neg_inds = self._sample_neg(assign_result, 0) @@ -163,7 +177,7 @@ class RandomSampler(BaseSampler): def __init__(self, num, pos_fraction, neg_pos_ub=-1, add_gt_as_proposals=True, **kwargs): super().__init__(num, pos_fraction, neg_pos_ub, add_gt_as_proposals, **kwargs) - def random_choice(self, gallery, num): + def random_choice(self, gallery: Any, num: int) -> torch.Tensor | np.ndarray: """Random select some elements from the gallery.""" assert len(gallery) >= num @@ -203,7 +217,15 @@ def _sample_neg(self, assign_result, num_expected, **kwargs): class IoUBalancedNegSampler(RandomSampler): """IoU Balanced Sampling.""" - def __init__(self, num, pos_fraction, floor_thr=-1, floor_fraction=0, num_bins=3, **kwargs): + def __init__( + self, + num: int, + pos_fraction: float, + floor_thr: int = -1, + floor_fraction: float = 0.0, + num_bins: int = 3, + **kwargs: Any, + ) -> None: super().__init__(num, pos_fraction, **kwargs) self.floor_thr = floor_thr self.floor_fraction = floor_fraction @@ -305,14 +327,17 @@ def _sample_pos(self, assign_result, num_expected, **kwargs): unique_gt_inds = gt_inds.unique() num_gts = unique_gt_inds.numel() - pos_inds_from_instance = [] + pos_inds_from_instance: list[torch.Tensor] = [] if num_expected >= num_gts: # at least one sample per instance avg_num_per_instance = int(num_expected / num_gts) for i in range(num_gts): instance_pos_inds = pos_inds[gt_inds == unique_gt_inds[i]] if instance_pos_inds.numel() > avg_num_per_instance: - sampled_instance_pos_inds = self.random_choice(instance_pos_inds, avg_num_per_instance) + sampled_instance_pos_inds = cast( + torch.Tensor, + self.random_choice(instance_pos_inds, avg_num_per_instance), + ) else: sampled_instance_pos_inds = instance_pos_inds pos_inds_from_instance.append(sampled_instance_pos_inds) @@ -326,13 +351,23 @@ def _sample_pos(self, assign_result, num_expected, **kwargs): extra_inds = self.random_choice(gallery_inds, remaining_num) else: extra_inds = np.array(gallery_inds, dtype=int) - pos_inds_from_instance.append(torch.from_numpy(extra_inds).to(pos_inds.device)) + + if isinstance(extra_inds, torch.Tensor): + extra_inds_tensor = extra_inds.to(pos_inds.device) + else: + extra_inds_tensor = torch.from_numpy(extra_inds).to(pos_inds.device) + pos_inds_from_instance.append(extra_inds_tensor) else: # sample instances and then sample one bbox per instance - sampled_instance_indices = self.random_choice(range(num_gts), num_expected) - for i in sampled_instance_indices: + sampled_instance_indices = self.random_choice(list(range(num_gts)), num_expected) + if isinstance(sampled_instance_indices, torch.Tensor): + sampled_instance_indices_list = [int(i) for i in sampled_instance_indices.cpu().tolist()] + else: + sampled_instance_indices_list = [int(i) for i in sampled_instance_indices.tolist()] + + for i in sampled_instance_indices_list: instance_pos_inds = pos_inds[gt_inds == unique_gt_inds[i]] - pos_inds_from_instance.append(self.random_choice(instance_pos_inds, 1)) + pos_inds_from_instance.append(cast(torch.Tensor, self.random_choice(instance_pos_inds, 1))) return torch.cat(pos_inds_from_instance) diff --git a/visdet/models/utils/image.py b/visdet/models/utils/image.py index c09d197b..109b0ef2 100644 --- a/visdet/models/utils/image.py +++ b/visdet/models/utils/image.py @@ -1,3 +1,5 @@ +from typing import cast + import numpy as np import torch from torch import Tensor @@ -20,7 +22,7 @@ def imrenormalize(img: Tensor | np.ndarray, img_norm_cfg: dict, new_img_norm_cfg if isinstance(img, torch.Tensor): assert img.ndim == 4 and img.shape[0] == 1 new_img = img.squeeze(0).cpu().numpy().transpose(1, 2, 0) - new_img = _imrenormalize(new_img, img_norm_cfg, new_img_norm_cfg) + new_img = cast(np.ndarray, _imrenormalize(new_img, img_norm_cfg, new_img_norm_cfg)) new_img = new_img.transpose(2, 0, 1)[None] return torch.from_numpy(new_img).to(img) else: @@ -31,16 +33,18 @@ def _imrenormalize(img: Tensor | np.ndarray, img_norm_cfg: dict, new_img_norm_cf """Re-normalize the image.""" img_norm_cfg = img_norm_cfg.copy() new_img_norm_cfg = new_img_norm_cfg.copy() - for k, v in img_norm_cfg.items(): - if (k == "mean" or k == "std") and not isinstance(v, np.ndarray): - img_norm_cfg[k] = np.array(v, dtype=img.dtype) + if isinstance(img, np.ndarray): + for k, v in img_norm_cfg.items(): + if (k == "mean" or k == "std") and not isinstance(v, np.ndarray): + img_norm_cfg[k] = np.array(v, dtype=img.dtype) # reverse cfg if "bgr_to_rgb" in img_norm_cfg: img_norm_cfg["rgb_to_bgr"] = img_norm_cfg["bgr_to_rgb"] img_norm_cfg.pop("bgr_to_rgb") - for k, v in new_img_norm_cfg.items(): - if (k == "mean" or k == "std") and not isinstance(v, np.ndarray): - new_img_norm_cfg[k] = np.array(v, dtype=img.dtype) + if isinstance(img, np.ndarray): + for k, v in new_img_norm_cfg.items(): + if (k == "mean" or k == "std") and not isinstance(v, np.ndarray): + new_img_norm_cfg[k] = np.array(v, dtype=img.dtype) # Denormalize mean = img_norm_cfg.get("mean", [0, 0, 0]) std = img_norm_cfg.get("std", [1, 1, 1]) diff --git a/visdet/models/utils/res_layer.py b/visdet/models/utils/res_layer.py index 5b9c192f..7c6a9388 100644 --- a/visdet/models/utils/res_layer.py +++ b/visdet/models/utils/res_layer.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +from typing import Any + from torch import nn as nn from visdet.cv.cnn import build_conv_layer, build_norm_layer @@ -26,26 +28,26 @@ class ResLayer(Sequential): def __init__( self, - block, - inplanes, - planes, - num_blocks, - stride=1, - avg_down=False, - conv_cfg=None, - norm_cfg=dict(type="BN"), - downsample_first=True, - **kwargs, - ): + block: Any, + inplanes: int, + planes: int, + num_blocks: int, + stride: int = 1, + avg_down: bool = False, + conv_cfg: dict | None = None, + norm_cfg: dict = dict(type="BN"), + downsample_first: bool = True, + **kwargs: Any, + ) -> None: self.block = block - downsample = None + downsample: nn.Module | None = None if stride != 1 or inplanes != planes * block.expansion: - downsample = [] + downsample_modules: list[nn.Module] = [] conv_stride = stride if avg_down: conv_stride = 1 - downsample.append( + downsample_modules.append( nn.AvgPool2d( kernel_size=stride, stride=stride, @@ -53,7 +55,7 @@ def __init__( count_include_pad=False, ) ) - downsample.extend( + downsample_modules.extend( [ build_conv_layer( conv_cfg, @@ -66,9 +68,9 @@ def __init__( build_norm_layer(norm_cfg, planes * block.expansion)[1], ] ) - downsample = nn.Sequential(*downsample) + downsample = nn.Sequential(*downsample_modules) - layers = [] + layers: list[nn.Module] = [] if downsample_first: layers.append( block( @@ -192,12 +194,16 @@ def forward(self, x): out = self.conv1(x) if self.with_norm: - out = self.norm1(out) + norm1 = self.norm1 + assert norm1 is not None + out = norm1(out) out = self.relu(out) out = self.conv2(out) if self.with_norm: - out = self.norm2(out) + norm2 = self.norm2 + assert norm2 is not None + out = norm2(out) if self.downsample is not None: identity = self.downsample(x)