├── nanodet ├── data │ ├── transform │ │ ├── mosaic.py │ │ ├── __init__.py │ │ ├── pipeline.py │ │ └── color.py │ ├── batch_process.py │ ├── dataset │ │ └── __init__.py │ └── collate.py ├── optim │ ├── __init__.py │ └── builder.py ├── __init__.py ├── model │ ├── head │ │ ├── assigner │ │ │ └── base_assigner.py │ │ └── __init__.py │ ├── module │ │ ├── scale.py │ │ ├── activation.py │ │ ├── init_weights.py │ │ └── norm.py │ ├── weight_averager │ │ ├── __init__.py │ │ └── ema.py │ ├── fpn │ │ ├── __init__.py │ │ ├── pan.py │ │ └── fpn.py │ ├── arch │ │ ├── __init__.py │ │ ├── nanodet_plus.py │ │ └── one_stage_detector.py │ ├── backbone │ │ ├── __init__.py │ │ └── timm_wrapper.py │ └── loss │ │ └── utils.py ├── __about__.py ├── trainer │ └── __init__.py ├── util │ ├── rank_filter.py │ ├── config.py │ ├── path.py │ ├── __init__.py │ ├── box_transform.py │ ├── misc.py │ └── env_utils.py └── evaluator │ └── __init__.py ├── demo_android_ncnn ├── app │ ├── .gitignore │ ├── src │ │ ├── main │ │ │ ├── res │ │ │ │ ├── values │ │ │ │ │ ├── strings.xml │ │ │ │ │ ├── colors.xml │ │ │ │ │ └── styles.xml │ │ │ │ ├── drawable-xxhdpi │ │ │ │ │ ├── cpu.png │ │ │ │ │ ├── gpu.png │ │ │ │ │ └── ncnn_icon.png │ │ │ │ ├── mipmap-hdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-mdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-xhdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-xxhdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-xxxhdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── drawable │ │ │ │ │ └── cpu_gpu_bg.xml │ │ │ │ ├── mipmap-anydpi-v26 │ │ │ │ │ ├── ic_launcher.xml │ │ │ │ │ └── ic_launcher_round.xml │ │ │ │ ├── drawable-v24 │ │ │ │ │ └── ic_launcher_foreground.xml │ │ │ │ └── layout │ │ │ │ │ └── activity_welcome.xml │ │ │ ├── cpp │ │ │ │ ├── CMakeLists.txt │ │ │ │ ├── YoloV4.h │ │ │ │ ├── YoloV4.cpp │ │ │ │ ├── YoloV5.h │ │ │ │ └── NanoDet.h │ │ │ ├── java │ │ │ │ └── com │ │ │ │ │ └── rangi │ │ │ │ │ └── nanodet │ │ │ │ │ ├── YOLOv5.java │ │ │ │ │ ├── NanoDet.java │ │ │ │ │ ├── YOLOv4.java │ │ │ │ │ ├── AppCrashHandler.java │ │ │ │ │ ├── Box.java │ │ │ │ │ ├── NcnnApp.java │ │ │ │ │ └── WelcomeActivity.java │ │ │ └── AndroidManifest.xml │ │ ├── test │ │ │ └── java │ │ │ │ └── com │ │ │ │ └── rangi │ │ │ │ └── nanodet │ │ │ │ └── ExampleUnitTest.java │ │ └── androidTest │ │ │ └── java │ │ │ └── com │ │ │ └── rangi │ │ │ └── nanodet │ │ │ └── ExampleInstrumentedTest.java │ ├── proguard-rules.pro │ └── build.gradle ├── settings.gradle ├── Android_demo.jpg ├── gradle │ └── wrapper │ │ ├── gradle-wrapper.jar │ │ └── gradle-wrapper.properties ├── .gitignore ├── build.gradle ├── gradle.properties ├── README.md └── gradlew.bat ├── .isort.cfg ├── demo └── demo_multi_backend_python.py ├── docs ├── imgs │ ├── Title.jpg │ ├── Model_arch.png │ ├── Android_demo.jpg │ └── nanodet-plus-arch.png └── update.md ├── demo_ncnn ├── benchmark.jpg ├── CMakeLists.txt └── nanodet.h ├── tests ├── data │ ├── test_img.jpg │ ├── batched_nms_data.pkl │ ├── test_img.txt │ ├── test_img.xml │ └── dummy_coco.json ├── test_utils │ ├── test_flops.py │ ├── test_logger.py │ └── test_env_utils.py ├── test_data │ ├── test_batch_process.py │ ├── test_dataset │ │ ├── test_yolodataset.py │ │ ├── test_xmldataset.py │ │ └── test_cocodataset.py │ ├── test_transform │ │ ├── test_color.py │ │ └── test_warp.py │ └── test_collate.py ├── test_models │ ├── test_modules │ │ ├── test_scale.py │ │ ├── test_dwconv.py │ │ ├── test_transformer.py │ │ ├── test_nms.py │ │ ├── test_repvgg_conv.py │ │ ├── test_norm.py │ │ ├── test_init_weights.py │ │ └── test_conv.py │ ├── test_head │ │ ├── test_simple_conv_head.py │ │ ├── test_nanodet_head.py │ │ └── test_gfl_head.py │ ├── test_loss │ │ ├── test_iou_loss.py │ │ └── test_gfocal_loss.py │ ├── test_weight_averager │ │ └── test_ema.py │ ├── test_backbone │ │ ├── test_timm_wrapper.py │ │ ├── test_custom_csp.py │ │ ├── test_mobilenetv2.py │ │ ├── test_repvgg.py │ │ ├── test_resnet.py │ │ ├── test_shufflenetv2.py │ │ ├── test_ghostnet.py │ │ └── test_efficient_lite.py │ └── test_fpn │ │ ├── test_tan.py │ │ ├── test_fpn.py │ │ ├── test_pan.py │ │ └── test_ghost_pan.py ├── test_optim │ └── test_builder.py ├── test_evaluator │ └── test_coco_detection.py ├── test_configs │ └── test_config.py └── test_trainer │ └── test_lightning_task.py ├── demo_mnn ├── imgs │ ├── 000252.jpg │ └── 000258.jpg ├── results │ ├── 000252.jpg │ └── 000258.jpg ├── CMakeLists.txt └── README.md ├── requirements.txt ├── .flake8 ├── demo_openvino ├── CMakeLists.txt └── nanodet_openvino.h ├── .pre-commit-config.yaml ├── demo_libtorch ├── CMakeLists.txt ├── README.md └── nanodet_libtorch.h ├── setup.py ├── tools ├── convert_old_checkpoint.py ├── flops.py ├── inference.py └── export_torchscript.py ├── .gitignore └── config ├── nanodet_custom_xml_dataset.yml └── legacy_v0.x_configs └── nanodet-m.yml /nanodet/data/transform/mosaic.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/.gitignore: -------------------------------------------------------------------------------- 1 | /build 2 | -------------------------------------------------------------------------------- /.isort.cfg: -------------------------------------------------------------------------------- 1 | [tool.isort] 2 | profile = "black" 3 | -------------------------------------------------------------------------------- /demo/demo_multi_backend_python.py: -------------------------------------------------------------------------------- 1 | # Working in progress 2 | -------------------------------------------------------------------------------- /demo_android_ncnn/settings.gradle: -------------------------------------------------------------------------------- 1 | include ':app' 2 | rootProject.name='NanoDet' 3 | -------------------------------------------------------------------------------- /docs/imgs/Title.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/docs/imgs/Title.jpg -------------------------------------------------------------------------------- /demo_ncnn/benchmark.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_ncnn/benchmark.jpg -------------------------------------------------------------------------------- /tests/data/test_img.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/tests/data/test_img.jpg -------------------------------------------------------------------------------- /demo_mnn/imgs/000252.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_mnn/imgs/000252.jpg -------------------------------------------------------------------------------- /demo_mnn/imgs/000258.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_mnn/imgs/000258.jpg -------------------------------------------------------------------------------- /docs/imgs/Model_arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/docs/imgs/Model_arch.png -------------------------------------------------------------------------------- /demo_mnn/results/000252.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_mnn/results/000252.jpg -------------------------------------------------------------------------------- /demo_mnn/results/000258.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_mnn/results/000258.jpg -------------------------------------------------------------------------------- /docs/imgs/Android_demo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/docs/imgs/Android_demo.jpg -------------------------------------------------------------------------------- /nanodet/optim/__init__.py: -------------------------------------------------------------------------------- 1 | from .builder import build_optimizer 2 | 3 | __all__ = ["build_optimizer"] 4 | -------------------------------------------------------------------------------- /docs/imgs/nanodet-plus-arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/docs/imgs/nanodet-plus-arch.png -------------------------------------------------------------------------------- /tests/data/batched_nms_data.pkl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/tests/data/batched_nms_data.pkl -------------------------------------------------------------------------------- /demo_android_ncnn/Android_demo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/Android_demo.jpg -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/values/strings.xml: -------------------------------------------------------------------------------- 1 | 2 | NanoDet 3 | 4 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradle/wrapper/gradle-wrapper.jar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/gradle/wrapper/gradle-wrapper.jar -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-xxhdpi/cpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/cpu.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-xxhdpi/gpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/gpu.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-xxhdpi/ncnn_icon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/ncnn_icon.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RangiLyu/nanodet/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /nanodet/__init__.py: -------------------------------------------------------------------------------- 1 | """package info.""" 2 | 3 | import os 4 | 5 | from nanodet.__about__ import * # noqa: F401 F403 6 | 7 | _PACKAGE_ROOT = os.path.dirname(__file__) 8 | _PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT) 9 | -------------------------------------------------------------------------------- /nanodet/model/head/assigner/base_assigner.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | 4 | class BaseAssigner(metaclass=ABCMeta): 5 | @abstractmethod 6 | def assign(self, bboxes, gt_bboxes, gt_bboxes_ignore=None, gt_labels=None): 7 | pass 8 | -------------------------------------------------------------------------------- /tests/data/test_img.txt: -------------------------------------------------------------------------------- 1 | 0 0.608987 0.354681 0.359542 0.404493 2 | 0 0.719387 0.691062 0.037075 0.074150 3 | 0 0.813105 0.692525 0.032876 0.038088 4 | 0 0.865956 0.690507 0.020801 0.060458 5 | 0 0.922998 0.677377 0.035114 0.085539 6 | 0 0.956160 0.656642 0.021013 0.041487 7 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradle/wrapper/gradle-wrapper.properties: -------------------------------------------------------------------------------- 1 | #Mon Jun 15 13:17:33 HKT 2020 2 | distributionBase=GRADLE_USER_HOME 3 | distributionPath=wrapper/dists 4 | zipStoreBase=GRADLE_USER_HOME 5 | zipStorePath=wrapper/dists 6 | distributionUrl=https\://services.gradle.org/distributions/gradle-6.1.1-all.zip 7 | -------------------------------------------------------------------------------- /demo_android_ncnn/.gitignore: -------------------------------------------------------------------------------- 1 | *.iml 2 | .gradle 3 | /local.properties 4 | /.idea 5 | /.idea/caches 6 | /.idea/libraries 7 | /.idea/modules.xml 8 | /.idea/workspace.xml 9 | /.idea/navEditor.xml 10 | /.idea/assetWizardSettings.xml 11 | .DS_Store 12 | /build 13 | /captures 14 | .externalNativeBuild 15 | .cxx 16 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | Cython 2 | imagesize 3 | matplotlib 4 | numpy 5 | omegaconf>=2.0.1 6 | onnx 7 | onnx-simplifier 8 | opencv-python 9 | pyaml 10 | pycocotools 11 | pytorch-lightning>=1.9.0,<2.0.0 12 | tabulate 13 | tensorboard 14 | termcolor 15 | torch>=1.10,<2.0 16 | torchmetrics 17 | torchvision 18 | tqdm 19 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/values/colors.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | #123F70 4 | #122F49 5 | #D81B60 6 | 7 | #DDDDDD 8 | 9 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable/cpu_gpu_bg.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-anydpi-v26/ic_launcher.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-anydpi-v26/ic_launcher_round.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /tests/test_utils/test_flops.py: -------------------------------------------------------------------------------- 1 | from nanodet.model.arch import build_model 2 | from nanodet.util import cfg, get_model_complexity_info, load_config 3 | 4 | 5 | def test_flops(): 6 | load_config(cfg, "./config/legacy_v0.x_configs/nanodet-m.yml") 7 | 8 | model = build_model(cfg.model) 9 | input_shape = (3, 320, 320) 10 | get_model_complexity_info(model, input_shape) 11 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.10) 2 | 3 | set(ncnn_DIR ${CMAKE_SOURCE_DIR}/ncnn-20211208-android-vulkan/${ANDROID_ABI}/lib/cmake/ncnn) 4 | find_package(ncnn REQUIRED) 5 | 6 | add_library(yolov5 SHARED 7 | jni_interface.cpp 8 | YoloV5.cpp 9 | YoloV4.cpp 10 | NanoDet.cpp 11 | ) 12 | 13 | target_link_libraries(yolov5 ncnn jnigraphics) 14 | -------------------------------------------------------------------------------- /nanodet/model/module/scale.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class Scale(nn.Module): 6 | """ 7 | A learnable scale parameter 8 | """ 9 | 10 | def __init__(self, scale=1.0): 11 | super(Scale, self).__init__() 12 | self.scale = nn.Parameter(torch.tensor(scale, dtype=torch.float)) 13 | 14 | def forward(self, x): 15 | return x * self.scale 16 | -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- 1 | # This is an example .flake8 config, used when developing *Black* itself. 2 | # Keep in sync with setup.cfg which is used for source packages. 3 | 4 | [flake8] 5 | ignore = W503, E203, E221, C901, C408, E741, C407, E741, B006, B007, B017, B950, C416, E203 6 | max-line-length = 88 7 | max-complexity = 18 8 | select = B,C,E,F,W,T4,B9 9 | exclude = build 10 | per-file-ignores = 11 | **/__init__.py:F401,F403,E402 12 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/YOLOv5.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.res.AssetManager; 4 | import android.graphics.Bitmap; 5 | 6 | public class YOLOv5 { 7 | static { 8 | System.loadLibrary("yolov5"); 9 | } 10 | 11 | public static native void init(AssetManager manager, boolean useGPU); 12 | public static native Box[] detect(Bitmap bitmap, double threshold, double nms_threshold); 13 | } 14 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/NanoDet.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.res.AssetManager; 4 | import android.graphics.Bitmap; 5 | 6 | public class NanoDet { 7 | static { 8 | System.loadLibrary("yolov5"); 9 | } 10 | 11 | public static native void init(AssetManager manager, boolean useGPU); 12 | public static native Box[] detect(Bitmap bitmap, double threshold, double nms_threshold); 13 | } 14 | -------------------------------------------------------------------------------- /demo_mnn/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.9) 2 | project(nanodet-mnn) 3 | 4 | set(CMAKE_CXX_STANDARD 17) 5 | 6 | # find_package(OpenCV REQUIRED PATHS "/work/dependence/opencv/opencv-3.4.3/build") 7 | find_package(OpenCV REQUIRED) 8 | include_directories( 9 | mnn/include 10 | . 11 | ) 12 | 13 | link_directories(mnn/lib) 14 | 15 | add_executable(nanodet-mnn main.cpp nanodet_mnn.cpp) 16 | target_link_libraries(nanodet-mnn MNN ${OpenCV_LIBS}) 17 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/YOLOv4.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.res.AssetManager; 4 | import android.graphics.Bitmap; 5 | 6 | public class YOLOv4 { 7 | static { 8 | System.loadLibrary("yolov5"); // 存放在yolov5.so中 9 | } 10 | 11 | public static native void init(AssetManager manager, boolean useGPU); 12 | public static native Box[] detect(Bitmap bitmap, double threshold, double nms_threshold); 13 | } 14 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/test/java/com/rangi/nanodet/ExampleUnitTest.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import org.junit.Test; 4 | 5 | import static org.junit.Assert.*; 6 | 7 | /** 8 | * Example local unit test, which will execute on the development machine (host). 9 | * 10 | * @see Testing documentation 11 | */ 12 | public class ExampleUnitTest { 13 | @Test 14 | public void addition_isCorrect() { 15 | assertEquals(4, 2 + 2); 16 | } 17 | } 18 | -------------------------------------------------------------------------------- /demo_openvino/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.4.1) 2 | set(CMAKE_CXX_STANDARD 14) 3 | 4 | project(nanodet_demo) 5 | 6 | find_package(OpenCV REQUIRED) 7 | find_package(InferenceEngine REQUIRED) 8 | find_package(ngraph REQUIRED) 9 | 10 | include_directories( 11 | ${OpenCV_INCLUDE_DIRS} 12 | ${CMAKE_CURRENT_SOURCE_DIR} 13 | ${CMAKE_CURRENT_BINARY_DIR} 14 | ) 15 | 16 | add_executable(nanodet_demo main.cpp nanodet_openvino.cpp) 17 | 18 | target_link_libraries( 19 | nanodet_demo 20 | ${InferenceEngine_LIBRARIES} 21 | ${NGRAPH_LIBRARIES} 22 | ${OpenCV_LIBS} 23 | ) 24 | -------------------------------------------------------------------------------- /tests/test_data/test_batch_process.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.data.batch_process import stack_batch_img 5 | 6 | 7 | def test_stack_batch_img(): 8 | with pytest.raises(AssertionError): 9 | dummy_imgs = [torch.rand((1, 30, 30)), torch.rand((3, 28, 10))] 10 | stack_batch_img(dummy_imgs, divisible=32, pad_value=0) 11 | 12 | dummy_imgs = [ 13 | torch.rand((3, 300, 300)), 14 | torch.rand((3, 288, 180)), 15 | torch.rand((3, 169, 256)), 16 | ] 17 | batch_tensor = stack_batch_img(dummy_imgs, divisible=32, pad_value=0) 18 | assert batch_tensor.shape == (3, 3, 320, 320) 19 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_scale.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nanodet.model.module.scale import Scale 4 | 5 | 6 | def test_scale(): 7 | # test default scale 8 | scale = Scale() 9 | assert scale.scale.data == 1.0 10 | assert scale.scale.dtype == torch.float 11 | x = torch.rand(1, 3, 64, 64) 12 | output = scale(x) 13 | assert output.shape == (1, 3, 64, 64) 14 | 15 | # test given scale 16 | scale = Scale(10.0) 17 | assert scale.scale.data == 10.0 18 | assert scale.scale.dtype == torch.float 19 | x = torch.rand(1, 3, 64, 64) 20 | output = scale(x) 21 | assert output.shape == (1, 3, 64, 64) 22 | -------------------------------------------------------------------------------- /tests/test_data/test_dataset/test_yolodataset.py: -------------------------------------------------------------------------------- 1 | from nanodet.data.dataset import YoloDataset 2 | 3 | 4 | def test_yolodataset(): 5 | ann_path = "tests/data" 6 | yolodataset = YoloDataset( 7 | img_path=ann_path, 8 | ann_path=ann_path, 9 | class_names=["class1"], 10 | input_size=[320, 320], # [w,h] 11 | keep_ratio=False, 12 | pipeline=dict(normalize=[[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]), 13 | ) 14 | assert len(yolodataset) == 1 15 | for i, data in enumerate(yolodataset): 16 | assert data["img"].shape == (3, 320, 320) 17 | assert data["gt_bboxes"].shape == (6, 4) 18 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/AppCrashHandler.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import androidx.annotation.NonNull; 4 | 5 | class AppCrashHandler implements Thread.UncaughtExceptionHandler { 6 | 7 | private Thread.UncaughtExceptionHandler uncaughtExceptionHandler = Thread.getDefaultUncaughtExceptionHandler(); 8 | 9 | @Override 10 | public void uncaughtException(@NonNull Thread t, @NonNull Throwable e) { 11 | uncaughtExceptionHandler.uncaughtException(t, e); 12 | } 13 | 14 | public static void register() { 15 | Thread.setDefaultUncaughtExceptionHandler(new AppCrashHandler()); 16 | } 17 | 18 | } 19 | -------------------------------------------------------------------------------- /tests/test_utils/test_logger.py: -------------------------------------------------------------------------------- 1 | import tempfile 2 | 3 | from torch.utils.tensorboard import SummaryWriter 4 | 5 | from nanodet.util import NanoDetLightningLogger, cfg, load_config 6 | 7 | 8 | def test_logger(): 9 | tmp_dir = tempfile.TemporaryDirectory() 10 | logger = NanoDetLightningLogger(tmp_dir.name) 11 | 12 | writer = logger.experiment 13 | assert isinstance(writer, SummaryWriter) 14 | 15 | logger.info("test") 16 | 17 | logger.log_hyperparams({"lr": 1}) 18 | 19 | logger.log_metrics({"mAP": 30.1}, 1) 20 | 21 | load_config(cfg, "./config/legacy_v0.x_configs/nanodet-m.yml") 22 | logger.dump_cfg(cfg) 23 | 24 | logger.finalize(None) 25 | -------------------------------------------------------------------------------- /nanodet/__about__.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | _this_year = time.strftime("%Y") 4 | __version__ = "1.0.0" 5 | __author__ = "RangiLyu" 6 | __author_email__ = "lyuchqi@gmail.com" 7 | __license__ = "Apache-2.0" 8 | __copyright__ = f"Copyright (c) 2020-{_this_year}, {__author__}." 9 | __homepage__ = "https://github.com/RangiLyu/nanodet" 10 | 11 | __docs__ = ( 12 | "NanoDet: Deep learning object detection toolbox for super fast and " 13 | "lightweight anchor-free object detection models." 14 | ) 15 | 16 | __all__ = [ 17 | "__author__", 18 | "__author_email__", 19 | "__copyright__", 20 | "__docs__", 21 | "__homepage__", 22 | "__license__", 23 | "__version__", 24 | ] 25 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v2.5.0 4 | hooks: 5 | - id: trailing-whitespace 6 | - id: end-of-file-fixer 7 | - id: check-docstring-first 8 | - id: check-yaml 9 | - id: debug-statements 10 | - id: requirements-txt-fixer 11 | 12 | - repo: https://github.com/pycqa/isort 13 | rev: 5.10.1 14 | hooks: 15 | - id: isort 16 | args: ["--profile", "black"] 17 | 18 | - repo: https://github.com/psf/black 19 | rev: 22.6.0 20 | hooks: 21 | - id: black 22 | 23 | - repo: https://github.com/pycqa/flake8 24 | rev: 5.0.4 25 | hooks: 26 | - id: flake8 27 | -------------------------------------------------------------------------------- /nanodet/model/head/__init__.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | from .gfl_head import GFLHead 4 | from .nanodet_head import NanoDetHead 5 | from .nanodet_plus_head import NanoDetPlusHead 6 | from .simple_conv_head import SimpleConvHead 7 | 8 | 9 | def build_head(cfg): 10 | head_cfg = copy.deepcopy(cfg) 11 | name = head_cfg.pop("name") 12 | if name == "GFLHead": 13 | return GFLHead(**head_cfg) 14 | elif name == "NanoDetHead": 15 | return NanoDetHead(**head_cfg) 16 | elif name == "NanoDetPlusHead": 17 | return NanoDetPlusHead(**head_cfg) 18 | elif name == "SimpleConvHead": 19 | return SimpleConvHead(**head_cfg) 20 | else: 21 | raise NotImplementedError 22 | -------------------------------------------------------------------------------- /demo_android_ncnn/build.gradle: -------------------------------------------------------------------------------- 1 | // Top-level build file where you can add configuration options common to all sub-projects/modules. 2 | 3 | buildscript { 4 | repositories { 5 | google() 6 | jcenter() 7 | 8 | } 9 | dependencies { 10 | classpath 'com.android.tools.build:gradle:4.0.0' 11 | 12 | // NOTE: Do not place your application dependencies here; they belong 13 | // in the individual module build.gradle files 14 | } 15 | } 16 | 17 | allprojects { 18 | repositories { 19 | google() 20 | jcenter() 21 | maven { url "https://jitpack.io" } 22 | } 23 | } 24 | 25 | task clean(type: Delete) { 26 | delete rootProject.buildDir 27 | } 28 | -------------------------------------------------------------------------------- /nanodet/trainer/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from .task import TrainingTask 15 | 16 | __all__ = ["TrainingTask"] 17 | -------------------------------------------------------------------------------- /nanodet/data/transform/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | from .pipeline import Pipeline 16 | 17 | __all__ = ["Pipeline"] 18 | -------------------------------------------------------------------------------- /tests/test_models/test_head/test_simple_conv_head.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nanodet.model.head import build_head 4 | from nanodet.util.yacs import CfgNode 5 | 6 | 7 | def test_simple_conv_head_forward(): 8 | head_cfg = dict( 9 | name="SimpleConvHead", 10 | num_classes=80, 11 | input_channel=1, 12 | feat_channels=96, 13 | stacked_convs=2, 14 | conv_type="DWConv", 15 | reg_max=8, 16 | strides=[8, 16, 32], 17 | ) 18 | cfg = CfgNode(head_cfg) 19 | head = build_head(cfg) 20 | feat = [torch.rand(1, 1, 320 // stride, 320 // stride) for stride in [8, 16, 32]] 21 | out = head.forward(feat) 22 | num_points = sum([(320 // stride) ** 2 for stride in [8, 16, 32]]) 23 | assert out.shape == (1, num_points, 80 + (8 + 1) * 4) 24 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/values/styles.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 10 | 16 | 17 | 18 | -------------------------------------------------------------------------------- /tests/test_models/test_loss/test_iou_loss.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.loss.iou_loss import ( 5 | BoundedIoULoss, 6 | CIoULoss, 7 | DIoULoss, 8 | GIoULoss, 9 | IoULoss, 10 | ) 11 | 12 | 13 | @pytest.mark.parametrize( 14 | "loss_class", [IoULoss, BoundedIoULoss, GIoULoss, DIoULoss, CIoULoss] 15 | ) 16 | def test_iou_type_loss_zeros_weight(loss_class): 17 | pred = torch.rand((10, 4)) 18 | target = torch.rand((10, 4)) 19 | weight = torch.zeros(10) 20 | 21 | with pytest.raises(AssertionError): 22 | loss_class()(pred, target, reduction_override="2333") 23 | 24 | loss = loss_class()(pred, target, weight) 25 | assert loss == 0.0 26 | 27 | weight = torch.rand(10) 28 | loss = loss_class()(pred, target, weight) 29 | assert loss != 0.0 30 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/proguard-rules.pro: -------------------------------------------------------------------------------- 1 | # Add project specific ProGuard rules here. 2 | # You can control the set of applied configuration files using the 3 | # proguardFiles setting in build.gradle. 4 | # 5 | # For more details, see 6 | # http://developer.android.com/guide/developing/tools/proguard.html 7 | 8 | # If your project uses WebView with JS, uncomment the following 9 | # and specify the fully qualified class name to the JavaScript interface 10 | # class: 11 | #-keepclassmembers class fqcn.of.javascript.interface.for.webview { 12 | # public *; 13 | #} 14 | 15 | # Uncomment this to preserve the line number information for 16 | # debugging stack traces. 17 | #-keepattributes SourceFile,LineNumberTable 18 | 19 | # If you keep the line number information, uncomment this to 20 | # hide the original source file name. 21 | #-renamesourcefileattribute SourceFile 22 | -------------------------------------------------------------------------------- /nanodet/util/rank_filter.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | 16 | def rank_filter(func): 17 | def func_filter(local_rank=-1, *args, **kwargs): 18 | if local_rank < 1: 19 | return func(*args, **kwargs) 20 | else: 21 | pass 22 | 23 | return func_filter 24 | -------------------------------------------------------------------------------- /tests/test_data/test_transform/test_color.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from nanodet.data.transform.color import ( 4 | normalize, 5 | random_brightness, 6 | random_contrast, 7 | random_saturation, 8 | ) 9 | 10 | 11 | def test_random_color_aug(): 12 | img = np.ones((10, 10, 3), dtype=np.float32) 13 | 14 | res = random_brightness(img, 0.6) 15 | assert np.max(res) <= 1.6 16 | assert np.min(res) >= 0.4 17 | 18 | img = np.ones((10, 10, 3), dtype=np.float32) 19 | res = random_contrast(img, 0.5, 1.5) 20 | assert np.max(res) <= 1.5 21 | assert np.min(res) >= 0.5 22 | 23 | img = np.ones((10, 10, 3), dtype=np.float32) 24 | random_saturation(img, 0.5, 1.5) 25 | 26 | img = np.ones((10, 10, 3), dtype=np.uint8) * 255 27 | meta = dict(img=img) 28 | res = normalize(meta, [100, 100, 100], [155, 155, 155]) 29 | assert np.array_equal(res["img"], np.ones((10, 10, 3))) 30 | -------------------------------------------------------------------------------- /demo_libtorch/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.4.1) 2 | set(CMAKE_CXX_STANDARD 14) 3 | 4 | project(nanodet_demo) 5 | 6 | find_package(OpenCV REQUIRED) 7 | find_package(Torch REQUIRED) 8 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") 9 | 10 | include_directories( 11 | ${OpenCV_INCLUDE_DIRS} 12 | ${CMAKE_CURRENT_SOURCE_DIR} 13 | ${CMAKE_CURRENT_BINARY_DIR} 14 | ) 15 | 16 | add_executable(nanodet_demo main.cpp nanodet_libtorch.cpp) 17 | 18 | target_link_libraries( 19 | nanodet_demo 20 | ${TORCH_LIBRARIES} 21 | ${OpenCV_LIBS} 22 | ) 23 | 24 | if (MSVC) 25 | file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll") 26 | add_custom_command(TARGET nanodet_demo 27 | POST_BUILD 28 | COMMAND ${CMAKE_COMMAND} -E copy_if_different 29 | ${TORCH_DLLS} 30 | $) 31 | endif (MSVC) 32 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/androidTest/java/com/rangi/nanodet/ExampleInstrumentedTest.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.Context; 4 | 5 | import androidx.test.platform.app.InstrumentationRegistry; 6 | import androidx.test.ext.junit.runners.AndroidJUnit4; 7 | 8 | import org.junit.Test; 9 | import org.junit.runner.RunWith; 10 | 11 | import static org.junit.Assert.*; 12 | 13 | /** 14 | * Instrumented test, which will execute on an Android device. 15 | * 16 | * @see Testing documentation 17 | */ 18 | @RunWith(AndroidJUnit4.class) 19 | public class ExampleInstrumentedTest { 20 | @Test 21 | public void useAppContext() { 22 | // Context of the app under test. 23 | Context appContext = InstrumentationRegistry.getInstrumentation().getTargetContext(); 24 | 25 | assertEquals("gd.hq.yolov5", appContext.getPackageName()); 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /demo_libtorch/README.md: -------------------------------------------------------------------------------- 1 | # NanoDet TorchScript / LibTorch Demo 2 | 3 | This folder provides NanoDet inference code using for LibTorch. 4 | 5 | ## Install dependencies 6 | 7 | This project needs OpenCV and CMake to work. 8 | 9 | Install CMake using a package manager of your choice. For example, the following command will install CMake on Ubuntu: 10 | 11 | ```bash 12 | sudo apt install cmake libopencv-dev 13 | ``` 14 | 15 | Also, you'll need to download LibTorch. Refer to [this page](https://pytorch.org/cppdocs/installing.html) for more info. 16 | 17 | ## Convert model 18 | 19 | Export TorchScript model using `tools/export_torchscript.py`: 20 | 21 | ```shell 22 | python ./tools/export_torchscript.py --cfg_path ${CONFIG_PATH} --model_path ${PYTORCH_MODEL_PATH} --input_shape ${MO} 23 | ``` 24 | ## Build 25 | 26 | ### Linux 27 | ```shell 28 | mkdir build 29 | cd build 30 | cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch .. 31 | make 32 | ``` 33 | -------------------------------------------------------------------------------- /tests/test_data/test_dataset/test_xmldataset.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from nanodet.data.dataset import XMLDataset, build_dataset 4 | 5 | 6 | def test_xmldataset(): 7 | class_names = ["asuka", "head"] 8 | cfg = dict( 9 | name="XMLDataset", 10 | class_names=class_names, 11 | img_path="./tests/data", 12 | ann_path="./tests/data", 13 | input_size=[320, 320], # [w,h] 14 | keep_ratio=True, 15 | pipeline=dict(normalize=[[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]), 16 | ) 17 | dataset = build_dataset(cfg, "train") 18 | assert isinstance(dataset, XMLDataset) 19 | 20 | for i, data in enumerate(dataset): 21 | assert data["img"].shape == (3, 320, 285) 22 | 23 | dataset = build_dataset(cfg, "val") 24 | for i, data in enumerate(dataset): 25 | assert data["img"].shape == (3, 320, 285) 26 | 27 | with pytest.raises(AssertionError): 28 | build_dataset(cfg, "2333") 29 | -------------------------------------------------------------------------------- /tests/test_optim/test_builder.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | from nanodet.optim import build_optimizer 7 | 8 | 9 | class ToyModel(nn.Module): 10 | def __init__(self) -> None: 11 | super().__init__() 12 | self.conv = nn.Conv2d(3, 14, 3) 13 | self.bn = nn.BatchNorm2d(32) 14 | self.act = nn.ReLU() 15 | 16 | def forward(self, x): 17 | return self.act(self.bn(self.conv(x))) 18 | 19 | 20 | class TestOptimBuilder(TestCase): 21 | def test_build_optimizer(self): 22 | model = ToyModel() 23 | config = dict( 24 | name="SGD", 25 | lr=0.001, 26 | momentum=0.9, 27 | param_level_cfg=dict(conv=dict(lr_mult=0.3, decay_mult=0.7)), 28 | no_bias_decay=True, 29 | no_norm_decay=True, 30 | ) 31 | optimizer = build_optimizer(model, config) 32 | assert isinstance(optimizer, torch.optim.SGD) 33 | -------------------------------------------------------------------------------- /demo_ncnn/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.4.1) 2 | set(CMAKE_CXX_STANDARD 17) 3 | 4 | project(nanodet_demo) 5 | 6 | find_package(OpenMP REQUIRED) 7 | if(OPENMP_FOUND) 8 | message("OPENMP FOUND") 9 | set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}") 10 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}") 11 | set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}") 12 | endif() 13 | 14 | find_package(OpenCV REQUIRED) 15 | 16 | find_package(ncnn REQUIRED) 17 | if(NOT TARGET ncnn) 18 | message(WARNING "ncnn NOT FOUND! Please set ncnn_DIR environment variable") 19 | else() 20 | message("ncnn FOUND ") 21 | endif() 22 | 23 | include_directories( 24 | ${OpenCV_INCLUDE_DIRS} 25 | ${CMAKE_CURRENT_SOURCE_DIR} 26 | ${CMAKE_CURRENT_BINARY_DIR} 27 | ) 28 | 29 | 30 | add_executable(nanodet_demo main.cpp nanodet.cpp) 31 | 32 | target_link_libraries( 33 | nanodet_demo 34 | ncnn 35 | ${OpenCV_LIBS} 36 | ) 37 | -------------------------------------------------------------------------------- /nanodet/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import copy 15 | 16 | from .coco_detection import CocoDetectionEvaluator 17 | 18 | 19 | def build_evaluator(cfg, dataset): 20 | evaluator_cfg = copy.deepcopy(cfg) 21 | name = evaluator_cfg.pop("name") 22 | if name == "CocoDetectionEvaluator": 23 | return CocoDetectionEvaluator(dataset) 24 | else: 25 | raise NotImplementedError 26 | -------------------------------------------------------------------------------- /nanodet/model/weight_averager/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | 17 | from .ema import ExpMovingAverager 18 | 19 | 20 | def build_weight_averager(cfg, device="cpu"): 21 | cfg = copy.deepcopy(cfg) 22 | name = cfg.pop("name") 23 | if name == "ExpMovingAverager": 24 | return ExpMovingAverager(**cfg, device=device) 25 | else: 26 | raise NotImplementedError(f"{name} is not implemented") 27 | -------------------------------------------------------------------------------- /demo_libtorch/nanodet_libtorch.h: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | 5 | typedef struct BoxInfo 6 | { 7 | float x1; 8 | float y1; 9 | float x2; 10 | float y2; 11 | float score; 12 | int label; 13 | } BoxInfo; 14 | 15 | class NanoDet 16 | { 17 | public: 18 | NanoDet(const char* model_path); 19 | ~NanoDet(); 20 | torch::jit::script::Module Net; 21 | std::vector detect(cv::Mat image, float score_threshold, float nms_threshold); 22 | 23 | private: 24 | torch::Tensor preprocess(cv::Mat& image); 25 | void decode_infer(torch::Tensor& cls_pred, torch::Tensor& dis_pred, float threshold, std::vector>& results); 26 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride); 27 | static void nms(std::vector& result, float nms_threshold); 28 | std::vector strides_{ 8, 16, 32, 64 }; 29 | int input_size_ = 416; 30 | int num_class_ = 80; 31 | int reg_max_ = 7; 32 | 33 | }; 34 | -------------------------------------------------------------------------------- /tests/test_models/test_loss/test_gfocal_loss.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.loss.gfocal_loss import DistributionFocalLoss, QualityFocalLoss 5 | 6 | 7 | def test_qfl(): 8 | with pytest.raises(AssertionError): 9 | QualityFocalLoss(use_sigmoid=False) 10 | 11 | label = torch.randint(low=0, high=7, size=(10,)) 12 | score = torch.rand((10,)) 13 | pred = torch.rand((10, 7)) 14 | target = (label, score) 15 | weight = torch.zeros(10) 16 | 17 | loss = QualityFocalLoss()(pred, target, weight) 18 | assert loss == 0.0 19 | 20 | loss = QualityFocalLoss()(pred, target, weight, reduction_override="sum") 21 | assert loss == 0.0 22 | 23 | 24 | def test_dfl(): 25 | 26 | pred = torch.rand((10, 7)) 27 | target = torch.rand((10,)) 28 | weight = torch.zeros(10) 29 | 30 | loss = DistributionFocalLoss()(pred, target, weight) 31 | assert loss == 0.0 32 | 33 | loss = DistributionFocalLoss()(pred, target, weight, reduction_override="sum") 34 | assert loss == 0.0 35 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | from setuptools import find_packages, setup 3 | 4 | from nanodet import __author__, __author_email__, __docs__, __homepage__, __version__ 5 | 6 | if __name__ == "__main__": 7 | setup( 8 | name="nanodet", 9 | version=__version__, 10 | description=__docs__, 11 | url=__homepage__, 12 | author=__author__, 13 | author_email=__author_email__, 14 | keywords="deep learning", 15 | packages=find_packages(exclude=("config", "tools", "demo")), 16 | classifiers=[ 17 | "Development Status :: Beta", 18 | "License :: OSI Approved :: Apache Software License", 19 | "Operating System :: OS Independent", 20 | "Programming Language :: Python :: 3.5", 21 | "Programming Language :: Python :: 3.6", 22 | "Programming Language :: Python :: 3.7", 23 | "Programming Language :: Python :: 3.8", 24 | ], 25 | license="Apache License 2.0", 26 | zip_safe=False, 27 | ) 28 | -------------------------------------------------------------------------------- /tests/test_evaluator/test_coco_detection.py: -------------------------------------------------------------------------------- 1 | import tempfile 2 | 3 | from nanodet.data.dataset import build_dataset 4 | from nanodet.evaluator import build_evaluator 5 | 6 | 7 | def test_coco_detection(): 8 | dummy_results = { 9 | 0: {0: [[0, 0, 20, 20, 1]], 1: [[0, 0, 20, 20, 1]]}, 10 | 1: {0: [[0, 0, 20, 20, 1]]}, 11 | } 12 | 13 | cfg = dict( 14 | name="CocoDataset", 15 | img_path="./tests/data", 16 | ann_path="./tests/data/dummy_coco.json", 17 | input_size=[320, 320], # [w,h] 18 | keep_ratio=True, 19 | pipeline=dict(normalize=[[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]), 20 | ) 21 | dataset = build_dataset(cfg, "train") 22 | 23 | eval_cfg = dict(name="CocoDetectionEvaluator", save_key="mAP") 24 | 25 | evaluator = build_evaluator(eval_cfg, dataset) 26 | tmp_dir = tempfile.TemporaryDirectory() 27 | eval_results = evaluator.evaluate( 28 | results=dummy_results, save_dir=tmp_dir.name, rank=-1 29 | ) 30 | assert eval_results["mAP"] == 1 31 | -------------------------------------------------------------------------------- /tests/test_data/test_collate.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from nanodet.data.collate import collate_function 5 | 6 | 7 | def test_collate(): 8 | batch = [1.2, 2.3, 3.4] 9 | out = collate_function(batch) 10 | assert isinstance(out, torch.Tensor) 11 | 12 | batch = [1, 2, 3] 13 | out = collate_function(batch) 14 | assert isinstance(out, torch.Tensor) 15 | 16 | batch = ["1", "2", "3"] 17 | out = collate_function(batch) 18 | assert out == batch 19 | 20 | batch = [{"1": 1, "2": 1.2, "3": 1.2}, {"1": 2, "2": 1.3, "3": 1.4}] 21 | out = collate_function(batch) 22 | assert isinstance(out, dict) 23 | for k, v in out.items(): 24 | assert isinstance(v, torch.Tensor) 25 | 26 | batch = [np.array([1, 2, 3]), np.array([4, 6, 8])] 27 | out = collate_function(batch) 28 | assert out == batch 29 | 30 | batch = [torch.randn((3, 20, 20)), torch.randn((3, 20, 20))] 31 | out = collate_function(batch) 32 | assert isinstance(out, torch.Tensor) 33 | assert out.shape == (2, 3, 20, 20) 34 | -------------------------------------------------------------------------------- /tests/test_models/test_weight_averager/test_ema.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. All rights reserved. 2 | import copy 3 | 4 | import pytest 5 | import torch 6 | import torch.nn as nn 7 | 8 | from nanodet.model.weight_averager import ExpMovingAverager 9 | 10 | 11 | def test_ema(): 12 | # test invalid input 13 | with pytest.raises(ValueError): 14 | ExpMovingAverager(-1) 15 | # test invalid input 16 | with pytest.raises(ValueError): 17 | ExpMovingAverager(10) 18 | 19 | class DummyModel(nn.Module): 20 | def __init__(self): 21 | super(DummyModel, self).__init__() 22 | self.fc = nn.Linear(10, 10) 23 | 24 | def forward(self, x): 25 | return self.fc(x) 26 | 27 | averager = ExpMovingAverager(0.5) 28 | model = DummyModel() 29 | ema_model = copy.deepcopy(model) 30 | 31 | averager.load_from(model) 32 | assert averager.has_inited() 33 | 34 | averager.update(model, 1) 35 | averager.apply_to(ema_model) 36 | assert torch.allclose(ema_model.fc.weight, model.fc.weight) 37 | -------------------------------------------------------------------------------- /nanodet/util/config.py: -------------------------------------------------------------------------------- 1 | from .yacs import CfgNode 2 | 3 | cfg = CfgNode(new_allowed=True) 4 | cfg.save_dir = "./" 5 | # common params for NETWORK 6 | cfg.model = CfgNode(new_allowed=True) 7 | cfg.model.arch = CfgNode(new_allowed=True) 8 | cfg.model.arch.backbone = CfgNode(new_allowed=True) 9 | cfg.model.arch.fpn = CfgNode(new_allowed=True) 10 | cfg.model.arch.head = CfgNode(new_allowed=True) 11 | 12 | # DATASET related params 13 | cfg.data = CfgNode(new_allowed=True) 14 | cfg.data.train = CfgNode(new_allowed=True) 15 | cfg.data.val = CfgNode(new_allowed=True) 16 | cfg.device = CfgNode(new_allowed=True) 17 | cfg.device.precision = 32 18 | # train 19 | cfg.schedule = CfgNode(new_allowed=True) 20 | 21 | # logger 22 | cfg.log = CfgNode() 23 | cfg.log.interval = 50 24 | 25 | # testing 26 | cfg.test = CfgNode() 27 | # size of images for each device 28 | 29 | 30 | def load_config(cfg, args_cfg): 31 | cfg.defrost() 32 | cfg.merge_from_file(args_cfg) 33 | cfg.freeze() 34 | 35 | 36 | if __name__ == "__main__": 37 | import sys 38 | 39 | with open(sys.argv[1], "w") as f: 40 | print(cfg, file=f) 41 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradle.properties: -------------------------------------------------------------------------------- 1 | # Project-wide Gradle settings. 2 | # IDE (e.g. Android Studio) users: 3 | # Gradle settings configured through the IDE *will override* 4 | # any settings specified in this file. 5 | # For more details on how to configure your build environment visit 6 | # http://www.gradle.org/docs/current/userguide/build_environment.html 7 | # Specifies the JVM arguments used for the daemon process. 8 | # The setting is particularly useful for tweaking memory settings. 9 | org.gradle.jvmargs=-Xmx1536m 10 | # When configured, Gradle will run in incubating parallel mode. 11 | # This option should only be used with decoupled projects. More details, visit 12 | # http://www.gradle.org/docs/current/userguide/multi_project_builds.html#sec:decoupled_projects 13 | # org.gradle.parallel=true 14 | # AndroidX package structure to make it clearer which packages are bundled with the 15 | # Android operating system, and which are packaged with your app's APK 16 | # https://developer.android.com/topic/libraries/support-library/androidx-rn 17 | android.useAndroidX=true 18 | # Automatically convert third-party libraries to use AndroidX 19 | android.enableJetifier=true 20 | -------------------------------------------------------------------------------- /nanodet/model/fpn/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | 17 | from .fpn import FPN 18 | from .ghost_pan import GhostPAN 19 | from .pan import PAN 20 | from .tan import TAN 21 | 22 | 23 | def build_fpn(cfg): 24 | fpn_cfg = copy.deepcopy(cfg) 25 | name = fpn_cfg.pop("name") 26 | if name == "FPN": 27 | return FPN(**fpn_cfg) 28 | elif name == "PAN": 29 | return PAN(**fpn_cfg) 30 | elif name == "TAN": 31 | return TAN(**fpn_cfg) 32 | elif name == "GhostPAN": 33 | return GhostPAN(**fpn_cfg) 34 | else: 35 | raise NotImplementedError 36 | -------------------------------------------------------------------------------- /nanodet/util/path.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | 17 | from .rank_filter import rank_filter 18 | 19 | 20 | @rank_filter 21 | def mkdir(path): 22 | if not os.path.exists(path): 23 | os.makedirs(path) 24 | 25 | 26 | def collect_files(path, exts): 27 | file_paths = [] 28 | for maindir, subdir, filename_list in os.walk(path): 29 | for filename in filename_list: 30 | file_path = os.path.join(maindir, filename) 31 | ext = os.path.splitext(file_path)[1] 32 | if ext in exts: 33 | file_paths.append(file_path) 34 | return file_paths 35 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_timm_wrapper.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nanodet.model.backbone import build_backbone 4 | from nanodet.model.backbone.timm_wrapper import TIMMWrapper 5 | 6 | 7 | def test_timm_wrapper(): 8 | cfg = dict( 9 | name="TIMMWrapper", 10 | model_name="resnet18", 11 | features_only=True, 12 | pretrained=False, 13 | output_stride=32, 14 | out_indices=(1, 2, 3, 4), 15 | ) 16 | model = build_backbone(cfg) 17 | 18 | input = torch.rand(1, 3, 64, 64) 19 | output = model(input) 20 | assert len(output) == 4 21 | assert output[0].shape == (1, 64, 16, 16) 22 | assert output[1].shape == (1, 128, 8, 8) 23 | assert output[2].shape == (1, 256, 4, 4) 24 | assert output[3].shape == (1, 512, 2, 2) 25 | 26 | model = TIMMWrapper( 27 | model_name="mobilenetv3_large_100", 28 | features_only=True, 29 | pretrained=False, 30 | output_stride=32, 31 | out_indices=(1, 2, 3, 4), 32 | ) 33 | output = model(input) 34 | 35 | assert len(output) == 4 36 | assert output[0].shape == (1, 24, 16, 16) 37 | assert output[1].shape == (1, 40, 8, 8) 38 | assert output[2].shape == (1, 112, 4, 4) 39 | assert output[3].shape == (1, 960, 2, 2) 40 | -------------------------------------------------------------------------------- /nanodet/data/batch_process.py: -------------------------------------------------------------------------------- 1 | from typing import Sequence 2 | 3 | import torch 4 | import torch.nn.functional as F 5 | 6 | 7 | def stack_batch_img( 8 | img_tensors: Sequence[torch.Tensor], divisible: int = 0, pad_value: float = 0.0 9 | ) -> torch.Tensor: 10 | """ 11 | Args: 12 | img_tensors (Sequence[torch.Tensor]): 13 | divisible (int): 14 | pad_value (float): value to pad 15 | 16 | Returns: 17 | torch.Tensor. 18 | """ 19 | assert len(img_tensors) > 0 20 | assert isinstance(img_tensors, (tuple, list)) 21 | assert divisible >= 0 22 | img_heights = [] 23 | img_widths = [] 24 | for img in img_tensors: 25 | assert img.shape[:-2] == img_tensors[0].shape[:-2] 26 | img_heights.append(img.shape[-2]) 27 | img_widths.append(img.shape[-1]) 28 | max_h, max_w = max(img_heights), max(img_widths) 29 | if divisible > 0: 30 | max_h = (max_h + divisible - 1) // divisible * divisible 31 | max_w = (max_w + divisible - 1) // divisible * divisible 32 | 33 | batch_imgs = [] 34 | for img in img_tensors: 35 | padding_size = [0, max_w - img.shape[-1], 0, max_h - img.shape[-2]] 36 | batch_imgs.append(F.pad(img, padding_size, value=pad_value)) 37 | return torch.stack(batch_imgs, dim=0).contiguous() 38 | -------------------------------------------------------------------------------- /nanodet/util/__init__.py: -------------------------------------------------------------------------------- 1 | from .box_transform import bbox2distance, distance2bbox 2 | from .check_point import ( 3 | convert_avg_params, 4 | convert_old_model, 5 | load_model_weight, 6 | save_model, 7 | ) 8 | from .config import cfg, load_config 9 | from .flops_counter import get_model_complexity_info 10 | from .logger import AverageMeter, Logger, MovingAverage, NanoDetLightningLogger 11 | from .misc import images_to_levels, multi_apply, unmap 12 | from .path import collect_files, mkdir 13 | from .rank_filter import rank_filter 14 | from .scatter_gather import gather_results, scatter_kwargs 15 | from .util_mixins import NiceRepr 16 | from .visualization import Visualizer, overlay_bbox_cv 17 | 18 | __all__ = [ 19 | "distance2bbox", 20 | "bbox2distance", 21 | "convert_old_model", 22 | "load_model_weight", 23 | "save_model", 24 | "cfg", 25 | "load_config", 26 | "get_model_complexity_info", 27 | "AverageMeter", 28 | "Logger", 29 | "MovingAverage", 30 | "images_to_levels", 31 | "multi_apply", 32 | "unmap", 33 | "mkdir", 34 | "rank_filter", 35 | "gather_results", 36 | "scatter_kwargs", 37 | "NiceRepr", 38 | "Visualizer", 39 | "overlay_bbox_cv", 40 | "collect_files", 41 | "NanoDetLightningLogger", 42 | "convert_avg_params", 43 | ] 44 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_custom_csp.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import CustomCspNet, build_backbone 5 | 6 | 7 | def test_custom_csp(): 8 | with pytest.raises(AssertionError): 9 | cfg = dict( 10 | name="CustomCspNet", net_cfg=[["Conv", 3, 32, 3, 2]], out_stages=(8, 9) 11 | ) 12 | build_backbone(cfg) 13 | 14 | with pytest.raises(AssertionError): 15 | CustomCspNet(net_cfg=dict(a=1), out_stages=(0, 1), activation="ReLU6") 16 | 17 | input = torch.rand(1, 3, 64, 64) 18 | out_stages = (0, 1, 2, 3, 4, 5) 19 | net_cfg = [ 20 | ["Conv", 3, 32, 3, 2], # 1/2 21 | ["MaxPool", 3, 2], # 1/4 22 | ["CspBlock", 32, 1, 3, 1], # 1/4 23 | ["CspBlock", 64, 2, 3, 2], # 1/8 24 | ["CspBlock", 128, 2, 3, 2], # 1/16 25 | ["CspBlock", 256, 3, 3, 2], # 1/32 26 | ] 27 | model = CustomCspNet(net_cfg=net_cfg, out_stages=out_stages, activation="ReLU6") 28 | output = model(input) 29 | 30 | assert output[0].shape == (1, 32, 32, 32) 31 | assert output[1].shape == (1, 32, 16, 16) 32 | assert output[2].shape == (1, 64, 16, 16) 33 | assert output[3].shape == (1, 128, 8, 8) 34 | assert output[4].shape == (1, 256, 4, 4) 35 | assert output[5].shape == (1, 512, 2, 2) 36 | -------------------------------------------------------------------------------- /tests/test_models/test_fpn/test_tan.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.fpn.tan import TAN 5 | 6 | 7 | def test_tan(): 8 | """Tests TAN.""" 9 | s = 64 10 | in_channels = [8, 16, 32] 11 | feat_sizes = [s // 2**i for i in range(3)] # [64, 32, 16] 12 | out_channels = 8 13 | 14 | with pytest.raises(AssertionError): 15 | TAN( 16 | in_channels=[8, 16, 32, 64], 17 | out_channels=[8, 16, 32, 64], 18 | feature_hw=[32, 32], 19 | num_heads=4, 20 | num_encoders=1, 21 | mlp_ratio=2, 22 | dropout_ratio=0.9, 23 | ) 24 | 25 | pan_model = TAN( 26 | in_channels=in_channels, 27 | out_channels=out_channels, 28 | feature_hw=[32, 32], 29 | num_heads=4, 30 | num_encoders=1, 31 | mlp_ratio=2, 32 | dropout_ratio=0.9, 33 | ) 34 | 35 | # TAN expects a multiple levels of features per image 36 | feats = [ 37 | torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i]) 38 | for i in range(len(in_channels)) 39 | ] 40 | outs = pan_model(feats) 41 | assert len(outs) == 3 42 | for i in range(3): 43 | assert outs[i].shape[1] == out_channels 44 | assert outs[i].shape[2] == outs[i].shape[3] == s // (2**i) 45 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/AndroidManifest.xml: -------------------------------------------------------------------------------- 1 | 2 | 4 | 5 | 6 | 7 | 8 | 9 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 28 | 29 | 30 | 31 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_mobilenetv2.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import MobileNetV2, build_backbone 5 | 6 | 7 | def test_mobilenetv2(): 8 | with pytest.raises(AssertionError): 9 | cfg = dict(name="MobileNetV2", width_mult=1.0, out_stages=(8, 9)) 10 | build_backbone(cfg) 11 | 12 | input = torch.rand(1, 3, 64, 64) 13 | out_stages = (0, 1, 2, 3, 4, 5, 6) 14 | model = MobileNetV2(width_mult=1.0, out_stages=out_stages, activation="ReLU6") 15 | output = model(input) 16 | 17 | assert output[0].shape == (1, 16, 32, 32) 18 | assert output[1].shape == (1, 24, 16, 16) 19 | assert output[2].shape == (1, 32, 8, 8) 20 | assert output[3].shape == (1, 64, 4, 4) 21 | assert output[4].shape == (1, 96, 4, 4) 22 | assert output[5].shape == (1, 160, 2, 2) 23 | assert output[6].shape == (1, 1280, 2, 2) 24 | 25 | model = MobileNetV2(width_mult=0.75, out_stages=out_stages, activation="LeakyReLU") 26 | output = model(input) 27 | 28 | assert output[0].shape == (1, 12, 32, 32) 29 | assert output[1].shape == (1, 18, 16, 16) 30 | assert output[2].shape == (1, 24, 8, 8) 31 | assert output[3].shape == (1, 48, 4, 4) 32 | assert output[4].shape == (1, 72, 4, 4) 33 | assert output[5].shape == (1, 120, 2, 2) 34 | assert output[6].shape == (1, 1280, 2, 2) 35 | -------------------------------------------------------------------------------- /nanodet/model/module/activation.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import torch.nn as nn 16 | 17 | activations = { 18 | "ReLU": nn.ReLU, 19 | "LeakyReLU": nn.LeakyReLU, 20 | "ReLU6": nn.ReLU6, 21 | "SELU": nn.SELU, 22 | "ELU": nn.ELU, 23 | "GELU": nn.GELU, 24 | "PReLU": nn.PReLU, 25 | "SiLU": nn.SiLU, 26 | "HardSwish": nn.Hardswish, 27 | "Hardswish": nn.Hardswish, 28 | None: nn.Identity, 29 | } 30 | 31 | 32 | def act_layers(name): 33 | assert name in activations.keys() 34 | if name == "LeakyReLU": 35 | return nn.LeakyReLU(negative_slope=0.1, inplace=True) 36 | elif name == "GELU": 37 | return nn.GELU() 38 | elif name == "PReLU": 39 | return nn.PReLU() 40 | else: 41 | return activations[name](inplace=True) 42 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_dwconv.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | import torch.nn as nn 4 | 5 | from nanodet.model.module.conv import DepthwiseConvModule 6 | 7 | 8 | def test_depthwise_conv(): 9 | with pytest.raises(AssertionError): 10 | activation = dict(type="softmax") 11 | DepthwiseConvModule(4, 5, 2, activation=activation) 12 | 13 | with pytest.raises(AssertionError): 14 | DepthwiseConvModule(3, 5, 2, order=("norm", "conv", "act")) 15 | 16 | # test default config 17 | conv = DepthwiseConvModule(3, 5, 2) 18 | assert conv.with_norm 19 | assert conv.depthwise.groups == 3 20 | assert conv.pointwise.kernel_size == (1, 1) 21 | assert conv.act.__class__.__name__ == "ReLU" 22 | x = torch.rand(1, 3, 16, 16) 23 | output = conv(x) 24 | assert output.shape == (1, 5, 15, 15) 25 | 26 | # test norm_cfg 27 | conv = DepthwiseConvModule(3, 5, 2, norm_cfg=dict(type="BN")) 28 | assert isinstance(conv.dwnorm, nn.BatchNorm2d) 29 | assert isinstance(conv.pwnorm, nn.BatchNorm2d) 30 | 31 | x = torch.rand(1, 3, 16, 16) 32 | output = conv(x) 33 | assert output.shape == (1, 5, 15, 15) 34 | 35 | # test act_cfg 36 | conv = DepthwiseConvModule(3, 5, 3, padding=1, activation="LeakyReLU") 37 | assert conv.act.__class__.__name__ == "LeakyReLU" 38 | output = conv(x) 39 | assert output.shape == (1, 5, 16, 16) 40 | -------------------------------------------------------------------------------- /tests/data/test_img.xml: -------------------------------------------------------------------------------- 1 | 2 | data 3 | test_img.jpg 4 | tests/data/test_img.jpg 5 | 6 | Unknown 7 | 8 | 9 | 360 10 | 404 11 | 3 12 | 13 | 0 14 | 15 | asuka 16 | Unspecified 17 | 0 18 | 0 19 | 20 | 7 21 | 3 22 | 189 23 | 401 24 | 25 | 26 | 27 | asuka 28 | Unspecified 29 | 1 30 | 0 31 | 32 | 171 33 | 20 34 | 360 35 | 404 36 | 37 | 38 | 39 | head 40 | Unspecified 41 | 0 42 | 0 43 | 44 | 50 45 | 4 46 | 181 47 | 140 48 | 49 | 50 | 51 | head 52 | Unspecified 53 | 0 54 | 0 55 | 56 | 170 57 | 20 58 | 303 59 | 133 60 | 61 | 62 | 63 | -------------------------------------------------------------------------------- /tools/convert_old_checkpoint.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import argparse 16 | 17 | import torch 18 | 19 | from nanodet.util import convert_old_model 20 | 21 | 22 | def parse_args(): 23 | parser = argparse.ArgumentParser( 24 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, 25 | description="Convert .pth model to onnx.", 26 | ) 27 | parser.add_argument("--file_path", type=str, help="Path to .pth checkpoint.") 28 | parser.add_argument("--out_path", type=str, help="Path to .ckpt checkpoint.") 29 | return parser.parse_args() 30 | 31 | 32 | if __name__ == "__main__": 33 | args = parse_args() 34 | file_path = args.file_path 35 | out_path = args.out_path 36 | old_check_point = torch.load(file_path) 37 | new_check_point = convert_old_model(old_check_point) 38 | torch.save(new_check_point, out_path) 39 | print("Checkpoint saved to:", out_path) 40 | -------------------------------------------------------------------------------- /docs/update.md: -------------------------------------------------------------------------------- 1 | # Update Notes 2 | 3 | * [2021.08.28] Refactor data processing pipeline and support multi-scale training (#311). 4 | 5 | * [2021.05.30] Release ncnn int8 models, and new pre-trained models with ShuffleNetV2-1.5x backbone. Much higher mAP but still realtime(**26.8mAP 21.53ms**). 6 | 7 | * [2021.03.12] Apply the **Transformer** encoder to NanoDet! Introducing **NanoDet-t**, which replaces the PAN in NanoDet-m with a **TAN(Transformer Attention Net)**, gets 21.7 mAP(+1.1) on COCO val 2017. Check [nanodet-t.yml](config/Transformer/nanodet-t.yml) for more details. 8 | 9 | * [2021.03.03] Update **Nanodet-m-416** COCO pretrained model. **COCO mAP(0.5:0.95)=23.5**. Download in [Model Zoo](#model-zoo). 10 | 11 | * [2021.02.03] Support [EfficientNet-Lite](https://github.com/RangiLyu/EfficientNet-Lite) and [Rep-VGG](https://github.com/DingXiaoH/RepVGG) backbone. Please check the [config folder](config/). Download models in [Model Zoo](#model-zoo) 12 | 13 | * [2021.01.10] **NanoDet-g** with lower memory access cost, which designed for edge NPU or GPU, is now available! 14 | Check [config/nanodet-g.yml](config/nanodet-g.yml) and download in [Model Zoo](#model-zoo). 15 | 16 | * [2020.12.19] [MNN python and cpp demos](demo_mnn/) are available. 17 | 18 | * [2020.12.05] Support voc .xml format dataset! Refer to [config/nanodet_custom_xml_dataset.yml](config/nanodet_custom_xml_dataset.yml). 19 | 20 | * [2020.12.01] Great thanks to nihui, now you can try NanoDet running in web browser! 👉 https://nihui.github.io/ncnn-webassembly-nanodet/ 21 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_repvgg.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import RepVGG, build_backbone 5 | from nanodet.model.backbone.repvgg import repvgg_model_convert 6 | 7 | 8 | def test_repvgg(): 9 | with pytest.raises(AssertionError): 10 | cfg = dict(name="RepVGG", arch="A3") 11 | build_backbone(cfg) 12 | 13 | with pytest.raises(AssertionError): 14 | RepVGG(arch="A0", out_stages=(4, 5, 6)) 15 | 16 | input = torch.rand(1, 3, 64, 64) 17 | 18 | model = RepVGG(arch="A0", out_stages=(1, 2, 3, 4), activation="PReLU") 19 | output = model(input) 20 | 21 | assert output[0].shape == (1, 48, 16, 16) 22 | assert output[1].shape == (1, 96, 8, 8) 23 | assert output[2].shape == (1, 192, 4, 4) 24 | assert output[3].shape == (1, 1280, 2, 2) 25 | 26 | # test last channel 27 | model = RepVGG(arch="A1", out_stages=(1, 2, 3, 4), last_channel=640) 28 | output = model(input) 29 | 30 | assert output[0].shape == (1, 64, 16, 16) 31 | assert output[1].shape == (1, 128, 8, 8) 32 | assert output[2].shape == (1, 256, 4, 4) 33 | assert output[3].shape == (1, 640, 2, 2) 34 | 35 | deploy_model = RepVGG(arch="A0", deploy=True) 36 | deploy_model = repvgg_model_convert(model, deploy_model, save_path=None) 37 | dep_output = deploy_model(input) 38 | assert dep_output[0].shape == (1, 64, 16, 16) 39 | assert dep_output[1].shape == (1, 128, 8, 8) 40 | assert dep_output[2].shape == (1, 256, 4, 4) 41 | assert dep_output[3].shape == (1, 640, 2, 2) 42 | -------------------------------------------------------------------------------- /nanodet/model/module/init_weights.py: -------------------------------------------------------------------------------- 1 | # Modification 2020 RangiLyu 2 | # Copyright 2018-2019 Open-MMLab. 3 | 4 | import torch.nn as nn 5 | 6 | 7 | def kaiming_init( 8 | module, a=0, mode="fan_out", nonlinearity="relu", bias=0, distribution="normal" 9 | ): 10 | assert distribution in ["uniform", "normal"] 11 | if distribution == "uniform": 12 | nn.init.kaiming_uniform_( 13 | module.weight, a=a, mode=mode, nonlinearity=nonlinearity 14 | ) 15 | else: 16 | nn.init.kaiming_normal_( 17 | module.weight, a=a, mode=mode, nonlinearity=nonlinearity 18 | ) 19 | if hasattr(module, "bias") and module.bias is not None: 20 | nn.init.constant_(module.bias, bias) 21 | 22 | 23 | def xavier_init(module, gain=1, bias=0, distribution="normal"): 24 | assert distribution in ["uniform", "normal"] 25 | if distribution == "uniform": 26 | nn.init.xavier_uniform_(module.weight, gain=gain) 27 | else: 28 | nn.init.xavier_normal_(module.weight, gain=gain) 29 | if hasattr(module, "bias") and module.bias is not None: 30 | nn.init.constant_(module.bias, bias) 31 | 32 | 33 | def normal_init(module, mean=0, std=1, bias=0): 34 | nn.init.normal_(module.weight, mean, std) 35 | if hasattr(module, "bias") and module.bias is not None: 36 | nn.init.constant_(module.bias, bias) 37 | 38 | 39 | def constant_init(module, val, bias=0): 40 | if hasattr(module, "weight") and module.weight is not None: 41 | nn.init.constant_(module.weight, val) 42 | if hasattr(module, "bias") and module.bias is not None: 43 | nn.init.constant_(module.bias, bias) 44 | -------------------------------------------------------------------------------- /nanodet/model/arch/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | import warnings 17 | 18 | from .nanodet_plus import NanoDetPlus 19 | from .one_stage_detector import OneStageDetector 20 | 21 | 22 | def build_model(model_cfg): 23 | model_cfg = copy.deepcopy(model_cfg) 24 | name = model_cfg.arch.pop("name") 25 | if name == "GFL": 26 | warnings.warn( 27 | "Model architecture name is changed to 'OneStageDetector'. " 28 | "The name 'GFL' is deprecated, please change the model->arch->name " 29 | "in your YAML config file to OneStageDetector." 30 | ) 31 | model = OneStageDetector( 32 | model_cfg.arch.backbone, model_cfg.arch.fpn, model_cfg.arch.head 33 | ) 34 | elif name == "OneStageDetector": 35 | model = OneStageDetector( 36 | model_cfg.arch.backbone, model_cfg.arch.fpn, model_cfg.arch.head 37 | ) 38 | elif name == "NanoDetPlus": 39 | model = NanoDetPlus(**model_cfg.arch) 40 | else: 41 | raise NotImplementedError 42 | return model 43 | -------------------------------------------------------------------------------- /tests/test_configs/test_config.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | from os.path import dirname, exists, join 17 | 18 | from nanodet.model.arch import build_model 19 | from nanodet.util import cfg, collect_files, load_config 20 | 21 | 22 | def test_config_files(): 23 | root_path = join(dirname(__file__), "../..") 24 | cfg_folder = join(root_path, "config") 25 | if not exists(cfg_folder): 26 | raise FileNotFoundError("Cannot find config folder.") 27 | 28 | cfg_paths = collect_files(cfg_folder, [".yml", ".yaml"]) 29 | for cfg_path in cfg_paths: 30 | print(f"Start testing {cfg_path}") 31 | config = copy.deepcopy(cfg) 32 | 33 | # test load cfg 34 | load_config(config, cfg_path) 35 | assert "save_dir" in config 36 | assert "model" in config 37 | assert "data" in config 38 | assert "device" in config 39 | assert "schedule" in config 40 | assert "log" in config 41 | 42 | # test build model 43 | model = build_model(config.model) 44 | assert config.model.arch.name == model.__class__.__name__ 45 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_resnet.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import ResNet, build_backbone 5 | 6 | 7 | def test_resnet(): 8 | with pytest.raises(KeyError): 9 | cfg = dict(name="ResNet", depth=15) 10 | build_backbone(cfg) 11 | 12 | with pytest.raises(AssertionError): 13 | ResNet(depth=18, out_stages=(4, 5, 6)) 14 | 15 | input = torch.rand(1, 3, 64, 64) 16 | 17 | model = ResNet(depth=18, out_stages=(1, 2, 3, 4), activation="PReLU", pretrain=True) 18 | output = model(input) 19 | 20 | assert output[0].shape == (1, 64, 16, 16) 21 | assert output[1].shape == (1, 128, 8, 8) 22 | assert output[2].shape == (1, 256, 4, 4) 23 | assert output[3].shape == (1, 512, 2, 2) 24 | 25 | model = ResNet( 26 | depth=34, out_stages=(1, 2, 3, 4), activation="LeakyReLU", pretrain=False 27 | ) 28 | output = model(input) 29 | assert output[0].shape == (1, 64, 16, 16) 30 | assert output[1].shape == (1, 128, 8, 8) 31 | assert output[2].shape == (1, 256, 4, 4) 32 | assert output[3].shape == (1, 512, 2, 2) 33 | 34 | model = ResNet(depth=50, out_stages=(1, 2, 3, 4), pretrain=False) 35 | output = model(input) 36 | assert output[0].shape == (1, 256, 16, 16) 37 | assert output[1].shape == (1, 512, 8, 8) 38 | assert output[2].shape == (1, 1024, 4, 4) 39 | assert output[3].shape == (1, 2048, 2, 2) 40 | 41 | model = ResNet(depth=101, out_stages=(1, 2, 3, 4), pretrain=False) 42 | output = model(input) 43 | assert output[0].shape == (1, 256, 16, 16) 44 | assert output[1].shape == (1, 512, 8, 8) 45 | assert output[2].shape == (1, 1024, 4, 4) 46 | assert output[3].shape == (1, 2048, 2, 2) 47 | -------------------------------------------------------------------------------- /nanodet/model/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | 17 | from .custom_csp import CustomCspNet 18 | from .efficientnet_lite import EfficientNetLite 19 | from .ghostnet import GhostNet 20 | from .mobilenetv2 import MobileNetV2 21 | from .repvgg import RepVGG 22 | from .resnet import ResNet 23 | from .shufflenetv2 import ShuffleNetV2 24 | from .timm_wrapper import TIMMWrapper 25 | 26 | 27 | def build_backbone(cfg): 28 | backbone_cfg = copy.deepcopy(cfg) 29 | name = backbone_cfg.pop("name") 30 | if name == "ResNet": 31 | return ResNet(**backbone_cfg) 32 | elif name == "ShuffleNetV2": 33 | return ShuffleNetV2(**backbone_cfg) 34 | elif name == "GhostNet": 35 | return GhostNet(**backbone_cfg) 36 | elif name == "MobileNetV2": 37 | return MobileNetV2(**backbone_cfg) 38 | elif name == "EfficientNetLite": 39 | return EfficientNetLite(**backbone_cfg) 40 | elif name == "CustomCspNet": 41 | return CustomCspNet(**backbone_cfg) 42 | elif name == "RepVGG": 43 | return RepVGG(**backbone_cfg) 44 | elif name == "TIMMWrapper": 45 | return TIMMWrapper(**backbone_cfg) 46 | else: 47 | raise NotImplementedError 48 | -------------------------------------------------------------------------------- /nanodet/data/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | import warnings 17 | 18 | from .coco import CocoDataset 19 | from .xml_dataset import XMLDataset 20 | from .yolo import YoloDataset 21 | 22 | 23 | def build_dataset(cfg, mode): 24 | dataset_cfg = copy.deepcopy(cfg) 25 | name = dataset_cfg.pop("name") 26 | if name == "coco": 27 | warnings.warn( 28 | "Dataset name coco has been deprecated. Please use CocoDataset instead." 29 | ) 30 | return CocoDataset(mode=mode, **dataset_cfg) 31 | elif name == "yolo": 32 | return YoloDataset(mode=mode, **dataset_cfg) 33 | elif name == "xml_dataset": 34 | warnings.warn( 35 | "Dataset name xml_dataset has been deprecated. " 36 | "Please use XMLDataset instead." 37 | ) 38 | return XMLDataset(mode=mode, **dataset_cfg) 39 | elif name == "CocoDataset": 40 | return CocoDataset(mode=mode, **dataset_cfg) 41 | elif name == "YoloDataset": 42 | return YoloDataset(mode=mode, **dataset_cfg) 43 | elif name == "XMLDataset": 44 | return XMLDataset(mode=mode, **dataset_cfg) 45 | else: 46 | raise NotImplementedError("Unknown dataset type!") 47 | -------------------------------------------------------------------------------- /tests/test_models/test_fpn/test_fpn.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.fpn.fpn import FPN 5 | 6 | 7 | def test_fpn(): 8 | """Tests fpn.""" 9 | s = 64 10 | in_channels = [8, 16, 32, 64] 11 | feat_sizes = [s // 2**i for i in range(4)] # [64, 32, 16, 8] 12 | out_channels = 8 13 | # `num_outs` is not equal to len(in_channels) - start_level 14 | with pytest.raises(AssertionError): 15 | FPN( 16 | in_channels=in_channels, 17 | out_channels=out_channels, 18 | start_level=1, 19 | num_outs=2, 20 | ) 21 | 22 | # `end_level` is larger than len(in_channels) - 1 23 | with pytest.raises(AssertionError): 24 | FPN( 25 | in_channels=in_channels, 26 | out_channels=out_channels, 27 | start_level=1, 28 | end_level=4, 29 | num_outs=2, 30 | ) 31 | 32 | # `num_outs` is not equal to end_level - start_level 33 | with pytest.raises(AssertionError): 34 | FPN( 35 | in_channels=in_channels, 36 | out_channels=out_channels, 37 | start_level=1, 38 | end_level=3, 39 | num_outs=1, 40 | ) 41 | 42 | fpn_model = FPN( 43 | in_channels=in_channels, out_channels=out_channels, start_level=1, num_outs=3 44 | ) 45 | 46 | # FPN expects a multiple levels of features per image 47 | feats = [ 48 | torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i]) 49 | for i in range(len(in_channels)) 50 | ] 51 | outs = fpn_model(feats) 52 | assert len(outs) == fpn_model.num_outs 53 | for i in range(fpn_model.num_outs): 54 | assert outs[i].shape[1] == out_channels 55 | assert outs[i].shape[2] == outs[i].shape[3] == s // (2 ** (i + 1)) 56 | -------------------------------------------------------------------------------- /tests/test_models/test_fpn/test_pan.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.fpn.pan import PAN 5 | 6 | 7 | def test_pan(): 8 | """Tests PAN.""" 9 | s = 64 10 | in_channels = [8, 16, 32, 64] 11 | feat_sizes = [s // 2**i for i in range(4)] # [64, 32, 16, 8] 12 | out_channels = 8 13 | # `num_outs` is not equal to len(in_channels) - start_level 14 | with pytest.raises(AssertionError): 15 | PAN( 16 | in_channels=in_channels, 17 | out_channels=out_channels, 18 | start_level=1, 19 | num_outs=2, 20 | ) 21 | 22 | # `end_level` is larger than len(in_channels) - 1 23 | with pytest.raises(AssertionError): 24 | PAN( 25 | in_channels=in_channels, 26 | out_channels=out_channels, 27 | start_level=1, 28 | end_level=4, 29 | num_outs=2, 30 | ) 31 | 32 | # `num_outs` is not equal to end_level - start_level 33 | with pytest.raises(AssertionError): 34 | PAN( 35 | in_channels=in_channels, 36 | out_channels=out_channels, 37 | start_level=1, 38 | end_level=3, 39 | num_outs=1, 40 | ) 41 | 42 | pan_model = PAN( 43 | in_channels=in_channels, out_channels=out_channels, start_level=1, num_outs=3 44 | ) 45 | 46 | # PAN expects a multiple levels of features per image 47 | feats = [ 48 | torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i]) 49 | for i in range(len(in_channels)) 50 | ] 51 | outs = pan_model(feats) 52 | assert len(outs) == pan_model.num_outs 53 | for i in range(pan_model.num_outs): 54 | assert outs[i].shape[1] == out_channels 55 | assert outs[i].shape[2] == outs[i].shape[3] == s // (2 ** (i + 1)) 56 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_transformer.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | import torch.nn as nn 4 | 5 | from nanodet.model.module.transformer import MLP, TransformerBlock, TransformerEncoder 6 | 7 | 8 | def test_mlp(): 9 | mlp = MLP(in_dim=10) 10 | assert mlp.fc1.in_features == 10 11 | assert mlp.fc1.out_features == 10 12 | assert mlp.fc2.in_features == 10 13 | assert mlp.fc2.out_features == 10 14 | 15 | mlp = MLP(in_dim=10, hidden_dim=5, out_dim=12, drop=0.2) 16 | assert mlp.fc1.in_features == 10 17 | assert mlp.fc1.out_features == 5 18 | assert mlp.fc2.in_features == 5 19 | assert mlp.fc2.out_features == 12 20 | assert mlp.drop.p == 0.2 21 | input = torch.rand(64, 10) 22 | output = mlp(input) 23 | assert output.shape == (64, 12) 24 | 25 | 26 | def test_tansformer_encoder(): 27 | # embed_dim must be divisible by num_heads 28 | with pytest.raises(AssertionError): 29 | TransformerEncoder(10, 6, 0.5) 30 | 31 | encoder = TransformerEncoder(10, 5, 1) 32 | assert isinstance(encoder.norm1, nn.LayerNorm) 33 | assert isinstance(encoder.norm2, nn.LayerNorm) 34 | assert encoder.attn.embed_dim == 10 35 | input = torch.rand(32, 1, 10) 36 | output = encoder(input) 37 | assert output.shape == (32, 1, 10) 38 | 39 | 40 | def test_transformer_block(): 41 | # out_channels must be divisible by num_heads 42 | with pytest.raises(AssertionError): 43 | TransformerBlock(10, 15, 4) 44 | 45 | # test in dim equal to out dim 46 | block = TransformerBlock(15, 15, 3) 47 | assert isinstance(block.conv, nn.Identity) 48 | 49 | block = TransformerBlock(12, 15, 3, 2) 50 | input = torch.rand(1, 12, 16, 16) 51 | pos_embed = torch.rand(16 * 16, 1, 15) 52 | out = block(input, pos_embed) 53 | assert out.shape == (1, 15, 16, 16) 54 | -------------------------------------------------------------------------------- /nanodet/util/box_transform.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def distance2bbox(points, distance, max_shape=None): 5 | """Decode distance prediction to bounding box. 6 | 7 | Args: 8 | points (Tensor): Shape (n, 2), [x, y]. 9 | distance (Tensor): Distance from the given point to 4 10 | boundaries (left, top, right, bottom). 11 | max_shape (tuple): Shape of the image. 12 | 13 | Returns: 14 | Tensor: Decoded bboxes. 15 | """ 16 | x1 = points[..., 0] - distance[..., 0] 17 | y1 = points[..., 1] - distance[..., 1] 18 | x2 = points[..., 0] + distance[..., 2] 19 | y2 = points[..., 1] + distance[..., 3] 20 | if max_shape is not None: 21 | x1 = x1.clamp(min=0, max=max_shape[1]) 22 | y1 = y1.clamp(min=0, max=max_shape[0]) 23 | x2 = x2.clamp(min=0, max=max_shape[1]) 24 | y2 = y2.clamp(min=0, max=max_shape[0]) 25 | return torch.stack([x1, y1, x2, y2], -1) 26 | 27 | 28 | def bbox2distance(points, bbox, max_dis=None, eps=0.1): 29 | """Decode bounding box based on distances. 30 | 31 | Args: 32 | points (Tensor): Shape (n, 2), [x, y]. 33 | bbox (Tensor): Shape (n, 4), "xyxy" format 34 | max_dis (float): Upper bound of the distance. 35 | eps (float): a small value to ensure target < max_dis, instead <= 36 | 37 | Returns: 38 | Tensor: Decoded distances. 39 | """ 40 | left = points[:, 0] - bbox[:, 0] 41 | top = points[:, 1] - bbox[:, 1] 42 | right = bbox[:, 2] - points[:, 0] 43 | bottom = bbox[:, 3] - points[:, 1] 44 | if max_dis is not None: 45 | left = left.clamp(min=0, max=max_dis - eps) 46 | top = top.clamp(min=0, max=max_dis - eps) 47 | right = right.clamp(min=0, max=max_dis - eps) 48 | bottom = bottom.clamp(min=0, max=max_dis - eps) 49 | return torch.stack([left, top, right, bottom], -1) 50 | -------------------------------------------------------------------------------- /nanodet/util/misc.py: -------------------------------------------------------------------------------- 1 | # Modification 2020 RangiLyu 2 | # Copyright 2018-2019 Open-MMLab. 3 | 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | 16 | from functools import partial 17 | 18 | import torch 19 | 20 | 21 | def multi_apply(func, *args, **kwargs): 22 | pfunc = partial(func, **kwargs) if kwargs else func 23 | map_results = map(pfunc, *args) 24 | return tuple(map(list, zip(*map_results))) 25 | 26 | 27 | def images_to_levels(target, num_level_anchors): 28 | """Convert targets by image to targets by feature level. 29 | 30 | [target_img0, target_img1] -> [target_level0, target_level1, ...] 31 | """ 32 | target = torch.stack(target, 0) 33 | level_targets = [] 34 | start = 0 35 | for n in num_level_anchors: 36 | end = start + n 37 | level_targets.append(target[:, start:end].squeeze(0)) 38 | start = end 39 | return level_targets 40 | 41 | 42 | def unmap(data, count, inds, fill=0): 43 | """Unmap a subset of item (data) back to the original set of items (of 44 | size count)""" 45 | if data.dim() == 1: 46 | ret = data.new_full((count,), fill) 47 | ret[inds.type(torch.bool)] = data 48 | else: 49 | new_size = (count,) + data.size()[1:] 50 | ret = data.new_full(new_size, fill) 51 | ret[inds.type(torch.bool), :] = data 52 | return ret 53 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_shufflenetv2.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import ShuffleNetV2, build_backbone 5 | 6 | 7 | def test_shufflenetv2(): 8 | 9 | with pytest.raises(NotImplementedError): 10 | cfg = dict(name="ShuffleNetV2", model_size="3.0x", pretrain=False) 11 | build_backbone(cfg) 12 | 13 | with pytest.raises(AssertionError): 14 | ShuffleNetV2("1.0x", out_stages=(1, 2, 3)) 15 | 16 | input = torch.rand(1, 3, 64, 64) 17 | model = ShuffleNetV2(model_size="0.5x", out_stages=(2, 3, 4), pretrain=True) 18 | output = model(input) 19 | assert output[0].shape == (1, 48, 8, 8) 20 | assert output[1].shape == (1, 96, 4, 4) 21 | assert output[2].shape == (1, 192, 2, 2) 22 | 23 | model = ShuffleNetV2(model_size="0.5x", out_stages=(3, 4), pretrain=True) 24 | output = model(input) 25 | assert output[0].shape == (1, 96, 4, 4) 26 | assert output[1].shape == (1, 192, 2, 2) 27 | 28 | model = ShuffleNetV2(model_size="1.0x", pretrain=False, with_last_conv=True) 29 | assert hasattr(model.stage4, "conv5") 30 | output = model(input) 31 | assert output[0].shape == (1, 116, 8, 8) 32 | assert output[1].shape == (1, 232, 4, 4) 33 | assert output[2].shape == (1, 1024, 2, 2) 34 | 35 | model = ShuffleNetV2( 36 | model_size="1.5x", pretrain=False, with_last_conv=False, activation="ReLU6" 37 | ) 38 | assert not hasattr(model.stage4, "conv5") 39 | output = model(input) 40 | assert output[0].shape == (1, 176, 8, 8) 41 | assert output[1].shape == (1, 352, 4, 4) 42 | assert output[2].shape == (1, 704, 2, 2) 43 | 44 | model = ShuffleNetV2(model_size="2.0x", pretrain=False, with_last_conv=False) 45 | output = model(input) 46 | assert output[0].shape == (1, 244, 8, 8) 47 | assert output[1].shape == (1, 488, 4, 4) 48 | assert output[2].shape == (1, 976, 2, 2) 49 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_nms.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | 3 | import torch 4 | 5 | from nanodet.model.module.nms import batched_nms, multiclass_nms 6 | 7 | 8 | def test_batched_nms(): 9 | file = open("./tests/data/batched_nms_data.pkl", "rb") 10 | results = pickle.load(file) 11 | 12 | nms_cfg = dict(iou_threshold=0.7) 13 | boxes, keep = batched_nms( 14 | torch.from_numpy(results["boxes"]), 15 | torch.from_numpy(results["scores"]), 16 | torch.from_numpy(results["idxs"]), 17 | nms_cfg, 18 | class_agnostic=False, 19 | ) 20 | 21 | nms_cfg.update(split_thr=100) 22 | seq_boxes, seq_keep = batched_nms( 23 | torch.from_numpy(results["boxes"]), 24 | torch.from_numpy(results["scores"]), 25 | torch.from_numpy(results["idxs"]), 26 | nms_cfg, 27 | class_agnostic=False, 28 | ) 29 | 30 | assert torch.equal(keep, seq_keep) 31 | assert torch.equal(boxes, seq_boxes) 32 | 33 | 34 | def test_multiclass_nms(): 35 | file = open("./tests/data/batched_nms_data.pkl", "rb") 36 | results = pickle.load(file) 37 | det_boxes = torch.from_numpy(results["boxes"]) 38 | 39 | socres = torch.rand(det_boxes.shape[0], 8) 40 | score_thr = 0.9 41 | max_num = 100 42 | nms_cfg = dict(iou_threshold=0.5) 43 | boxes, keep = multiclass_nms( 44 | det_boxes, socres, score_thr=score_thr, nms_cfg=nms_cfg, max_num=max_num 45 | ) 46 | 47 | assert boxes.shape[0] <= max_num 48 | assert keep.shape[0] <= max_num 49 | assert min(boxes[:, -1]) >= score_thr 50 | 51 | # test all zero score 52 | socres = torch.zeros(det_boxes.shape[0], 8) 53 | score_thr = 0.1 54 | nms_cfg = dict(iou_threshold=0.5) 55 | boxes, keep = multiclass_nms( 56 | det_boxes, socres, score_thr=score_thr, nms_cfg=nms_cfg, max_num=-1 57 | ) 58 | assert boxes.shape[0] == 0 59 | assert keep.shape[0] == 0 60 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_repvgg_conv.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pytest 3 | import torch 4 | 5 | from nanodet.model.module.conv import RepVGGConvModule 6 | 7 | 8 | def test_repvgg_conv(): 9 | # test activation type 10 | with pytest.raises(AssertionError): 11 | activation = dict(type="softmax") 12 | RepVGGConvModule(3, 5, 3, padding=1, activation=activation) 13 | 14 | # repvgg only support 3x3 conv 15 | with pytest.raises(AssertionError): 16 | RepVGGConvModule(3, 5, 2, activation="ReLU") 17 | 18 | with pytest.raises(AssertionError): 19 | RepVGGConvModule(3, 5, 3, padding=0, activation="ReLU") 20 | 21 | training_conv = RepVGGConvModule(3, 5, 3, deploy=False) 22 | assert hasattr(training_conv, "rbr_identity") 23 | assert hasattr(training_conv, "rbr_dense") 24 | assert hasattr(training_conv, "rbr_1x1") 25 | assert not hasattr(training_conv, "rbr_reparam") 26 | 27 | deploy_conv = RepVGGConvModule(3, 5, 3, deploy=True) 28 | assert not hasattr(deploy_conv, "rbr_identity") 29 | assert not hasattr(deploy_conv, "rbr_dense") 30 | assert not hasattr(deploy_conv, "rbr_1x1") 31 | assert hasattr(deploy_conv, "rbr_reparam") 32 | 33 | converted_weights = {} 34 | deploy_conv.load_state_dict(training_conv.state_dict(), strict=False) 35 | kernel, bias = training_conv.repvgg_convert() 36 | converted_weights["rbr_reparam.weight"] = kernel 37 | converted_weights["rbr_reparam.bias"] = bias 38 | for name, param in deploy_conv.named_parameters(): 39 | print("deploy param: ", name, param.size(), np.mean(converted_weights[name])) 40 | param.data = torch.from_numpy(converted_weights[name]).float() 41 | 42 | x = torch.rand(1, 3, 16, 16) 43 | train_out = training_conv(x) 44 | deploy_out = deploy_conv(x) 45 | assert train_out.shape == (1, 5, 16, 16) 46 | assert deploy_out.shape == (1, 5, 16, 16) 47 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_norm.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch.nn as nn 3 | 4 | from nanodet.model.module.norm import build_norm_layer 5 | 6 | 7 | def test_build_norm_layer(): 8 | with pytest.raises(AssertionError): 9 | # cfg must be a dict 10 | cfg = "BN" 11 | build_norm_layer(cfg, 3) 12 | 13 | with pytest.raises(AssertionError): 14 | # `type` must be in cfg 15 | cfg = dict() 16 | build_norm_layer(cfg, 3) 17 | 18 | with pytest.raises(KeyError): 19 | # unsupported norm type 20 | cfg = dict(type="FancyNorm") 21 | build_norm_layer(cfg, 3) 22 | 23 | with pytest.raises(AssertionError): 24 | # postfix must be int or str 25 | cfg = dict(type="BN") 26 | build_norm_layer(cfg, 3, postfix=[1, 2]) 27 | 28 | with pytest.raises(AssertionError): 29 | # `num_groups` must be in cfg when using 'GN' 30 | cfg = dict(type="GN") 31 | build_norm_layer(cfg, 3) 32 | 33 | # test each type of norm layer in norm_cfg 34 | abbr_mapping = { 35 | "BN": "bn", 36 | "SyncBN": "bn", 37 | "GN": "gn", 38 | } 39 | module_dict = { 40 | "BN": nn.BatchNorm2d, 41 | "SyncBN": nn.SyncBatchNorm, 42 | "GN": nn.GroupNorm, 43 | } 44 | for type_name, module in module_dict.items(): 45 | for postfix in ["_test", 1]: 46 | cfg = dict(type=type_name) 47 | if type_name == "GN": 48 | cfg["num_groups"] = 2 49 | name, layer = build_norm_layer(cfg, 4, postfix=postfix) 50 | assert name == abbr_mapping[type_name] + str(postfix) 51 | assert isinstance(layer, module) 52 | if type_name == "GN": 53 | assert layer.num_channels == 4 54 | assert layer.num_groups == cfg["num_groups"] 55 | elif type_name != "LN": 56 | assert layer.num_features == 4 57 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | .vscode 107 | .idea 108 | .DS_Store 109 | 110 | # custom 111 | *.pkl 112 | *.pkl.json 113 | *.log.json 114 | work_dirs/ 115 | 116 | # Pytorch 117 | *.pth 118 | *.py~ 119 | *.sh~ 120 | -------------------------------------------------------------------------------- /demo_openvino/nanodet_openvino.h: -------------------------------------------------------------------------------- 1 | // 2 | // Create by RangiLyu 3 | // 2021 / 1 / 12 4 | // 5 | 6 | #ifndef _NANODET_OPENVINO_H_ 7 | #define _NANODET_OPENVINO_H_ 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | 14 | typedef struct HeadInfo 15 | { 16 | std::string cls_layer; 17 | std::string dis_layer; 18 | int stride; 19 | } HeadInfo; 20 | 21 | struct CenterPrior 22 | { 23 | int x; 24 | int y; 25 | int stride; 26 | }; 27 | 28 | typedef struct BoxInfo 29 | { 30 | float x1; 31 | float y1; 32 | float x2; 33 | float y2; 34 | float score; 35 | int label; 36 | } BoxInfo; 37 | 38 | class NanoDet 39 | { 40 | public: 41 | NanoDet(const char* param); 42 | 43 | ~NanoDet(); 44 | 45 | InferenceEngine::ExecutableNetwork network_; 46 | InferenceEngine::InferRequest infer_request_; 47 | // static bool hasGPU; 48 | 49 | // modify these parameters to the same with your config if you want to use your own model 50 | int input_size[2] = {416, 416}; // input height and width 51 | int num_class = 80; // number of classes. 80 for COCO 52 | int reg_max = 7; // `reg_max` set in the training config. Default: 7. 53 | std::vector strides = { 8, 16, 32, 64 }; // strides of the multi-level feature. 54 | 55 | std::vector detect(cv::Mat image, float score_threshold, float nms_threshold); 56 | 57 | private: 58 | void preprocess(cv::Mat& image, InferenceEngine::Blob::Ptr& blob); 59 | void decode_infer(const float*& pred, std::vector& center_priors, float threshold, std::vector>& results); 60 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride); 61 | static void nms(std::vector& result, float nms_threshold); 62 | std::string input_name_ = "data"; 63 | std::string output_name_ = "output"; 64 | }; 65 | 66 | 67 | #endif //_NANODE_TOPENVINO_H_ 68 | -------------------------------------------------------------------------------- /nanodet/model/module/norm.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | norm_cfg = { 4 | # format: layer_type: (abbreviation, module) 5 | "BN": ("bn", nn.BatchNorm2d), 6 | "SyncBN": ("bn", nn.SyncBatchNorm), 7 | "GN": ("gn", nn.GroupNorm), 8 | # and potentially 'SN' 9 | } 10 | 11 | 12 | def build_norm_layer(cfg, num_features, postfix=""): 13 | """Build normalization layer 14 | 15 | Args: 16 | cfg (dict): cfg should contain: 17 | type (str): identify norm layer type. 18 | layer args: args needed to instantiate a norm layer. 19 | requires_grad (bool): [optional] whether stop gradient updates 20 | num_features (int): number of channels from input. 21 | postfix (int, str): appended into norm abbreviation to 22 | create named layer. 23 | 24 | Returns: 25 | name (str): abbreviation + postfix 26 | layer (nn.Module): created norm layer 27 | """ 28 | assert isinstance(cfg, dict) and "type" in cfg 29 | cfg_ = cfg.copy() 30 | 31 | layer_type = cfg_.pop("type") 32 | if layer_type not in norm_cfg: 33 | raise KeyError("Unrecognized norm type {}".format(layer_type)) 34 | else: 35 | abbr, norm_layer = norm_cfg[layer_type] 36 | if norm_layer is None: 37 | raise NotImplementedError 38 | 39 | assert isinstance(postfix, (int, str)) 40 | name = abbr + str(postfix) 41 | 42 | requires_grad = cfg_.pop("requires_grad", True) 43 | cfg_.setdefault("eps", 1e-5) 44 | if layer_type != "GN": 45 | layer = norm_layer(num_features, **cfg_) 46 | if layer_type == "SyncBN" and hasattr(layer, "_specify_ddp_gpu_num"): 47 | layer._specify_ddp_gpu_num(1) 48 | else: 49 | assert "num_groups" in cfg_ 50 | layer = norm_layer(num_channels=num_features, **cfg_) 51 | 52 | for param in layer.parameters(): 53 | param.requires_grad = requires_grad 54 | 55 | return name, layer 56 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_ghostnet.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import GhostNet, build_backbone 5 | 6 | 7 | def test_ghostnet(): 8 | with pytest.raises(AssertionError): 9 | cfg = dict(name="GhostNet", width_mult=1.0, out_stages=(11, 12), pretrain=False) 10 | build_backbone(cfg) 11 | 12 | input = torch.rand(1, 3, 64, 64) 13 | out_stages = [i for i in range(10)] 14 | model = GhostNet( 15 | width_mult=1.0, out_stages=out_stages, activation="ReLU6", pretrain=True 16 | ) 17 | output = model(input) 18 | 19 | assert output[0].shape == torch.Size([1, 16, 32, 32]) 20 | assert output[1].shape == torch.Size([1, 24, 16, 16]) 21 | assert output[2].shape == torch.Size([1, 24, 16, 16]) 22 | assert output[3].shape == torch.Size([1, 40, 8, 8]) 23 | assert output[4].shape == torch.Size([1, 40, 8, 8]) 24 | assert output[5].shape == torch.Size([1, 80, 4, 4]) 25 | assert output[6].shape == torch.Size([1, 112, 4, 4]) 26 | assert output[7].shape == torch.Size([1, 160, 2, 2]) 27 | assert output[8].shape == torch.Size([1, 160, 2, 2]) 28 | assert output[9].shape == torch.Size([1, 960, 2, 2]) 29 | 30 | model = GhostNet( 31 | width_mult=0.75, out_stages=out_stages, activation="LeakyReLU", pretrain=False 32 | ) 33 | output = model(input) 34 | 35 | assert output[0].shape == torch.Size([1, 12, 32, 32]) 36 | assert output[1].shape == torch.Size([1, 20, 16, 16]) 37 | assert output[2].shape == torch.Size([1, 20, 16, 16]) 38 | assert output[3].shape == torch.Size([1, 32, 8, 8]) 39 | assert output[4].shape == torch.Size([1, 32, 8, 8]) 40 | assert output[5].shape == torch.Size([1, 60, 4, 4]) 41 | assert output[6].shape == torch.Size([1, 84, 4, 4]) 42 | assert output[7].shape == torch.Size([1, 120, 2, 2]) 43 | assert output[8].shape == torch.Size([1, 120, 2, 2]) 44 | assert output[9].shape == torch.Size([1, 720, 2, 2]) 45 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_init_weights.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | import torch.nn as nn 4 | 5 | from nanodet.model.module.init_weights import ( 6 | constant_init, 7 | kaiming_init, 8 | normal_init, 9 | xavier_init, 10 | ) 11 | 12 | 13 | def test_constant_init(): 14 | conv_module = nn.Conv2d(3, 16, 3) 15 | constant_init(conv_module, 0.1) 16 | assert conv_module.weight.allclose(torch.full_like(conv_module.weight, 0.1)) 17 | assert conv_module.bias.allclose(torch.zeros_like(conv_module.bias)) 18 | conv_module_no_bias = nn.Conv2d(3, 16, 3, bias=False) 19 | constant_init(conv_module_no_bias, 0.1) 20 | assert conv_module.weight.allclose(torch.full_like(conv_module.weight, 0.1)) 21 | 22 | 23 | def test_xavier_init(): 24 | conv_module = nn.Conv2d(3, 16, 3) 25 | xavier_init(conv_module, bias=0.1) 26 | assert conv_module.bias.allclose(torch.full_like(conv_module.bias, 0.1)) 27 | xavier_init(conv_module, distribution="uniform") 28 | with pytest.raises(AssertionError): 29 | xavier_init(conv_module, distribution="student-t") 30 | conv_module_no_bias = nn.Conv2d(3, 16, 3, bias=False) 31 | xavier_init(conv_module_no_bias) 32 | 33 | 34 | def test_normal_init(): 35 | conv_module = nn.Conv2d(3, 16, 3) 36 | normal_init(conv_module, bias=0.1) 37 | assert conv_module.bias.allclose(torch.full_like(conv_module.bias, 0.1)) 38 | conv_module_no_bias = nn.Conv2d(3, 16, 3, bias=False) 39 | normal_init(conv_module_no_bias) 40 | 41 | 42 | def test_kaiming_init(): 43 | conv_module = nn.Conv2d(3, 16, 3) 44 | kaiming_init(conv_module, bias=0.1) 45 | assert conv_module.bias.allclose(torch.full_like(conv_module.bias, 0.1)) 46 | kaiming_init(conv_module, distribution="uniform") 47 | with pytest.raises(AssertionError): 48 | kaiming_init(conv_module, distribution="student-t") 49 | conv_module_no_bias = nn.Conv2d(3, 16, 3, bias=False) 50 | kaiming_init(conv_module_no_bias) 51 | -------------------------------------------------------------------------------- /tests/test_models/test_head/test_nanodet_head.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from nanodet.model.head import build_head 4 | from nanodet.util.yacs import CfgNode 5 | 6 | 7 | def test_gfl_head_loss(): 8 | head_cfg = dict( 9 | name="NanoDetHead", 10 | num_classes=80, 11 | input_channel=1, 12 | feat_channels=96, 13 | stacked_convs=2, 14 | conv_type="DWConv", 15 | reg_max=8, 16 | strides=[8, 16, 32], 17 | loss=dict( 18 | loss_qfl=dict( 19 | name="QualityFocalLoss", use_sigmoid=True, beta=2.0, loss_weight=1.0 20 | ), 21 | loss_dfl=dict(name="DistributionFocalLoss", loss_weight=0.25), 22 | loss_bbox=dict(name="GIoULoss", loss_weight=2.0), 23 | ), 24 | ) 25 | cfg = CfgNode(head_cfg) 26 | 27 | head = build_head(cfg) 28 | feat = [torch.rand(1, 1, 320 // stride, 320 // stride) for stride in [8, 16, 32]] 29 | 30 | preds = head.forward(feat) 31 | num_points = sum([(320 // stride) ** 2 for stride in [8, 16, 32]]) 32 | assert preds.shape == (1, num_points, 80 + (8 + 1) * 4) 33 | 34 | head_cfg = dict( 35 | name="NanoDetHead", 36 | num_classes=20, 37 | input_channel=1, 38 | feat_channels=96, 39 | stacked_convs=2, 40 | conv_type="Conv", 41 | reg_max=5, 42 | share_cls_reg=False, 43 | strides=[8, 16, 32], 44 | loss=dict( 45 | loss_qfl=dict( 46 | name="QualityFocalLoss", use_sigmoid=True, beta=2.0, loss_weight=1.0 47 | ), 48 | loss_dfl=dict(name="DistributionFocalLoss", loss_weight=0.25), 49 | loss_bbox=dict(name="GIoULoss", loss_weight=2.0), 50 | ), 51 | ) 52 | cfg = CfgNode(head_cfg) 53 | head = build_head(cfg) 54 | 55 | preds = head.forward(feat) 56 | num_points = sum([(320 // stride) ** 2 for stride in [8, 16, 32]]) 57 | assert preds.shape == (1, num_points, 20 + (5 + 1) * 4) 58 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/Box.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.graphics.Color; 4 | import android.graphics.RectF; 5 | 6 | import java.util.Random; 7 | 8 | public class Box { 9 | public float x0,y0,x1,y1; 10 | private int label; 11 | private float score; 12 | private static String[] labels={"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 13 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 14 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 15 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 16 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 17 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 18 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 19 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 20 | "hair drier", "toothbrush"}; 21 | public Box(float x0,float y0, float x1, float y1, int label, float score){ 22 | this.x0 = x0; 23 | this.y0 = y0; 24 | this.x1 = x1; 25 | this.y1 = y1; 26 | this.label = label; 27 | this.score = score; 28 | } 29 | 30 | public RectF getRect(){ 31 | return new RectF(x0,y0,x1,y1); 32 | } 33 | 34 | public String getLabel(){ 35 | return labels[label]; 36 | } 37 | 38 | public float getScore(){ 39 | return score; 40 | } 41 | 42 | public int getColor(){ 43 | Random random = new Random(label); 44 | return Color.argb(255,random.nextInt(256),random.nextInt(256),random.nextInt(256)); 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /demo_android_ncnn/README.md: -------------------------------------------------------------------------------- 1 | # NanoDet NCNN Android Demo 2 | 3 | This repo is an Android object detection demo of NanoDet using 4 | [Tencent's NCNN framework](https://github.com/Tencent/ncnn). 5 | 6 | # Tutorial 7 | 8 | ## Step1. 9 | Download ncnn-android-vulkan.zip from ncnn repo or build ncnn-android from source. 10 | 11 | - [ncnn-android-vulkan.zip download link](https://github.com/Tencent/ncnn/releases) 12 | 13 | ## Step2. 14 | Unzip ncnn-android-vulkan.zip into demo_android_ncnn/app/src/main/cpp or change the ncnn_DIR path to yours in demo_android_ncnn/app/src/main/cpp/CMakeLists.txt 15 | 16 | ```bash 17 | # e.g. change to `ncnn-20211208-android-vulkan` if download version 200211208 18 | set(ncnn_DIR ${CMAKE_SOURCE_DIR}/ncnn-20211208-android-vulkan/${ANDROID_ABI}/lib/cmake/ncnn) 19 | ``` 20 | 21 | ## Step3. 22 | Copy the NanoDet ncnn model file and rename to nanodet.param and nanodet.bin from models folder into demo_android_ncnn/app/src/main/assets 23 | 24 | * [NanoDet ncnn model download link](https://drive.google.com/file/d/1cuVBJiFKwyq1-l3AwHoP2boTesUQP-6K/view?usp=sharing) 25 | 26 | If you want to run yolov4-tiny and yolov5s, download them and also put in demo_android_ncnn/app/src/main/assets. 27 | 28 | * [Yolov4 and v5 ncnn model download link](https://drive.google.com/file/d/1Qk_1fDvOcFmNppDnaMFW-xFpMgLDyeAs/view?usp=sharing) 29 | 30 | ## Step4. 31 | Open demo_android_ncnn folder with Android Studio and then build it. 32 | 33 | # Screenshot 34 | ![](Android_demo.jpg) 35 | 36 | # Notice 37 | 38 | * The FPS in the app includes pre-process, post-process and visualization, not equal to the model inference time. 39 | 40 | * If meet error like `No version of NDK matched the requested version`, set `android { ndkVersion` to your ndk version. 41 | 42 | * If you want to use custom model, remember to change the hyperparams in `demo_android_ncnn/app/src/main/cpp/NanoDet.h` the same with your training config. 43 | 44 | # Reference 45 | 46 | * [ncnn](https://github.com/tencent/ncnn) 47 | * [YOLOv5_NCNN](https://github.com/WZTENG/YOLOv5_NCNN) 48 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/cpp/YoloV4.h: -------------------------------------------------------------------------------- 1 | #ifndef YOLOV4_H 2 | #define YOLOV4_H 3 | 4 | #include "net.h" 5 | #include "YoloV5.h" 6 | 7 | 8 | class YoloV4 { 9 | public: 10 | YoloV4(AAssetManager *mgr, const char *param, const char *bin, bool useGPU); 11 | 12 | ~YoloV4(); 13 | 14 | std::vector detect(JNIEnv *env, jobject image, float threshold, float nms_threshold); 15 | std::vector labels{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 16 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 17 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 18 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 19 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 20 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 21 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 22 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 23 | "hair drier", "toothbrush"}; 24 | private: 25 | static std::vector 26 | decode_infer(ncnn::Mat &data, const yolocv::YoloSize &frame_size, int net_size, int num_classes, float threshold); 27 | 28 | // static void nms(std::vector& result,float nms_threshold); 29 | ncnn::Net *Net; 30 | int input_size = 640 / 2; 31 | int num_class = 80; 32 | public: 33 | static YoloV4 *detector; 34 | static bool hasGPU; 35 | }; 36 | 37 | 38 | #endif //YOLOV4_H 39 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-v24/ic_launcher_foreground.xml: -------------------------------------------------------------------------------- 1 | 7 | 12 | 13 | 19 | 22 | 25 | 26 | 27 | 28 | 34 | 35 | -------------------------------------------------------------------------------- /tests/test_data/test_dataset/test_cocodataset.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from nanodet.data.dataset import CocoDataset, build_dataset 4 | 5 | 6 | def test_cocodataset(): 7 | cfg = dict( 8 | name="CocoDataset", 9 | img_path="./tests/data", 10 | ann_path="./tests/data/dummy_coco.json", 11 | input_size=[320, 320], # [w,h] 12 | keep_ratio=True, 13 | use_instance_mask=True, 14 | pipeline=dict(normalize=[[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]), 15 | ) 16 | dataset = build_dataset(cfg, "train") 17 | assert isinstance(dataset, CocoDataset) 18 | 19 | for i, data in enumerate(dataset): 20 | assert data["img"].shape == (3, 320, 285) 21 | for mask in data["gt_masks"]: 22 | assert mask.shape == (320, 285) 23 | 24 | dataset = build_dataset(cfg, "val") 25 | for i, data in enumerate(dataset): 26 | assert data["img"].shape == (3, 320, 285) 27 | for mask in data["gt_masks"]: 28 | assert mask.shape == (320, 285) 29 | 30 | cfg["keep_ratio"] = False 31 | dataset = build_dataset(cfg, "train") 32 | for i, data in enumerate(dataset): 33 | assert data["img"].shape == (3, 320, 320) 34 | for mask in data["gt_masks"]: 35 | assert mask.shape == (320, 320) 36 | 37 | with pytest.raises(AssertionError): 38 | build_dataset(cfg, "2333") 39 | 40 | 41 | def test_multi_scale(): 42 | cfg = dict( 43 | name="CocoDataset", 44 | img_path="./tests/data", 45 | ann_path="./tests/data/dummy_coco.json", 46 | input_size=[320, 320], # [w,h] 47 | multi_scale=[1.5, 1.5], 48 | keep_ratio=True, 49 | use_instance_mask=True, 50 | pipeline=dict(normalize=[[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]), 51 | ) 52 | dataset = build_dataset(cfg, "train") 53 | 54 | for i, data in enumerate(dataset): 55 | assert data["img"].shape == (3, 480, 427) 56 | for mask in data["gt_masks"]: 57 | assert mask.shape == (480, 427) 58 | -------------------------------------------------------------------------------- /tools/flops.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import argparse 16 | 17 | import torch 18 | 19 | from nanodet.model.arch import build_model 20 | from nanodet.util import cfg, load_config 21 | 22 | 23 | def main(config, input_shape=(320, 320)): 24 | model = build_model(config.model) 25 | try: 26 | import mobile_cv.lut.lib.pt.flops_utils as flops_utils 27 | except ImportError: 28 | print("mobile-cv is not installed. Skip flops calculation.") 29 | return 30 | first_batch = torch.rand((1, 3, input_shape[0], input_shape[1])) 31 | input_args = (first_batch,) 32 | flops_utils.print_model_flops(model, input_args) 33 | 34 | 35 | def parse_args(): 36 | parser = argparse.ArgumentParser( 37 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, 38 | description="Convert .pth model to onnx.", 39 | ) 40 | parser.add_argument("cfg", type=str, help="Path to .yml config file.") 41 | parser.add_argument( 42 | "--input_shape", type=str, default=None, help="Model intput shape." 43 | ) 44 | return parser.parse_args() 45 | 46 | 47 | if __name__ == "__main__": 48 | args = parse_args() 49 | 50 | cfg_path = args.cfg 51 | load_config(cfg, cfg_path) 52 | 53 | input_shape = args.input_shape 54 | if input_shape is None: 55 | input_shape = cfg.data.train.input_size 56 | else: 57 | input_shape = tuple(map(int, input_shape.split(","))) 58 | assert len(input_shape) == 2 59 | main(config=cfg, input_shape=input_shape) 60 | -------------------------------------------------------------------------------- /tests/test_utils/test_env_utils.py: -------------------------------------------------------------------------------- 1 | import multiprocessing as mp 2 | import os 3 | import platform 4 | 5 | import cv2 6 | 7 | from nanodet.util.env_utils import set_multi_processing 8 | 9 | 10 | def test_setup_multi_processes(): 11 | # temp save system setting 12 | sys_start_mehod = mp.get_start_method(allow_none=True) 13 | sys_cv_threads = cv2.getNumThreads() 14 | # pop and temp save system env vars 15 | sys_omp_threads = os.environ.pop("OMP_NUM_THREADS", default=None) 16 | sys_mkl_threads = os.environ.pop("MKL_NUM_THREADS", default=None) 17 | 18 | # test distributed 19 | set_multi_processing(distributed=True) 20 | assert os.getenv("OMP_NUM_THREADS") == "1" 21 | assert os.getenv("MKL_NUM_THREADS") == "1" 22 | # when set to 0, the num threads will be 1 23 | assert cv2.getNumThreads() == 1 24 | if platform.system() != "Windows": 25 | assert mp.get_start_method() == "fork" 26 | 27 | # test num workers <= 1 28 | os.environ.pop("OMP_NUM_THREADS") 29 | os.environ.pop("MKL_NUM_THREADS") 30 | set_multi_processing(distributed=False) 31 | assert "OMP_NUM_THREADS" not in os.environ 32 | assert "MKL_NUM_THREADS" not in os.environ 33 | 34 | # test manually set env var 35 | os.environ["OMP_NUM_THREADS"] = "4" 36 | set_multi_processing(distributed=False) 37 | assert os.getenv("OMP_NUM_THREADS") == "4" 38 | 39 | # test manually set opencv threads and mp start method 40 | config = dict(mp_start_method="spawn", opencv_num_threads=4, distributed=True) 41 | set_multi_processing(**config) 42 | assert cv2.getNumThreads() == 4 43 | assert mp.get_start_method() == "spawn" 44 | 45 | # revert setting to avoid affecting other programs 46 | if sys_start_mehod: 47 | mp.set_start_method(sys_start_mehod, force=True) 48 | cv2.setNumThreads(sys_cv_threads) 49 | if sys_omp_threads: 50 | os.environ["OMP_NUM_THREADS"] = sys_omp_threads 51 | else: 52 | os.environ.pop("OMP_NUM_THREADS") 53 | if sys_mkl_threads: 54 | os.environ["MKL_NUM_THREADS"] = sys_mkl_threads 55 | else: 56 | os.environ.pop("MKL_NUM_THREADS") 57 | -------------------------------------------------------------------------------- /nanodet/model/arch/nanodet_plus.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import copy 16 | 17 | import torch 18 | 19 | from ..head import build_head 20 | from .one_stage_detector import OneStageDetector 21 | 22 | 23 | class NanoDetPlus(OneStageDetector): 24 | def __init__( 25 | self, 26 | backbone, 27 | fpn, 28 | aux_head, 29 | head, 30 | detach_epoch=0, 31 | ): 32 | super(NanoDetPlus, self).__init__( 33 | backbone_cfg=backbone, fpn_cfg=fpn, head_cfg=head 34 | ) 35 | self.aux_fpn = copy.deepcopy(self.fpn) 36 | self.aux_head = build_head(aux_head) 37 | self.detach_epoch = detach_epoch 38 | 39 | def forward_train(self, gt_meta): 40 | img = gt_meta["img"] 41 | feat = self.backbone(img) 42 | fpn_feat = self.fpn(feat) 43 | if self.epoch >= self.detach_epoch: 44 | aux_fpn_feat = self.aux_fpn([f.detach() for f in feat]) 45 | dual_fpn_feat = ( 46 | torch.cat([f.detach(), aux_f], dim=1) 47 | for f, aux_f in zip(fpn_feat, aux_fpn_feat) 48 | ) 49 | else: 50 | aux_fpn_feat = self.aux_fpn(feat) 51 | dual_fpn_feat = ( 52 | torch.cat([f, aux_f], dim=1) for f, aux_f in zip(fpn_feat, aux_fpn_feat) 53 | ) 54 | head_out = self.head(fpn_feat) 55 | aux_head_out = self.aux_head(dual_fpn_feat) 56 | loss, loss_states = self.head.loss(head_out, gt_meta, aux_preds=aux_head_out) 57 | return head_out, loss, loss_states 58 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/NcnnApp.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.app.Application; 4 | import android.content.Context; 5 | import android.util.Log; 6 | 7 | import androidx.multidex.MultiDex; 8 | 9 | import com.zxy.recovery.callback.RecoveryCallback; 10 | import com.zxy.recovery.core.Recovery; 11 | 12 | 13 | public class NcnnApp extends Application { 14 | 15 | @Override 16 | public void onCreate() { 17 | super.onCreate(); 18 | 19 | //崩溃界面 20 | initRecovery(); 21 | } 22 | 23 | @Override 24 | protected void attachBaseContext(Context base) { 25 | super.attachBaseContext(base); 26 | MultiDex.install(base); 27 | } 28 | 29 | private void initRecovery() { 30 | Recovery.getInstance() 31 | .debug(BuildConfig.DEBUG) 32 | .recoverInBackground(true) 33 | .recoverStack(true) 34 | .mainPage(MainActivity.class) 35 | .recoverEnabled(true) 36 | .callback(new MyCrashCallback()) 37 | .silent(false, Recovery.SilentMode.RECOVER_ACTIVITY_STACK) 38 | // .skip(TestActivity.class) 39 | .init(this); 40 | AppCrashHandler.register(); 41 | } 42 | 43 | static final class MyCrashCallback implements RecoveryCallback { 44 | @Override 45 | public void stackTrace(String exceptionMessage) { 46 | Log.e("wzt", "exceptionMessage:" + exceptionMessage); 47 | } 48 | 49 | @Override 50 | public void cause(String cause) { 51 | Log.e("wzt", "cause:" + cause); 52 | } 53 | 54 | @Override 55 | public void exception(String exceptionType, String throwClassName, String throwMethodName, int throwLineNumber) { 56 | Log.e("wzt", "exceptionClassName:" + exceptionType); 57 | Log.e("wzt", "throwClassName:" + throwClassName); 58 | Log.e("wzt", "throwMethodName:" + throwMethodName); 59 | Log.e("wzt", "throwLineNumber:" + throwLineNumber); 60 | } 61 | 62 | @Override 63 | public void throwable(Throwable throwable) { 64 | 65 | } 66 | } 67 | 68 | } 69 | -------------------------------------------------------------------------------- /nanodet/data/transform/pipeline.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import functools 16 | import warnings 17 | from typing import Dict, Tuple 18 | 19 | from torch.utils.data import Dataset 20 | 21 | from .color import color_aug_and_norm 22 | from .warp import ShapeTransform, warp_and_resize 23 | 24 | 25 | class LegacyPipeline: 26 | def __init__(self, cfg, keep_ratio): 27 | warnings.warn( 28 | "Deprecated warning! Pipeline from nanodet v0.x has been deprecated," 29 | "Please use new Pipeline and update your config!" 30 | ) 31 | self.warp = functools.partial( 32 | warp_and_resize, warp_kwargs=cfg, keep_ratio=keep_ratio 33 | ) 34 | self.color = functools.partial(color_aug_and_norm, kwargs=cfg) 35 | 36 | def __call__(self, meta, dst_shape): 37 | meta = self.warp(meta, dst_shape=dst_shape) 38 | meta = self.color(meta=meta) 39 | return meta 40 | 41 | 42 | class Pipeline: 43 | """Data process pipeline. Apply augmentation and pre-processing on 44 | meta_data from dataset. 45 | 46 | Args: 47 | cfg (Dict): Data pipeline config. 48 | keep_ratio (bool): Whether to keep aspect ratio when resizing image. 49 | 50 | """ 51 | 52 | def __init__(self, cfg: Dict, keep_ratio: bool): 53 | self.shape_transform = ShapeTransform(keep_ratio, **cfg) 54 | self.color = functools.partial(color_aug_and_norm, kwargs=cfg) 55 | 56 | def __call__(self, dataset: Dataset, meta: Dict, dst_shape: Tuple[int, int]): 57 | meta = self.shape_transform(meta, dst_shape=dst_shape) 58 | meta = self.color(meta=meta) 59 | return meta 60 | -------------------------------------------------------------------------------- /tests/data/dummy_coco.json: -------------------------------------------------------------------------------- 1 | { 2 | "images": [ 3 | { 4 | "file_name": "test_img.jpg", 5 | "height": 404, 6 | "width": 360, 7 | "id": 0 8 | }, 9 | { 10 | "file_name": "test_img.jpg", 11 | "height": 404, 12 | "width": 360, 13 | "id": 1 14 | }, 15 | { 16 | "file_name": "test_img.jpg", 17 | "height": 404, 18 | "width": 360, 19 | "id": 2 20 | } 21 | ], 22 | "annotations": [ 23 | { 24 | "segmentation": [ 25 | [0,0,0,20,20,20,20,0] 26 | ], 27 | "bbox": [ 28 | 0, 29 | 0, 30 | 20, 31 | 20 32 | ], 33 | "area": 400.00, 34 | "score": 1.0, 35 | "category_id": 1, 36 | "id": 1, 37 | "iscrowd": 0, 38 | "image_id": 0 39 | }, 40 | { 41 | "segmentation": [ 42 | [0,0,0,20,20,20,20,0] 43 | ], 44 | "bbox": [ 45 | 0, 46 | 0, 47 | 20, 48 | 20 49 | ], 50 | "area": 400.00, 51 | "score": 1.0, 52 | "category_id": 2, 53 | "id": 2, 54 | "iscrowd": 0, 55 | "image_id": 0 56 | }, 57 | { 58 | "segmentation": [ 59 | [0,0,0,20,20,20,20,0] 60 | ], 61 | "bbox": [ 62 | 0, 63 | 0, 64 | 20, 65 | 20 66 | ], 67 | "area": 400.00, 68 | "score": 1.0, 69 | "category_id": 1, 70 | "id": 3, 71 | "iscrowd": 0, 72 | "image_id": 1 73 | } 74 | ], 75 | "categories": [ 76 | { 77 | "id": 1, 78 | "name": "bus", 79 | "supercategory": "none" 80 | }, 81 | { 82 | "id": 2, 83 | "name": "car", 84 | "supercategory": "none" 85 | } 86 | ], 87 | "licenses": [], 88 | "info": null 89 | } 90 | -------------------------------------------------------------------------------- /nanodet/model/backbone/timm_wrapper.py: -------------------------------------------------------------------------------- 1 | # Copyright 2022 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import logging 16 | 17 | import torch.nn as nn 18 | 19 | logger = logging.getLogger("NanoDet") 20 | 21 | 22 | class TIMMWrapper(nn.Module): 23 | """Wrapper to use backbones in timm 24 | https://github.com/rwightman/pytorch-image-models.""" 25 | 26 | def __init__( 27 | self, 28 | model_name, 29 | features_only=True, 30 | pretrained=True, 31 | checkpoint_path="", 32 | in_channels=3, 33 | **kwargs, 34 | ): 35 | try: 36 | import timm 37 | except ImportError as exc: 38 | raise RuntimeError( 39 | "timm is not installed, please install it first" 40 | ) from exc 41 | super(TIMMWrapper, self).__init__() 42 | self.timm = timm.create_model( 43 | model_name=model_name, 44 | features_only=features_only, 45 | pretrained=pretrained, 46 | in_chans=in_channels, 47 | checkpoint_path=checkpoint_path, 48 | **kwargs, 49 | ) 50 | 51 | # Remove unused layers 52 | self.timm.global_pool = None 53 | self.timm.fc = None 54 | self.timm.classifier = None 55 | 56 | feature_info = getattr(self.timm, "feature_info", None) 57 | if feature_info: 58 | logger.info(f"TIMM backbone feature channels: {feature_info.channels()}") 59 | 60 | def forward(self, x): 61 | outs = self.timm(x) 62 | if isinstance(outs, (list, tuple)): 63 | features = tuple(outs) 64 | else: 65 | features = (outs,) 66 | return features 67 | -------------------------------------------------------------------------------- /nanodet/data/transform/color.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import random 16 | 17 | import cv2 18 | import numpy as np 19 | 20 | 21 | def random_brightness(img, delta): 22 | img += random.uniform(-delta, delta) 23 | return img 24 | 25 | 26 | def random_contrast(img, alpha_low, alpha_up): 27 | img *= random.uniform(alpha_low, alpha_up) 28 | return img 29 | 30 | 31 | def random_saturation(img, alpha_low, alpha_up): 32 | hsv_img = cv2.cvtColor(img.astype(np.float32), cv2.COLOR_BGR2HSV) 33 | hsv_img[..., 1] *= random.uniform(alpha_low, alpha_up) 34 | img = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR) 35 | return img 36 | 37 | 38 | def normalize(meta, mean, std): 39 | img = meta["img"].astype(np.float32) 40 | mean = np.array(mean, dtype=np.float64).reshape(1, -1) 41 | stdinv = 1 / np.array(std, dtype=np.float64).reshape(1, -1) 42 | cv2.subtract(img, mean, img) 43 | cv2.multiply(img, stdinv, img) 44 | meta["img"] = img 45 | return meta 46 | 47 | 48 | def _normalize(img, mean, std): 49 | mean = np.array(mean, dtype=np.float32).reshape(1, 1, 3) / 255 50 | std = np.array(std, dtype=np.float32).reshape(1, 1, 3) / 255 51 | img = (img - mean) / std 52 | return img 53 | 54 | 55 | def color_aug_and_norm(meta, kwargs): 56 | img = meta["img"].astype(np.float32) / 255 57 | 58 | if "brightness" in kwargs and random.randint(0, 1): 59 | img = random_brightness(img, kwargs["brightness"]) 60 | 61 | if "contrast" in kwargs and random.randint(0, 1): 62 | img = random_contrast(img, *kwargs["contrast"]) 63 | 64 | if "saturation" in kwargs and random.randint(0, 1): 65 | img = random_saturation(img, *kwargs["saturation"]) 66 | # cv2.imshow('trans', img) 67 | # cv2.waitKey(0) 68 | img = _normalize(img, *kwargs["normalize"]) 69 | meta["img"] = img 70 | return meta 71 | -------------------------------------------------------------------------------- /tests/test_models/test_fpn/test_ghost_pan.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.fpn.ghost_pan import GhostPAN 5 | 6 | 7 | def test_ghost_pan(): 8 | """Tests GhostPAN.""" 9 | s = 64 10 | in_channels = [8, 16, 32, 64] 11 | feat_sizes = [s // 2**i for i in range(4)] # [64, 32, 16, 8] 12 | out_channels = 8 13 | # `num_extra_level` >= 0 14 | with pytest.raises(AssertionError): 15 | GhostPAN(in_channels=in_channels, out_channels=out_channels, num_extra_level=-1) 16 | 17 | # `num_blocks` > 0 18 | with pytest.raises(AssertionError): 19 | GhostPAN(in_channels=in_channels, out_channels=out_channels, num_blocks=0) 20 | 21 | pan_model = GhostPAN( 22 | in_channels=in_channels, 23 | out_channels=out_channels, 24 | num_extra_level=1, 25 | num_blocks=1, 26 | ) 27 | 28 | # PAN expects a multiple levels of features per image 29 | feats = [ 30 | torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i]) 31 | for i in range(len(in_channels)) 32 | ] 33 | outs = pan_model(feats) 34 | assert len(outs) == len(in_channels) + 1 35 | for i in range(len(in_channels)): 36 | assert outs[i].shape[1] == out_channels 37 | assert outs[i].shape[2] == outs[i].shape[3] == s // (2**i) 38 | assert outs[-1].shape[1] == out_channels 39 | assert outs[-1].shape[2] == outs[-1].shape[3] == s // 2 ** len(in_channels) 40 | 41 | # test non-default values 42 | pan_model = GhostPAN( 43 | in_channels=in_channels, 44 | out_channels=out_channels, 45 | num_extra_level=1, 46 | num_blocks=1, 47 | use_depthwise=True, 48 | expand=2, 49 | use_res=True, 50 | ) 51 | # PAN expects a multiple levels of features per image 52 | assert ( 53 | pan_model.downsamples[0].depthwise.groups 54 | == pan_model.downsamples[0].depthwise.in_channels 55 | ) 56 | feats = [ 57 | torch.rand(1, in_channels[i], feat_sizes[i], feat_sizes[i]) 58 | for i in range(len(in_channels)) 59 | ] 60 | outs = pan_model(feats) 61 | assert len(outs) == len(in_channels) + 1 62 | for i in range(len(in_channels)): 63 | assert outs[i].shape[1] == out_channels 64 | assert outs[i].shape[2] == outs[i].shape[3] == s // (2**i) 65 | assert outs[-1].shape[1] == out_channels 66 | assert outs[-1].shape[2] == outs[-1].shape[3] == s // 2 ** len(in_channels) 67 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/build.gradle: -------------------------------------------------------------------------------- 1 | apply plugin: 'com.android.application' 2 | 3 | android { 4 | compileSdkVersion 29 5 | buildToolsVersion "29.0.3" 6 | defaultConfig { 7 | applicationId "com.rangi.nanodet" 8 | minSdkVersion 26 9 | targetSdkVersion 29 10 | versionCode 1 11 | versionName "1.0" 12 | testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner" 13 | externalNativeBuild { 14 | cmake { 15 | cppFlags "" 16 | arguments '-DANDROID_PLATFORM=android-24', '-DANDROID_STL=c++_static', '-DANDROID_STL=c++_shared' 17 | } 18 | } 19 | 20 | ndk { 21 | moduleName "NcnnJniLog" 22 | ldLibs "log", "z", "m" 23 | abiFilters "armeabi-v7a", "arm64-v8a" 24 | } 25 | 26 | multiDexEnabled true 27 | } 28 | buildTypes { 29 | release { 30 | minifyEnabled false 31 | proguardFiles getDefaultProguardFile('proguard-android-optimize.txt'), 'proguard-rules.pro' 32 | } 33 | } 34 | externalNativeBuild { 35 | cmake { 36 | path "src/main/cpp/CMakeLists.txt" 37 | version "3.10.2" 38 | } 39 | } 40 | sourceSets { 41 | main { 42 | jniLibs.srcDirs = ['libs'] 43 | } 44 | } 45 | 46 | repositories { 47 | flatDir { 48 | dirs 'libs' 49 | } 50 | } 51 | } 52 | 53 | dependencies { 54 | implementation fileTree(dir: 'libs', include: ['*.jar']) 55 | implementation 'androidx.appcompat:appcompat:1.1.0' 56 | implementation 'androidx.constraintlayout:constraintlayout:1.1.3' 57 | testImplementation 'junit:junit:4.12' 58 | androidTestImplementation 'androidx.test.ext:junit:1.1.1' 59 | androidTestImplementation 'androidx.test.espresso:espresso-core:3.2.0' 60 | 61 | // Use the most recent version of CameraX, currently that is alpha04 62 | def camerax_version = "1.0.0-alpha05" 63 | //noinspection GradleDependency 64 | implementation "androidx.camera:camera-core:${camerax_version}" 65 | //noinspection GradleDependency 66 | implementation "androidx.camera:camera-camera2:${camerax_version}" 67 | 68 | implementation 'com.android.support:multidex:1.0.3' 69 | // crash 70 | implementation 'com.zxy.android:recovery:1.0.0' 71 | // photoview 72 | implementation 'com.github.chrisbanes:PhotoView:2.3.0' 73 | // implementation 'com.bm.photoview:library:1.4.1' 74 | 75 | } 76 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradlew.bat: -------------------------------------------------------------------------------- 1 | @if "%DEBUG%" == "" @echo off 2 | @rem ########################################################################## 3 | @rem 4 | @rem Gradle startup script for Windows 5 | @rem 6 | @rem ########################################################################## 7 | 8 | @rem Set local scope for the variables with windows NT shell 9 | if "%OS%"=="Windows_NT" setlocal 10 | 11 | set DIRNAME=%~dp0 12 | if "%DIRNAME%" == "" set DIRNAME=. 13 | set APP_BASE_NAME=%~n0 14 | set APP_HOME=%DIRNAME% 15 | 16 | @rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. 17 | set DEFAULT_JVM_OPTS= 18 | 19 | @rem Find java.exe 20 | if defined JAVA_HOME goto findJavaFromJavaHome 21 | 22 | set JAVA_EXE=java.exe 23 | %JAVA_EXE% -version >NUL 2>&1 24 | if "%ERRORLEVEL%" == "0" goto init 25 | 26 | echo. 27 | echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 28 | echo. 29 | echo Please set the JAVA_HOME variable in your environment to match the 30 | echo location of your Java installation. 31 | 32 | goto fail 33 | 34 | :findJavaFromJavaHome 35 | set JAVA_HOME=%JAVA_HOME:"=% 36 | set JAVA_EXE=%JAVA_HOME%/bin/java.exe 37 | 38 | if exist "%JAVA_EXE%" goto init 39 | 40 | echo. 41 | echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 42 | echo. 43 | echo Please set the JAVA_HOME variable in your environment to match the 44 | echo location of your Java installation. 45 | 46 | goto fail 47 | 48 | :init 49 | @rem Get command-line arguments, handling Windows variants 50 | 51 | if not "%OS%" == "Windows_NT" goto win9xME_args 52 | 53 | :win9xME_args 54 | @rem Slurp the command line arguments. 55 | set CMD_LINE_ARGS= 56 | set _SKIP=2 57 | 58 | :win9xME_args_slurp 59 | if "x%~1" == "x" goto execute 60 | 61 | set CMD_LINE_ARGS=%* 62 | 63 | :execute 64 | @rem Setup the command line 65 | 66 | set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar 67 | 68 | @rem Execute Gradle 69 | "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS% 70 | 71 | :end 72 | @rem End local scope for the variables with windows NT shell 73 | if "%ERRORLEVEL%"=="0" goto mainEnd 74 | 75 | :fail 76 | rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of 77 | rem the _cmd.exe /c_ return code! 78 | if not "" == "%GRADLE_EXIT_CONSOLE%" exit 1 79 | exit /b 1 80 | 81 | :mainEnd 82 | if "%OS%"=="Windows_NT" endlocal 83 | 84 | :omega 85 | -------------------------------------------------------------------------------- /tests/test_trainer/test_lightning_task.py: -------------------------------------------------------------------------------- 1 | import tempfile 2 | from unittest.mock import Mock 3 | 4 | import numpy as np 5 | import torch 6 | 7 | from nanodet.trainer.task import TrainingTask 8 | from nanodet.util import NanoDetLightningLogger, cfg, load_config 9 | 10 | 11 | class DummyRunner: 12 | def __init__(self, task): 13 | self.task = task 14 | 15 | def test(self): 16 | trainer = Mock() 17 | trainer.current_epoch = 0 18 | trainer.global_step = 0 19 | trainer.local_rank = 0 20 | trainer.use_ddp = False 21 | trainer.loggers = [NanoDetLightningLogger(tempfile.TemporaryDirectory().name)] 22 | trainer.num_val_batches = [1] 23 | 24 | optimizer = self.task.configure_optimizers()["optimizer"] 25 | 26 | trainer.optimizers = [optimizer] 27 | self.task._trainer = trainer 28 | 29 | self.task.on_train_start() 30 | assert self.task.current_epoch == 0 31 | 32 | dummy_batch = { 33 | "img": torch.randn((2, 3, 32, 32)), 34 | "img_info": { 35 | "height": torch.randn(2), 36 | "width": torch.randn(2), 37 | "id": torch.from_numpy(np.array([0, 1])), 38 | }, 39 | "gt_bboxes": [ 40 | np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), 41 | np.array( 42 | [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]], dtype=np.float32 43 | ), 44 | ], 45 | "gt_bboxes_ignore": [ 46 | np.array([[3.0, 4.0, 5.0, 6.0]], dtype=np.float32), 47 | np.array( 48 | [[7.0, 8.0, 9.0, 10.0], [7.0, 8.0, 9.0, 10.0]], dtype=np.float32 49 | ), 50 | ], 51 | "gt_labels": [np.array([1]), np.array([1, 2])], 52 | "warp_matrix": [np.eye(3), np.eye(3)], 53 | } 54 | 55 | def func(*args, **kwargs): 56 | pass 57 | 58 | self.task.scalar_summary = func 59 | self.task.training_step(dummy_batch, 0) 60 | 61 | self.task.optimizer_step(optimizer=optimizer) 62 | self.task.training_epoch_end([]) 63 | 64 | self.task.validation_step(dummy_batch, 0) 65 | self.task.validation_epoch_end([]) 66 | 67 | self.task.test_step(dummy_batch, 0) 68 | self.task.test_epoch_end([]) 69 | 70 | 71 | def test_lightning_training_task(): 72 | load_config(cfg, "./config/legacy_v0.x_configs/nanodet-m.yml") 73 | task = TrainingTask(cfg) 74 | runner = DummyRunner(task) 75 | runner.test() 76 | -------------------------------------------------------------------------------- /nanodet/model/arch/one_stage_detector.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import time 16 | 17 | import torch 18 | import torch.nn as nn 19 | 20 | from ..backbone import build_backbone 21 | from ..fpn import build_fpn 22 | from ..head import build_head 23 | 24 | 25 | class OneStageDetector(nn.Module): 26 | def __init__( 27 | self, 28 | backbone_cfg, 29 | fpn_cfg=None, 30 | head_cfg=None, 31 | ): 32 | super(OneStageDetector, self).__init__() 33 | self.backbone = build_backbone(backbone_cfg) 34 | if fpn_cfg is not None: 35 | self.fpn = build_fpn(fpn_cfg) 36 | if head_cfg is not None: 37 | self.head = build_head(head_cfg) 38 | self.epoch = 0 39 | 40 | def forward(self, x): 41 | x = self.backbone(x) 42 | if hasattr(self, "fpn"): 43 | x = self.fpn(x) 44 | if hasattr(self, "head"): 45 | x = self.head(x) 46 | return x 47 | 48 | def inference(self, meta): 49 | with torch.no_grad(): 50 | is_cuda_available = torch.cuda.is_available() 51 | if is_cuda_available: 52 | torch.cuda.synchronize() 53 | 54 | time1 = time.time() 55 | preds = self(meta["img"]) 56 | 57 | if is_cuda_available: 58 | torch.cuda.synchronize() 59 | 60 | time2 = time.time() 61 | print("forward time: {:.3f}s".format((time2 - time1)), end=" | ") 62 | results = self.head.post_process(preds, meta) 63 | 64 | if is_cuda_available: 65 | torch.cuda.synchronize() 66 | 67 | print("decode time: {:.3f}s".format((time.time() - time2)), end=" | ") 68 | return results 69 | 70 | def forward_train(self, gt_meta): 71 | preds = self(gt_meta["img"]) 72 | loss, loss_states = self.head.loss(preds, gt_meta) 73 | 74 | return preds, loss, loss_states 75 | 76 | def set_epoch(self, epoch): 77 | self.epoch = epoch 78 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/cpp/YoloV4.cpp: -------------------------------------------------------------------------------- 1 | #include "YoloV4.h" 2 | 3 | bool YoloV4::hasGPU = true; 4 | YoloV4 *YoloV4::detector = nullptr; 5 | 6 | YoloV4::YoloV4(AAssetManager *mgr, const char *param, const char *bin, bool useGPU) { 7 | Net = new ncnn::Net(); 8 | // opt 需要在加载前设置 9 | hasGPU = ncnn::get_gpu_count() > 0; 10 | Net->opt.use_vulkan_compute = hasGPU && useGPU; // gpu 11 | Net->opt.use_fp16_arithmetic = true; // fp16运算加速 12 | Net->load_param(mgr, param); 13 | Net->load_model(mgr, bin); 14 | } 15 | 16 | YoloV4::~YoloV4() { 17 | delete Net; 18 | } 19 | 20 | std::vector YoloV4::detect(JNIEnv *env, jobject image, float threshold, float nms_threshold) { 21 | AndroidBitmapInfo img_size; 22 | AndroidBitmap_getInfo(env, image, &img_size); 23 | ncnn::Mat in_net = ncnn::Mat::from_android_bitmap_resize(env, image, ncnn::Mat::PIXEL_RGBA2RGB, input_size, 24 | input_size); 25 | float norm[3] = {1 / 255.f, 1 / 255.f, 1 / 255.f}; 26 | float mean[3] = {0, 0, 0}; 27 | in_net.substract_mean_normalize(mean, norm); 28 | auto ex = Net->create_extractor(); 29 | ex.set_light_mode(true); 30 | ex.set_num_threads(4); 31 | hasGPU = ncnn::get_gpu_count() > 0; 32 | ex.set_vulkan_compute(hasGPU); 33 | ex.input(0, in_net); 34 | std::vector result; 35 | ncnn::Mat blob; 36 | ex.extract("output", blob); 37 | auto boxes = decode_infer(blob, {(int) img_size.width, (int) img_size.height}, input_size, num_class, threshold); 38 | result.insert(result.begin(), boxes.begin(), boxes.end()); 39 | // nms(result,nms_threshold); 40 | return result; 41 | } 42 | 43 | inline float fast_exp(float x) { 44 | union { 45 | uint32_t i; 46 | float f; 47 | } v{}; 48 | v.i = (1 << 23) * (1.4426950409 * x + 126.93490512f); 49 | return v.f; 50 | } 51 | 52 | inline float sigmoid(float x) { 53 | return 1.0f / (1.0f + fast_exp(-x)); 54 | } 55 | 56 | std::vector 57 | YoloV4::decode_infer(ncnn::Mat &data, const yolocv::YoloSize &frame_size, int net_size, int num_classes, float threshold) { 58 | std::vector result; 59 | for (int i = 0; i < data.h; i++) { 60 | BoxInfo box; 61 | const float *values = data.row(i); 62 | box.label = values[0] - 1; 63 | box.score = values[1]; 64 | box.x1 = values[2] * (float) frame_size.width; 65 | box.y1 = values[3] * (float) frame_size.height; 66 | box.x2 = values[4] * (float) frame_size.width; 67 | box.y2 = values[5] * (float) frame_size.height; 68 | result.push_back(box); 69 | } 70 | return result; 71 | } 72 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/cpp/YoloV5.h: -------------------------------------------------------------------------------- 1 | // 2 | // Created by 邓昊晴 on 14/6/2020. 3 | // 4 | 5 | #ifndef YOLOV5_H 6 | #define YOLOV5_H 7 | 8 | #include "net.h" 9 | 10 | namespace yolocv { 11 | typedef struct { 12 | int width; 13 | int height; 14 | } YoloSize; 15 | } 16 | 17 | typedef struct { 18 | std::string name; 19 | int stride; 20 | std::vector anchors; 21 | } YoloLayerData; 22 | 23 | typedef struct BoxInfo { 24 | float x1; 25 | float y1; 26 | float x2; 27 | float y2; 28 | float score; 29 | int label; 30 | } BoxInfo; 31 | 32 | class YoloV5 { 33 | public: 34 | YoloV5(AAssetManager *mgr, const char *param, const char *bin, bool useGPU); 35 | 36 | ~YoloV5(); 37 | 38 | std::vector detect(JNIEnv *env, jobject image, float threshold, float nms_threshold); 39 | std::vector labels{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 40 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 41 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 42 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 43 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 44 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 45 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 46 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 47 | "hair drier", "toothbrush"}; 48 | private: 49 | static std::vector 50 | decode_infer(ncnn::Mat &data, int stride, const yolocv::YoloSize &frame_size, int net_size, int num_classes, 51 | const std::vector &anchors, float threshold); 52 | 53 | static void nms(std::vector &result, float nms_threshold); 54 | 55 | ncnn::Net *Net; 56 | int input_size = 640; 57 | int num_class = 80; 58 | std::vector layers{ 59 | {"394", 32, {{116, 90}, {156, 198}, {373, 326}}}, 60 | {"375", 16, {{30, 61}, {62, 45}, {59, 119}}}, 61 | {"output", 8, {{10, 13}, {16, 30}, {33, 23}}}, 62 | }; 63 | 64 | public: 65 | static YoloV5 *detector; 66 | static bool hasGPU; 67 | }; 68 | 69 | 70 | #endif //YOLOV5_H 71 | -------------------------------------------------------------------------------- /tests/test_models/test_head/test_gfl_head.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | 4 | from nanodet.model.head import build_head 5 | from nanodet.util.yacs import CfgNode 6 | 7 | 8 | def test_gfl_head_loss(): 9 | head_cfg = dict( 10 | name="GFLHead", 11 | num_classes=80, 12 | input_channel=1, 13 | feat_channels=96, 14 | stacked_convs=2, 15 | strides=[8, 16, 32], 16 | loss=dict( 17 | loss_qfl=dict( 18 | name="QualityFocalLoss", use_sigmoid=True, beta=2.0, loss_weight=1.0 19 | ), 20 | loss_dfl=dict(name="DistributionFocalLoss", loss_weight=0.25), 21 | loss_bbox=dict(name="GIoULoss", loss_weight=2.0), 22 | ), 23 | ) 24 | cfg = CfgNode(head_cfg) 25 | 26 | head = build_head(cfg) 27 | feat = [torch.rand(1, 1, 320 // stride, 320 // stride) for stride in [8, 16, 32]] 28 | preds = head.forward(feat) 29 | 30 | # Test that empty ground truth encourages the network to predict background 31 | meta = dict( 32 | img=torch.rand((2, 3, 64, 64)), 33 | gt_bboxes=[np.random.random((0, 4))], 34 | gt_bboxes_ignore=[np.random.random((0, 4))], 35 | gt_labels=[np.array([])], 36 | ) 37 | loss, empty_gt_losses = head.loss(preds, meta) 38 | # When there is no truth, the cls loss should be nonzero but there should 39 | # be no box loss. 40 | empty_qfl_loss = empty_gt_losses["loss_qfl"] 41 | empty_box_loss = empty_gt_losses["loss_bbox"] 42 | empty_dfl_loss = empty_gt_losses["loss_dfl"] 43 | assert empty_qfl_loss.item() == 0 44 | assert ( 45 | empty_box_loss.item() == 0 46 | ), "there should be no box loss when there are no true boxes" 47 | assert ( 48 | empty_dfl_loss.item() == 0 49 | ), "there should be no dfl loss when there are no true boxes" 50 | 51 | # When truth is non-empty then both cls and box loss should be nonzero for 52 | # random inputs 53 | gt_bboxes = [ 54 | np.array([[23.6667, 23.8757, 238.6326, 151.8874]], dtype=np.float32), 55 | ] 56 | gt_bboxes_ignore = [ 57 | np.array([[29.6667, 29.8757, 244.6326, 160.8874]], dtype=np.float32), 58 | ] 59 | gt_labels = [np.array([2])] 60 | meta = dict( 61 | img=torch.rand((2, 3, 64, 64)), 62 | gt_bboxes=gt_bboxes, 63 | gt_labels=gt_labels, 64 | gt_bboxes_ignore=gt_bboxes_ignore, 65 | ) 66 | loss, one_gt_losses = head.loss(preds, meta) 67 | onegt_qfl_loss = one_gt_losses["loss_qfl"] 68 | onegt_box_loss = one_gt_losses["loss_bbox"] 69 | onegt_dfl_loss = one_gt_losses["loss_dfl"] 70 | assert onegt_qfl_loss.item() > 0, "qfl loss should be non-zero" 71 | assert onegt_box_loss.item() > 0, "box loss should be non-zero" 72 | assert onegt_dfl_loss.item() > 0, "dfl loss should be non-zero" 73 | -------------------------------------------------------------------------------- /tools/inference.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import os 16 | import time 17 | 18 | import cv2 19 | import torch 20 | 21 | from nanodet.data.transform import Pipeline 22 | from nanodet.model.arch import build_model 23 | from nanodet.util import load_model_weight 24 | 25 | 26 | class Predictor(object): 27 | def __init__(self, cfg, model_path, logger, device="cuda:0"): 28 | self.cfg = cfg 29 | self.device = device 30 | model = build_model(cfg.model) 31 | ckpt = torch.load(model_path, map_location=lambda storage, loc: storage) 32 | load_model_weight(model, ckpt, logger) 33 | if cfg.model.arch.backbone.name == "RepVGG": 34 | deploy_config = cfg.model 35 | deploy_config.arch.backbone.update({"deploy": True}) 36 | deploy_model = build_model(deploy_config) 37 | from nanodet.model.backbone.repvgg import repvgg_det_model_convert 38 | 39 | model = repvgg_det_model_convert(model, deploy_model) 40 | self.model = model.to(device).eval() 41 | self.pipeline = Pipeline(cfg.data.val.pipeline, cfg.data.val.keep_ratio) 42 | 43 | def inference(self, img): 44 | img_info = {} 45 | if isinstance(img, str): 46 | img_info["file_name"] = os.path.basename(img) 47 | img = cv2.imread(img) 48 | else: 49 | img_info["file_name"] = None 50 | 51 | height, width = img.shape[:2] 52 | img_info["height"] = height 53 | img_info["width"] = width 54 | meta = dict(img_info=img_info, raw_img=img, img=img) 55 | meta = self.pipeline(meta, self.cfg.data.val.input_size) 56 | meta["img"] = ( 57 | torch.from_numpy(meta["img"].transpose(2, 0, 1)) 58 | .unsqueeze(0) 59 | .to(self.device) 60 | ) 61 | with torch.no_grad(): 62 | results = self.model.inference(meta) 63 | return meta, results 64 | 65 | def visualize(self, dets, meta, class_names, score_thres, wait=0): 66 | time1 = time.time() 67 | self.model.head.show_result( 68 | meta["raw_img"], dets, class_names, score_thres=score_thres, show=True 69 | ) 70 | print("viz time: {:.3f}s".format(time.time() - time1)) 71 | -------------------------------------------------------------------------------- /nanodet/optim/builder.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import logging 3 | 4 | import torch 5 | from torch.nn import GroupNorm, LayerNorm 6 | from torch.nn.modules.batchnorm import _BatchNorm 7 | 8 | NORMS = (GroupNorm, LayerNorm, _BatchNorm) 9 | 10 | 11 | def build_optimizer(model, config): 12 | """Build optimizer from config. 13 | 14 | Supports customised parameter-level hyperparameters. 15 | The config should be like: 16 | >>> optimizer: 17 | >>> name: AdamW 18 | >>> lr: 0.001 19 | >>> weight_decay: 0.05 20 | >>> no_norm_decay: True 21 | >>> param_level_cfg: # parameter-level config 22 | >>> backbone: 23 | >>> lr_mult: 0.1 24 | """ 25 | config = copy.deepcopy(config) 26 | param_dict = {} 27 | no_norm_decay = config.pop("no_norm_decay", False) 28 | no_bias_decay = config.pop("no_bias_decay", False) 29 | param_level_cfg = config.pop("param_level_cfg", {}) 30 | base_lr = config.get("lr", None) 31 | base_wd = config.get("weight_decay", None) 32 | 33 | name = config.pop("name") 34 | optim_cls = getattr(torch.optim, name) 35 | 36 | logger = logging.getLogger("NanoDet") 37 | 38 | # custom param-wise lr and weight_decay 39 | for name, p in model.named_parameters(): 40 | if not p.requires_grad: 41 | continue 42 | param_dict[p] = {"name": name} 43 | 44 | for key in param_level_cfg: 45 | if key in name: 46 | if "lr_mult" in param_level_cfg[key] and base_lr: 47 | param_dict[p].update( 48 | {"lr": base_lr * param_level_cfg[key]["lr_mult"]} 49 | ) 50 | if "decay_mult" in param_level_cfg[key] and base_wd: 51 | param_dict[p].update( 52 | {"weight_decay": base_wd * param_level_cfg[key]["decay_mult"]} 53 | ) 54 | break 55 | if no_norm_decay: 56 | # update norms decay 57 | for name, m in model.named_modules(): 58 | if isinstance(m, NORMS): 59 | param_dict[m.bias].update({"weight_decay": 0}) 60 | param_dict[m.weight].update({"weight_decay": 0}) 61 | if no_bias_decay: 62 | # update bias decay 63 | for name, m in model.named_modules(): 64 | if hasattr(m, "bias"): 65 | param_dict[m.bias].update({"weight_decay": 0}) 66 | 67 | # convert param dict to optimizer's param groups 68 | param_groups = [] 69 | for p, pconfig in param_dict.items(): 70 | name = pconfig.pop("name", None) 71 | if "weight_decay" in pconfig or "lr" in pconfig: 72 | logger.info(f"special optimizer hyperparameter: {name} - {pconfig}") 73 | param_groups += [{"params": p, **pconfig}] 74 | 75 | optimizer = optim_cls(param_groups, **config) 76 | return optimizer 77 | -------------------------------------------------------------------------------- /tests/test_models/test_modules/test_conv.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | import torch.nn as nn 4 | 5 | from nanodet.model.module.conv import ConvModule 6 | 7 | 8 | def test_conv_module(): 9 | with pytest.raises(AssertionError): 10 | # conv_cfg must be a dict or None 11 | conv_cfg = "conv" 12 | ConvModule(3, 5, 2, conv_cfg=conv_cfg) 13 | 14 | with pytest.raises(AssertionError): 15 | # norm_cfg must be a dict or None 16 | norm_cfg = "norm" 17 | ConvModule(3, 5, 2, norm_cfg=norm_cfg) 18 | 19 | with pytest.raises(AssertionError): 20 | # softmax is not supported 21 | activation = "softmax" 22 | ConvModule(3, 5, 2, activation=activation) 23 | 24 | with pytest.raises(AssertionError): 25 | # softmax is not supported 26 | activation = dict(type="softmax") 27 | ConvModule(3, 5, 2, activation=activation) 28 | 29 | # conv + norm + act 30 | conv = ConvModule(3, 5, 2, norm_cfg=dict(type="BN")) 31 | assert hasattr(conv, "act") 32 | assert conv.with_norm 33 | assert hasattr(conv, "norm") 34 | x = torch.rand(1, 3, 16, 16) 35 | output = conv(x) 36 | assert output.shape == (1, 5, 15, 15) 37 | 38 | # conv + act 39 | conv = ConvModule(3, 5, 2) 40 | assert hasattr(conv, "act") 41 | assert not conv.with_norm 42 | assert conv.norm is None 43 | x = torch.rand(1, 3, 16, 16) 44 | output = conv(x) 45 | assert output.shape == (1, 5, 15, 15) 46 | 47 | # conv 48 | conv = ConvModule(3, 5, 2, activation=None) 49 | assert not conv.with_norm 50 | assert conv.norm is None 51 | assert not hasattr(conv, "act") 52 | x = torch.rand(1, 3, 16, 16) 53 | output = conv(x) 54 | assert output.shape == (1, 5, 15, 15) 55 | 56 | # leaky relu 57 | conv = ConvModule(3, 5, 3, padding=1, activation="LeakyReLU") 58 | assert isinstance(conv.act, nn.LeakyReLU) 59 | output = conv(x) 60 | assert output.shape == (1, 5, 16, 16) 61 | 62 | # PReLU 63 | conv = ConvModule(3, 5, 3, padding=1, activation="PReLU") 64 | assert isinstance(conv.act, nn.PReLU) 65 | output = conv(x) 66 | assert output.shape == (1, 5, 16, 16) 67 | 68 | 69 | def test_bias(): 70 | # bias: auto, without norm 71 | conv = ConvModule(3, 5, 2) 72 | assert conv.conv.bias is not None 73 | 74 | # bias: auto, with norm 75 | conv = ConvModule(3, 5, 2, norm_cfg=dict(type="BN")) 76 | assert conv.conv.bias is None 77 | 78 | # bias: False, without norm 79 | conv = ConvModule(3, 5, 2, bias=False) 80 | assert conv.conv.bias is None 81 | 82 | # bias: True, with norm 83 | with pytest.warns(UserWarning) as record: 84 | ConvModule(3, 5, 2, bias=True, norm_cfg=dict(type="BN")) 85 | assert len(record) == 1 86 | assert record[0].message.args[0] == "ConvModule has norm and bias at the same time" 87 | -------------------------------------------------------------------------------- /nanodet/util/env_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import platform 3 | import warnings 4 | 5 | import torch.multiprocessing as mp 6 | 7 | 8 | def set_multi_processing( 9 | mp_start_method: str = "fork", opencv_num_threads: int = 0, distributed: bool = True 10 | ) -> None: 11 | """Set multi-processing related environment. 12 | 13 | This function is refered from https://github.com/open-mmlab/mmengine/blob/main/mmengine/utils/dl_utils/setup_env.py 14 | 15 | Args: 16 | mp_start_method (str): Set the method which should be used to start 17 | child processes. Defaults to 'fork'. 18 | opencv_num_threads (int): Number of threads for opencv. 19 | Defaults to 0. 20 | distributed (bool): True if distributed environment. 21 | Defaults to False. 22 | """ # noqa 23 | # set multi-process start method as `fork` to speed up the training 24 | if platform.system() != "Windows": 25 | current_method = mp.get_start_method(allow_none=True) 26 | if current_method is not None and current_method != mp_start_method: 27 | warnings.warn( 28 | f"Multi-processing start method `{mp_start_method}` is " 29 | f"different from the previous setting `{current_method}`." 30 | f"It will be force set to `{mp_start_method}`. You can " 31 | "change this behavior by changing `mp_start_method` in " 32 | "your config." 33 | ) 34 | mp.set_start_method(mp_start_method, force=True) 35 | 36 | try: 37 | import cv2 38 | 39 | # disable opencv multithreading to avoid system being overloaded 40 | cv2.setNumThreads(opencv_num_threads) 41 | except ImportError: 42 | pass 43 | 44 | # setup OMP threads 45 | # This code is referred from https://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py # noqa 46 | if "OMP_NUM_THREADS" not in os.environ and distributed: 47 | omp_num_threads = 1 48 | warnings.warn( 49 | "Setting OMP_NUM_THREADS environment variable for each process" 50 | f" to be {omp_num_threads} in default, to avoid your system " 51 | "being overloaded, please further tune the variable for " 52 | "optimal performance in your application as needed." 53 | ) 54 | os.environ["OMP_NUM_THREADS"] = str(omp_num_threads) 55 | 56 | # setup MKL threads 57 | if "MKL_NUM_THREADS" not in os.environ and distributed: 58 | mkl_num_threads = 1 59 | warnings.warn( 60 | "Setting MKL_NUM_THREADS environment variable for each process" 61 | f" to be {mkl_num_threads} in default, to avoid your system " 62 | "being overloaded, please further tune the variable for " 63 | "optimal performance in your application as needed." 64 | ) 65 | os.environ["MKL_NUM_THREADS"] = str(mkl_num_threads) 66 | -------------------------------------------------------------------------------- /demo_ncnn/nanodet.h: -------------------------------------------------------------------------------- 1 | // 2 | // Create by RangiLyu 3 | // 2020 / 10 / 2 4 | // 5 | 6 | #ifndef NANODET_H 7 | #define NANODET_H 8 | 9 | #include 10 | #include 11 | 12 | typedef struct HeadInfo 13 | { 14 | std::string cls_layer; 15 | std::string dis_layer; 16 | int stride; 17 | }; 18 | 19 | struct CenterPrior 20 | { 21 | int x; 22 | int y; 23 | int stride; 24 | }; 25 | 26 | typedef struct BoxInfo 27 | { 28 | float x1; 29 | float y1; 30 | float x2; 31 | float y2; 32 | float score; 33 | int label; 34 | } BoxInfo; 35 | 36 | class NanoDet 37 | { 38 | public: 39 | NanoDet(const char* param, const char* bin, bool useGPU); 40 | 41 | ~NanoDet(); 42 | 43 | static NanoDet* detector; 44 | ncnn::Net* Net; 45 | static bool hasGPU; 46 | // modify these parameters to the same with your config if you want to use your own model 47 | int input_size[2] = {416, 416}; // input height and width 48 | int num_class = 80; // number of classes. 80 for COCO 49 | int reg_max = 7; // `reg_max` set in the training config. Default: 7. 50 | std::vector strides = { 8, 16, 32, 64 }; // strides of the multi-level feature. 51 | 52 | std::vector detect(cv::Mat image, float score_threshold, float nms_threshold); 53 | 54 | std::vector labels{ "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 55 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 56 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 57 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 58 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 59 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 60 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 61 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 62 | "hair drier", "toothbrush" }; 63 | private: 64 | void preprocess(cv::Mat& image, ncnn::Mat& in); 65 | void decode_infer(ncnn::Mat& feats, std::vector& center_priors, float threshold, std::vector>& results); 66 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride); 67 | static void nms(std::vector& result, float nms_threshold); 68 | 69 | }; 70 | 71 | 72 | #endif //NANODET_H 73 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/cpp/NanoDet.h: -------------------------------------------------------------------------------- 1 | // 2 | // Create by RangiLyu 3 | // 2020 / 10 / 2 4 | // 5 | 6 | #ifndef NANODET_H 7 | #define NANODET_H 8 | 9 | #include "net.h" 10 | #include "YoloV5.h" 11 | 12 | typedef struct HeadInfo_ 13 | { 14 | std::string cls_layer; 15 | std::string dis_layer; 16 | int stride; 17 | } HeadInfo; 18 | 19 | typedef struct CenterPrior_ 20 | { 21 | int x; 22 | int y; 23 | int stride; 24 | } CenterPrior; 25 | 26 | 27 | class NanoDet{ 28 | public: 29 | NanoDet(AAssetManager *mgr, const char *param, const char *bin, bool useGPU); 30 | 31 | ~NanoDet(); 32 | 33 | std::vector detect(JNIEnv *env, jobject image, float score_threshold, float nms_threshold); 34 | std::vector labels{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 35 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 36 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 37 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 38 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 39 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 40 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 41 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 42 | "hair drier", "toothbrush"}; 43 | private: 44 | void preprocess(JNIEnv *env, jobject image, ncnn::Mat& in); 45 | void decode_infer(ncnn::Mat& feats, std::vector& center_priors, float threshold, std::vector>& results, float width_ratio, float height_ratio); 46 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride, float width_ratio, float height_ratio); 47 | 48 | static void nms(std::vector& result, float nms_threshold); 49 | 50 | ncnn::Net *Net; 51 | // modify these parameters to the same with your config if you want to use your own model 52 | int input_size[2] = {320, 320}; // input height and width 53 | int num_class = 80; // number of classes. 80 for COCO 54 | int reg_max = 7; // `reg_max` set in the training config. Default: 7. 55 | std::vector strides = { 8, 16, 32, 64 }; // strides of the multi-level feature. 56 | 57 | 58 | public: 59 | static NanoDet *detector; 60 | static bool hasGPU; 61 | }; 62 | 63 | 64 | #endif //NANODET_H 65 | -------------------------------------------------------------------------------- /demo_mnn/README.md: -------------------------------------------------------------------------------- 1 | # NanoDet MNN Demo 2 | 3 | This fold provides NanoDet inference code using 4 | [Alibaba's MNN framework](https://github.com/alibaba/MNN). Most of the implements in 5 | this fold are same as *demo_ncnn*. 6 | 7 | ## Install MNN 8 | 9 | ### Python library 10 | 11 | Just run: 12 | 13 | ``` shell 14 | pip install MNN 15 | ``` 16 | 17 | ### C++ library 18 | 19 | Please follow the [official document](https://www.yuque.com/mnn/en/build_linux) to build MNN engine. 20 | 21 | ## Convert model 22 | 23 | 1. Export ONNX model 24 | 25 | ```shell 26 | python tools/export_onnx.py --cfg_path ${CONFIG_PATH} --model_path ${PYTORCH_MODEL_PATH} 27 | ``` 28 | 29 | 2. Convert to MNN 30 | 31 | ``` shell 32 | python -m MNN.tools.mnnconvert -f ONNX --modelFile sim.onnx --MNNModel nanodet.mnn 33 | ``` 34 | 35 | It should be note that the input size does not have to be fixed, it can be any integer multiple of strides, 36 | since NanoDet is anchor free. We can adapt the shape of `dummy_input` in *./tools/export_onnx.py* to get ONNX and MNN models 37 | with different input sizes. 38 | 39 | Here are converted model 40 | [Download Link](https://github.com/RangiLyu/nanodet/releases/download/v1.0.0-alpha-1/nanodet-plus-m_416_mnn.mnn). 41 | 42 | ## Build 43 | 44 | For C++ code, replace `libMNN.so` under *./mnn/lib* with the one you just compiled, modify OpenCV path at CMake file, 45 | and run 46 | 47 | ``` shell 48 | mkdir build && cd build 49 | cmake .. 50 | make 51 | ``` 52 | 53 | Note that a flag at `main.cpp` is used to control whether to show the detection result or save it into a fold. 54 | 55 | ``` c++ 56 | #define __SAVE_RESULT__ // if defined save drawed results to ../results, else show it in windows 57 | ``` 58 | 59 | ## Run 60 | 61 | ### Python 62 | 63 | The multi-backend python demo is still working in progress. 64 | 65 | ### C++ 66 | 67 | C++ inference interface is same with NCNN code, to detect images in a fold, run: 68 | 69 | ``` shell 70 | ./nanodet-mnn "1" "../imgs/*.jpg" 71 | ``` 72 | 73 | For speed benchmark 74 | 75 | ``` shell 76 | ./nanodet-mnn "3" "0" 77 | ``` 78 | 79 | ## Custom model 80 | 81 | If you want to use custom model, please make sure the hyperparameters 82 | in `nanodet_mnn.h` are the same with your training config file. 83 | 84 | ```cpp 85 | int input_size[2] = {416, 416}; // input height and width 86 | int num_class = 80; // number of classes. 80 for COCO 87 | int reg_max = 7; // `reg_max` set in the training config. Default: 7. 88 | std::vector strides = { 8, 16, 32, 64 }; // strides of the multi-level feature. 89 | ``` 90 | 91 | ## Reference 92 | 93 | [Ultra-Light-Fast-Generic-Face-Detector-1MB](https://github.com/Linzaer/Ultra-Light-Fast-Generic-Face-Detector-1MB/tree/master/MNN) 94 | 95 | [ONNX Simplifier](https://github.com/daquexian/onnx-simplifier) 96 | 97 | [NanoDet NCNN](https://github.com/RangiLyu/nanodet/tree/main/demo_ncnn) 98 | 99 | [MNN](https://github.com/alibaba/MNN) 100 | 101 | ## Example results 102 | 103 | ![screenshot](./results/000252.jpg?raw=true) 104 | ![screenshot](./results/000258.jpg?raw=true) 105 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/WelcomeActivity.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import androidx.appcompat.app.AlertDialog; 4 | import androidx.appcompat.app.AppCompatActivity; 5 | 6 | import android.content.Intent; 7 | import android.os.Bundle; 8 | import android.view.View; 9 | import android.widget.Button; 10 | import android.widget.CompoundButton; 11 | import android.widget.ToggleButton; 12 | 13 | 14 | public class WelcomeActivity extends AppCompatActivity { 15 | 16 | private ToggleButton tbUseGpu; 17 | private Button nanodet; 18 | private Button yolov5s; 19 | private Button yolov4tiny; 20 | 21 | private boolean useGPU = false; 22 | 23 | @Override 24 | protected void onCreate(Bundle savedInstanceState) { 25 | super.onCreate(savedInstanceState); 26 | setContentView(R.layout.activity_welcome); 27 | 28 | tbUseGpu = findViewById(R.id.tb_use_gpu); 29 | tbUseGpu.setOnCheckedChangeListener(new CompoundButton.OnCheckedChangeListener() { 30 | @Override 31 | public void onCheckedChanged(CompoundButton buttonView, boolean isChecked) { 32 | useGPU = isChecked; 33 | MainActivity.USE_GPU = useGPU; 34 | if (useGPU) { 35 | AlertDialog.Builder builder = new AlertDialog.Builder(WelcomeActivity.this); 36 | builder.setTitle("Warning"); 37 | builder.setMessage("It may not work well in GPU mode, or errors may occur."); 38 | builder.setCancelable(true); 39 | builder.setPositiveButton("OK", null); 40 | AlertDialog dialog = builder.create(); 41 | dialog.show(); 42 | } 43 | } 44 | }); 45 | 46 | nanodet = findViewById(R.id.btn_start_detect0); 47 | nanodet.setOnClickListener(new View.OnClickListener() { 48 | @Override 49 | public void onClick(View v) { 50 | MainActivity.USE_MODEL = MainActivity.NANODET; 51 | Intent intent = new Intent(WelcomeActivity.this, MainActivity.class); 52 | WelcomeActivity.this.startActivity(intent); 53 | } 54 | }); 55 | 56 | yolov5s = findViewById(R.id.btn_start_detect1); 57 | yolov5s.setOnClickListener(new View.OnClickListener() { 58 | @Override 59 | public void onClick(View v) { 60 | MainActivity.USE_MODEL = MainActivity.YOLOV5S; 61 | Intent intent = new Intent(WelcomeActivity.this, MainActivity.class); 62 | WelcomeActivity.this.startActivity(intent); 63 | } 64 | }); 65 | 66 | yolov4tiny = findViewById(R.id.btn_start_detect2); 67 | yolov4tiny.setOnClickListener(new View.OnClickListener() { 68 | @Override 69 | public void onClick(View v) { 70 | MainActivity.USE_MODEL = MainActivity.YOLOV4_TINY; 71 | Intent intent = new Intent(WelcomeActivity.this, MainActivity.class); 72 | WelcomeActivity.this.startActivity(intent); 73 | } 74 | }); 75 | 76 | } 77 | 78 | } 79 | -------------------------------------------------------------------------------- /tests/test_data/test_transform/test_warp.py: -------------------------------------------------------------------------------- 1 | import copy 2 | 3 | import numpy as np 4 | 5 | from nanodet.data.transform.warp import ( 6 | ShapeTransform, 7 | get_flip_matrix, 8 | get_perspective_matrix, 9 | get_rotation_matrix, 10 | get_scale_matrix, 11 | get_shear_matrix, 12 | get_stretch_matrix, 13 | get_translate_matrix, 14 | warp_and_resize, 15 | ) 16 | 17 | 18 | def test_get_matrix(): 19 | # TODO: better unit test 20 | height = 100 21 | width = 200 22 | 23 | # center 24 | C = np.eye(3) 25 | C[0, 2] = -width / 2 26 | C[1, 2] = -height / 2 27 | 28 | # do not change the order of mat mul 29 | P = get_perspective_matrix(0.1) 30 | C = P @ C 31 | 32 | Scl = get_scale_matrix((1, 2)) 33 | C = Scl @ C 34 | 35 | Str = get_stretch_matrix((0.5, 1.5), (0.5, 1.5)) 36 | C = Str @ C 37 | 38 | R = get_rotation_matrix(180) 39 | C = R @ C 40 | 41 | Sh = get_shear_matrix(60) 42 | C = Sh @ C 43 | 44 | F = get_flip_matrix(0.5) 45 | C = F @ C 46 | 47 | T = get_translate_matrix(0.5, width, height) 48 | 49 | M = T @ C 50 | 51 | assert M.shape == (3, 3) 52 | 53 | 54 | def test_warp(): 55 | dummy_meta = dict( 56 | img=np.random.randint(0, 255, size=(100, 200, 3), dtype=np.uint8), 57 | gt_bboxes=np.array([[0, 0, 20, 20]]), 58 | gt_masks=[np.zeros((100, 200), dtype=np.uint8)], 59 | ) 60 | warp_cfg = {} 61 | res = warp_and_resize( 62 | copy.deepcopy(dummy_meta), warp_cfg, dst_shape=(50, 50), keep_ratio=False 63 | ) 64 | assert res["img"].shape == (50, 50, 3) 65 | assert res["gt_masks"][0].shape == (50, 50) 66 | assert np.array_equal(res["gt_bboxes"], np.array([[0, 0, 5, 10]], dtype=np.float32)) 67 | 68 | res = warp_and_resize( 69 | copy.deepcopy(dummy_meta), warp_cfg, dst_shape=(50, 50), keep_ratio=True 70 | ) 71 | assert np.array_equal( 72 | res["gt_bboxes"], np.array([[0, 12.5, 5.0, 17.5]], dtype=np.float32) 73 | ) 74 | 75 | res = warp_and_resize( 76 | copy.deepcopy(dummy_meta), warp_cfg, dst_shape=(300, 300), keep_ratio=True 77 | ) 78 | assert np.array_equal( 79 | res["gt_bboxes"], np.array([[0, 75, 30, 105]], dtype=np.float32) 80 | ) 81 | 82 | 83 | def test_shape_transform(): 84 | dummy_meta = dict( 85 | img=np.random.randint(0, 255, size=(100, 200, 3), dtype=np.uint8), 86 | gt_bboxes=np.array([[0, 0, 20, 20]]), 87 | gt_masks=[np.zeros((100, 200), dtype=np.uint8)], 88 | ) 89 | # keep ratio 90 | transform = ShapeTransform(keep_ratio=True, divisible=32) 91 | res = transform(dummy_meta, dst_shape=(50, 50)) 92 | assert np.array_equal( 93 | res["gt_bboxes"], np.array([[0, 0, 6.4, 6.4]], dtype=np.float32) 94 | ) 95 | assert res["img"].shape[0] % 32 == 0 96 | assert res["img"].shape[1] % 32 == 0 97 | 98 | # not keep ratio 99 | transform = ShapeTransform(keep_ratio=False) 100 | res = transform(dummy_meta, dst_shape=(50, 50)) 101 | assert np.array_equal(res["gt_bboxes"], np.array([[0, 0, 5, 10]], dtype=np.float32)) 102 | -------------------------------------------------------------------------------- /tests/test_models/test_backbone/test_efficient_lite.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import torch 3 | 4 | from nanodet.model.backbone import EfficientNetLite, build_backbone 5 | 6 | 7 | def test_efficientnet_lite(): 8 | with pytest.raises(AssertionError): 9 | cfg = dict( 10 | name="EfficientNetLite", 11 | model_name="efficientnet_lite0", 12 | out_stages=(7, 8, 9), 13 | ) 14 | build_backbone(cfg) 15 | 16 | with pytest.raises(AssertionError): 17 | EfficientNetLite(model_name="efficientnet_lite9") 18 | 19 | input = torch.rand(1, 3, 64, 64) 20 | 21 | model = EfficientNetLite( 22 | model_name="efficientnet_lite0", out_stages=(0, 1, 2, 3, 4, 5, 6), pretrain=True 23 | ) 24 | output = model(input) 25 | assert output[0].shape == (1, 16, 32, 32) 26 | assert output[1].shape == (1, 24, 16, 16) 27 | assert output[2].shape == (1, 40, 8, 8) 28 | assert output[3].shape == (1, 80, 4, 4) 29 | assert output[4].shape == (1, 112, 4, 4) 30 | assert output[5].shape == (1, 192, 2, 2) 31 | assert output[6].shape == (1, 320, 2, 2) 32 | 33 | model = EfficientNetLite( 34 | model_name="efficientnet_lite1", 35 | out_stages=(0, 1, 2, 3, 4, 5, 6), 36 | pretrain=False, 37 | ) 38 | output = model(input) 39 | assert output[0].shape == (1, 16, 32, 32) 40 | assert output[1].shape == (1, 24, 16, 16) 41 | assert output[2].shape == (1, 40, 8, 8) 42 | assert output[3].shape == (1, 80, 4, 4) 43 | assert output[4].shape == (1, 112, 4, 4) 44 | assert output[5].shape == (1, 192, 2, 2) 45 | assert output[6].shape == (1, 320, 2, 2) 46 | 47 | model = EfficientNetLite( 48 | model_name="efficientnet_lite2", 49 | out_stages=(0, 1, 2, 3, 4, 5, 6), 50 | activation="ReLU", 51 | pretrain=False, 52 | ) 53 | output = model(input) 54 | assert output[0].shape == (1, 16, 32, 32) 55 | assert output[1].shape == (1, 24, 16, 16) 56 | assert output[2].shape == (1, 48, 8, 8) 57 | assert output[3].shape == (1, 88, 4, 4) 58 | assert output[4].shape == (1, 120, 4, 4) 59 | assert output[5].shape == (1, 208, 2, 2) 60 | assert output[6].shape == (1, 352, 2, 2) 61 | 62 | model = EfficientNetLite( 63 | model_name="efficientnet_lite3", 64 | out_stages=(0, 1, 2, 3, 4, 5, 6), 65 | pretrain=False, 66 | ) 67 | output = model(input) 68 | assert output[0].shape == (1, 24, 32, 32) 69 | assert output[1].shape == (1, 32, 16, 16) 70 | assert output[2].shape == (1, 48, 8, 8) 71 | assert output[3].shape == (1, 96, 4, 4) 72 | assert output[4].shape == (1, 136, 4, 4) 73 | assert output[5].shape == (1, 232, 2, 2) 74 | assert output[6].shape == (1, 384, 2, 2) 75 | 76 | model = EfficientNetLite( 77 | model_name="efficientnet_lite4", 78 | out_stages=(0, 1, 2, 3, 4, 5, 6), 79 | pretrain=False, 80 | ) 81 | output = model(input) 82 | assert output[0].shape == (1, 24, 32, 32) 83 | assert output[1].shape == (1, 32, 16, 16) 84 | assert output[2].shape == (1, 56, 8, 8) 85 | assert output[3].shape == (1, 112, 4, 4) 86 | assert output[4].shape == (1, 160, 4, 4) 87 | assert output[5].shape == (1, 272, 2, 2) 88 | assert output[6].shape == (1, 448, 2, 2) 89 | -------------------------------------------------------------------------------- /nanodet/model/loss/utils.py: -------------------------------------------------------------------------------- 1 | import functools 2 | 3 | import torch.nn.functional as F 4 | 5 | 6 | def reduce_loss(loss, reduction): 7 | """Reduce loss as specified. 8 | 9 | Args: 10 | loss (Tensor): Elementwise loss tensor. 11 | reduction (str): Options are "none", "mean" and "sum". 12 | 13 | Return: 14 | Tensor: Reduced loss tensor. 15 | """ 16 | reduction_enum = F._Reduction.get_enum(reduction) 17 | # none: 0, elementwise_mean:1, sum: 2 18 | if reduction_enum == 0: 19 | return loss 20 | elif reduction_enum == 1: 21 | return loss.mean() 22 | elif reduction_enum == 2: 23 | return loss.sum() 24 | 25 | 26 | def weight_reduce_loss(loss, weight=None, reduction="mean", avg_factor=None): 27 | """Apply element-wise weight and reduce loss. 28 | 29 | Args: 30 | loss (Tensor): Element-wise loss. 31 | weight (Tensor): Element-wise weights. 32 | reduction (str): Same as built-in losses of PyTorch. 33 | avg_factor (float): Avarage factor when computing the mean of losses. 34 | 35 | Returns: 36 | Tensor: Processed loss values. 37 | """ 38 | # if weight is specified, apply element-wise weight 39 | if weight is not None: 40 | loss = loss * weight 41 | 42 | # if avg_factor is not specified, just reduce the loss 43 | if avg_factor is None: 44 | loss = reduce_loss(loss, reduction) 45 | else: 46 | # if reduction is mean, then average the loss by avg_factor 47 | if reduction == "mean": 48 | loss = loss.sum() / avg_factor 49 | # if reduction is 'none', then do nothing, otherwise raise an error 50 | elif reduction != "none": 51 | raise ValueError('avg_factor can not be used with reduction="sum"') 52 | return loss 53 | 54 | 55 | def weighted_loss(loss_func): 56 | """Create a weighted version of a given loss function. 57 | 58 | To use this decorator, the loss function must have the signature like 59 | `loss_func(pred, target, **kwargs)`. The function only needs to compute 60 | element-wise loss without any reduction. This decorator will add weight 61 | and reduction arguments to the function. The decorated function will have 62 | the signature like `loss_func(pred, target, weight=None, reduction='mean', 63 | avg_factor=None, **kwargs)`. 64 | 65 | :Example: 66 | 67 | >>> import torch 68 | >>> @weighted_loss 69 | >>> def l1_loss(pred, target): 70 | >>> return (pred - target).abs() 71 | 72 | >>> pred = torch.Tensor([0, 2, 3]) 73 | >>> target = torch.Tensor([1, 1, 1]) 74 | >>> weight = torch.Tensor([1, 0, 1]) 75 | 76 | >>> l1_loss(pred, target) 77 | tensor(1.3333) 78 | >>> l1_loss(pred, target, weight) 79 | tensor(1.) 80 | >>> l1_loss(pred, target, reduction='none') 81 | tensor([1., 1., 2.]) 82 | >>> l1_loss(pred, target, weight, avg_factor=2) 83 | tensor(1.5000) 84 | """ 85 | 86 | @functools.wraps(loss_func) 87 | def wrapper(pred, target, weight=None, reduction="mean", avg_factor=None, **kwargs): 88 | # get element-wise loss 89 | loss = loss_func(pred, target, **kwargs) 90 | loss = weight_reduce_loss(loss, weight, reduction, avg_factor) 91 | return loss 92 | 93 | return wrapper 94 | -------------------------------------------------------------------------------- /nanodet/data/collate.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 RangiLyu. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | import collections 16 | import re 17 | 18 | import torch 19 | from torch._six import string_classes 20 | 21 | np_str_obj_array_pattern = re.compile(r"[SaUO]") 22 | 23 | default_collate_err_msg_format = ( 24 | "default_collate: batch must contain tensors, numpy arrays, numbers, " 25 | "dicts or lists; found {}" 26 | ) 27 | 28 | 29 | def collate_function(batch): 30 | r"""Puts each data field into a tensor with outer dimension batch size""" 31 | 32 | elem = batch[0] 33 | elem_type = type(elem) 34 | if isinstance(elem, torch.Tensor): 35 | out = None 36 | if torch.utils.data.get_worker_info() is not None: 37 | # If we're in a background process, concatenate directly into a 38 | # shared memory tensor to avoid an extra copy 39 | numel = sum([x.numel() for x in batch]) 40 | storage = elem.storage()._new_shared(numel) 41 | out = elem.new(storage) 42 | return torch.stack(batch, 0, out=out) 43 | elif ( 44 | elem_type.__module__ == "numpy" 45 | and elem_type.__name__ != "str_" 46 | and elem_type.__name__ != "string_" 47 | ): 48 | elem = batch[0] 49 | if elem_type.__name__ == "ndarray": 50 | # array of string classes and object 51 | if np_str_obj_array_pattern.search(elem.dtype.str) is not None: 52 | raise TypeError(default_collate_err_msg_format.format(elem.dtype)) 53 | 54 | return batch 55 | elif elem.shape == (): # scalars 56 | return batch 57 | elif isinstance(elem, float): 58 | return torch.tensor(batch, dtype=torch.float64) 59 | elif isinstance(elem, int): 60 | return torch.tensor(batch) 61 | elif isinstance(elem, string_classes): 62 | return batch 63 | elif isinstance(elem, collections.abc.Mapping): 64 | return {key: collate_function([d[key] for d in batch]) for key in elem} 65 | elif isinstance(elem, tuple) and hasattr(elem, "_fields"): # namedtuple 66 | return elem_type(*(collate_function(samples) for samples in zip(*batch))) 67 | elif isinstance(elem, collections.abc.Sequence): 68 | transposed = zip(*batch) 69 | return [collate_function(samples) for samples in transposed] 70 | 71 | raise TypeError(default_collate_err_msg_format.format(elem_type)) 72 | 73 | 74 | def naive_collate(batch): 75 | """Only collate dict value in to a list. E.g. meta data dict and img_info 76 | dict will be collated.""" 77 | 78 | elem = batch[0] 79 | if isinstance(elem, dict): 80 | return {key: naive_collate([d[key] for d in batch]) for key in elem} 81 | else: 82 | return batch 83 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/layout/activity_welcome.xml: -------------------------------------------------------------------------------- 1 | 2 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 30 | 31 | 38 | 39 | 47 | 48 |