├── .flake8
├── .github
└── workflows
│ └── workflow.yml
├── .gitignore
├── .isort.cfg
├── .pre-commit-config.yaml
├── LICENSE
├── README.md
├── config
├── EfficientNet-Lite
│ ├── nanodet-EfficientNet-Lite0_320.yml
│ ├── nanodet-EfficientNet-Lite1_416.yml
│ └── nanodet-EfficientNet-Lite2_512.yml
├── RepVGG
│ └── nanodet-RepVGG-A0_416.yml
├── Transformer
│ └── nanodet-t.yml
├── nanodet-g.yml
├── nanodet-m-0.5x.yml
├── nanodet-m-1.5x-416.yml
├── nanodet-m-1.5x.yml
├── nanodet-m-416.yml
├── nanodet-m.yml
└── nanodet_custom_xml_dataset.yml
├── demo
├── demo-inference-with-pytorch.ipynb
└── demo.py
├── demo_android_ncnn
├── .gitignore
├── Android_demo.jpg
├── LICENSE
├── README.md
├── app
│ ├── .gitignore
│ ├── build.gradle
│ ├── proguard-rules.pro
│ └── src
│ │ ├── androidTest
│ │ └── java
│ │ │ └── com
│ │ │ └── rangi
│ │ │ └── nanodet
│ │ │ └── ExampleInstrumentedTest.java
│ │ ├── main
│ │ ├── AndroidManifest.xml
│ │ ├── cpp
│ │ │ ├── CMakeLists.txt
│ │ │ ├── NanoDet.cpp
│ │ │ ├── NanoDet.h
│ │ │ ├── YoloV4.cpp
│ │ │ ├── YoloV4.h
│ │ │ ├── YoloV5.cpp
│ │ │ ├── YoloV5.h
│ │ │ └── jni_interface.cpp
│ │ ├── java
│ │ │ └── com
│ │ │ │ └── rangi
│ │ │ │ └── nanodet
│ │ │ │ ├── AppCrashHandler.java
│ │ │ │ ├── Box.java
│ │ │ │ ├── MainActivity.java
│ │ │ │ ├── NanoDet.java
│ │ │ │ ├── NcnnApp.java
│ │ │ │ ├── WelcomeActivity.java
│ │ │ │ ├── YOLOv4.java
│ │ │ │ └── YOLOv5.java
│ │ └── res
│ │ │ ├── drawable-v24
│ │ │ └── ic_launcher_foreground.xml
│ │ │ ├── drawable-xxhdpi
│ │ │ ├── cpu.png
│ │ │ ├── gpu.png
│ │ │ └── ncnn_icon.png
│ │ │ ├── drawable
│ │ │ ├── cpu_gpu_bg.xml
│ │ │ └── ic_launcher_background.xml
│ │ │ ├── layout
│ │ │ ├── activity_main.xml
│ │ │ └── activity_welcome.xml
│ │ │ ├── mipmap-anydpi-v26
│ │ │ ├── ic_launcher.xml
│ │ │ └── ic_launcher_round.xml
│ │ │ ├── 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
│ │ │ └── values
│ │ │ ├── colors.xml
│ │ │ ├── strings.xml
│ │ │ └── styles.xml
│ │ └── test
│ │ └── java
│ │ └── com
│ │ └── rangi
│ │ └── nanodet
│ │ └── ExampleUnitTest.java
├── build.gradle
├── gradle.properties
├── gradle
│ └── wrapper
│ │ ├── gradle-wrapper.jar
│ │ └── gradle-wrapper.properties
├── gradlew
├── gradlew.bat
└── settings.gradle
├── demo_libtorch
├── CMakeLists.txt
├── README.md
├── main.cpp
├── nanodet_libtorch.cpp
└── nanodet_libtorch.h
├── demo_mnn
├── CMakeLists.txt
├── README.md
├── imgs
│ ├── 000252.jpg
│ └── 000258.jpg
├── main.cpp
├── nanodet_mnn.cpp
├── nanodet_mnn.hpp
├── python
│ └── demo_mnn.py
└── results
│ ├── 000252.jpg
│ └── 000258.jpg
├── demo_ncnn
├── CMakeLists.txt
├── README.md
├── benchmark.jpg
├── main.cpp
├── nanodet.cpp
├── nanodet.h
└── python
│ └── demo_ncnn.py
├── demo_openvino
├── CMakeLists.txt
├── README.md
├── main.cpp
├── nanodet_openvino.cpp
└── nanodet_openvino.h
├── docs
├── config_file_detail.md
├── imgs
│ ├── Android_demo.jpg
│ ├── Model_arch.png
│ └── Title.jpg
└── update.md
├── nanodet
├── __about__.py
├── __init__.py
├── data
│ ├── batch_process.py
│ ├── collate.py
│ ├── dataset
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── coco.py
│ │ └── xml_dataset.py
│ └── transform
│ │ ├── __init__.py
│ │ ├── color.py
│ │ ├── mosaic.py
│ │ ├── pipeline.py
│ │ └── warp.py
├── evaluator
│ ├── __init__.py
│ └── coco_detection.py
├── model
│ ├── arch
│ │ ├── __init__.py
│ │ └── one_stage_detector.py
│ ├── backbone
│ │ ├── __init__.py
│ │ ├── custom_csp.py
│ │ ├── efficientnet_lite.py
│ │ ├── ghostnet.py
│ │ ├── mobilenetv2.py
│ │ ├── repvgg.py
│ │ ├── resnet.py
│ │ └── shufflenetv2.py
│ ├── fpn
│ │ ├── __init__.py
│ │ ├── fpn.py
│ │ ├── pan.py
│ │ └── tan.py
│ ├── head
│ │ ├── __init__.py
│ │ ├── assigner
│ │ │ ├── assign_result.py
│ │ │ ├── atss_assigner.py
│ │ │ └── base_assigner.py
│ │ ├── gfl_head.py
│ │ └── nanodet_head.py
│ ├── loss
│ │ ├── gfocal_loss.py
│ │ ├── iou_loss.py
│ │ └── utils.py
│ └── module
│ │ ├── activation.py
│ │ ├── conv.py
│ │ ├── init_weights.py
│ │ ├── nms.py
│ │ ├── norm.py
│ │ ├── scale.py
│ │ └── transformer.py
├── trainer
│ ├── __init__.py
│ ├── dist_trainer.py
│ ├── task.py
│ └── trainer.py
└── util
│ ├── __init__.py
│ ├── box_transform.py
│ ├── check_point.py
│ ├── config.py
│ ├── data_parallel.py
│ ├── distributed_data_parallel.py
│ ├── flops_counter.py
│ ├── logger.py
│ ├── misc.py
│ ├── path.py
│ ├── rank_filter.py
│ ├── scatter_gather.py
│ ├── util_mixins.py
│ ├── visualization.py
│ └── yacs.py
├── requirements.txt
├── setup.py
├── tests
├── data
│ ├── batched_nms_data.pkl
│ ├── dummy_coco.json
│ ├── test_img.jpg
│ └── test_img.xml
├── test_configs
│ └── test_config.py
├── test_data
│ ├── test_batch_process.py
│ ├── test_collate.py
│ ├── test_dataset
│ │ ├── test_cocodataset.py
│ │ └── test_xmldataset.py
│ └── test_transform
│ │ ├── test_color.py
│ │ └── test_warp.py
├── test_evaluator
│ └── test_coco_detection.py
├── test_models
│ ├── test_backbone
│ │ ├── test_custom_csp.py
│ │ ├── test_efficient_lite.py
│ │ ├── test_ghostnet.py
│ │ ├── test_mobilenetv2.py
│ │ ├── test_repvgg.py
│ │ ├── test_resnet.py
│ │ └── test_shufflenetv2.py
│ ├── test_fpn
│ │ ├── test_fpn.py
│ │ ├── test_pan.py
│ │ └── test_tan.py
│ ├── test_head
│ │ ├── test_gfl_head.py
│ │ └── test_nanodet_head.py
│ ├── test_loss
│ │ ├── test_gfocal_loss.py
│ │ └── test_iou_loss.py
│ └── test_modules
│ │ ├── test_conv.py
│ │ ├── test_dwconv.py
│ │ ├── test_init_weights.py
│ │ ├── test_nms.py
│ │ ├── test_norm.py
│ │ ├── test_repvgg_conv.py
│ │ ├── test_scale.py
│ │ └── test_transformer.py
├── test_trainer
│ ├── test_lightning_task.py
│ └── test_trainer.py
└── test_utils
│ ├── test_flops.py
│ └── test_logger.py
└── tools
├── convert_old_checkpoint.py
├── deprecated
├── test.py
└── train.py
├── export_onnx.py
├── export_torchscript.py
├── flops.py
├── inference.py
├── test.py
└── train.py
/.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
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 |
--------------------------------------------------------------------------------
/.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 |
--------------------------------------------------------------------------------
/.isort.cfg:
--------------------------------------------------------------------------------
1 | [tool.isort]
2 | profile = "black"
3 |
--------------------------------------------------------------------------------
/.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.8.0
14 | hooks:
15 | - id: isort
16 | args: ["--profile", "black"]
17 |
18 | - repo: https://github.com/psf/black
19 | rev: 21.6b0
20 | hooks:
21 | - id: black
22 |
23 | - repo: https://gitlab.com/pycqa/flake8
24 | rev: 3.9.2
25 | hooks:
26 | - id: flake8
27 |
--------------------------------------------------------------------------------
/config/RepVGG/nanodet-RepVGG-A0_416.yml:
--------------------------------------------------------------------------------
1 | # nanodet-EfficientNet-Lite1_416
2 | save_dir: workspace/RepVGG-A0-416
3 | model:
4 | arch:
5 | name: OneStageDetector
6 | backbone:
7 | name: RepVGG
8 | arch: A0
9 | out_stages: [2,3,4]
10 | activation: ReLU
11 | last_channel: 512
12 | deploy: False
13 | fpn:
14 | name: PAN
15 | in_channels: [96, 192, 512]
16 | out_channels: 128
17 | start_level: 0
18 | num_outs: 3
19 | head:
20 | name: NanoDetHead
21 | num_classes: 80
22 | conv_type: Conv
23 | input_channel: 128
24 | feat_channels: 128
25 | stacked_convs: 2
26 | activation: ReLU
27 | share_cls_reg: True
28 | octave_base_scale: 8
29 | scales_per_octave: 1
30 | strides: [8, 16, 32]
31 | reg_max: 10
32 | norm_cfg:
33 | type: BN
34 | loss:
35 | loss_qfl:
36 | name: QualityFocalLoss
37 | use_sigmoid: True
38 | beta: 2.0
39 | loss_weight: 1.0
40 | loss_dfl:
41 | name: DistributionFocalLoss
42 | loss_weight: 0.25
43 | loss_bbox:
44 | name: GIoULoss
45 | loss_weight: 2.0
46 | data:
47 | train:
48 | name: CocoDataset
49 | img_path: /coco/train2017
50 | ann_path: /coco/annotations/instances_train2017.json
51 | input_size: [416,416] #[w,h]
52 | keep_ratio: True
53 | pipeline:
54 | perspective: 0.0
55 | scale: [0.5, 1.5]
56 | stretch: [[1, 1], [1, 1]]
57 | rotation: 0
58 | shear: 0
59 | translate: 0.2
60 | flip: 0.5
61 | brightness: 0.2
62 | contrast: [0.6, 1.4]
63 | saturation: [0.5, 1.2]
64 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]
65 | val:
66 | name: CocoDataset
67 | img_path: /coco/val2017
68 | ann_path: /coco/annotations/instances_val2017.json
69 | input_size: [416,416] #[w,h]
70 | keep_ratio: True
71 | pipeline:
72 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]
73 | device:
74 | gpu_ids: [0]
75 | workers_per_gpu: 1
76 | batchsize_per_gpu: 100
77 | schedule:
78 | # resume:
79 | # load_model: YOUR_MODEL_PATH
80 | optimizer:
81 | name: SGD
82 | lr: 0.07
83 | momentum: 0.9
84 | weight_decay: 0.0001
85 | warmup:
86 | name: linear
87 | steps: 500
88 | ratio: 0.01
89 | total_epochs: 170
90 | lr_schedule:
91 | name: MultiStepLR
92 | milestones: [130,150,160,165]
93 | gamma: 0.1
94 | val_intervals: 5
95 | evaluator:
96 | name: CocoDetectionEvaluator
97 | save_key: mAP
98 |
99 | log:
100 | interval: 10
101 |
102 | class_names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
103 | 'train', 'truck', 'boat', 'traffic_light', 'fire_hydrant',
104 | 'stop_sign', 'parking_meter', 'bench', 'bird', 'cat', 'dog',
105 | 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
106 | 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
107 | 'skis', 'snowboard', 'sports_ball', 'kite', 'baseball_bat',
108 | 'baseball_glove', 'skateboard', 'surfboard', 'tennis_racket',
109 | 'bottle', 'wine_glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
110 | 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
111 | 'hot_dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
112 | 'potted_plant', 'bed', 'dining_table', 'toilet', 'tv', 'laptop',
113 | 'mouse', 'remote', 'keyboard', 'cell_phone', 'microwave',
114 | 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
115 | 'vase', 'scissors', 'teddy_bear', 'hair_drier', 'toothbrush']
116 |
--------------------------------------------------------------------------------
/config/nanodet-m.yml:
--------------------------------------------------------------------------------
1 | #Config File example
2 | save_dir: workspace/nanodet_m
3 | model:
4 | arch:
5 | name: OneStageDetector
6 | backbone:
7 | name: ShuffleNetV2
8 | model_size: 1.0x
9 | out_stages: [2,3,4]
10 | activation: LeakyReLU
11 | fpn:
12 | name: PAN
13 | in_channels: [116, 232, 464]
14 | out_channels: 96
15 | start_level: 0
16 | num_outs: 3
17 | head:
18 | name: NanoDetHead
19 | num_classes: 80
20 | input_channel: 96
21 | feat_channels: 96
22 | stacked_convs: 2
23 | share_cls_reg: True
24 | octave_base_scale: 5
25 | scales_per_octave: 1
26 | strides: [8, 16, 32]
27 | reg_max: 7
28 | norm_cfg:
29 | type: BN
30 | loss:
31 | loss_qfl:
32 | name: QualityFocalLoss
33 | use_sigmoid: True
34 | beta: 2.0
35 | loss_weight: 1.0
36 | loss_dfl:
37 | name: DistributionFocalLoss
38 | loss_weight: 0.25
39 | loss_bbox:
40 | name: GIoULoss
41 | loss_weight: 2.0
42 | data:
43 | train:
44 | name: CocoDataset
45 | img_path: coco/train2017
46 | ann_path: coco/annotations/instances_train2017.json
47 | input_size: [320,320] #[w,h]
48 | keep_ratio: True
49 | pipeline:
50 | perspective: 0.0
51 | scale: [0.6, 1.4]
52 | stretch: [[1, 1], [1, 1]]
53 | rotation: 0
54 | shear: 0
55 | translate: 0.2
56 | flip: 0.5
57 | brightness: 0.2
58 | contrast: [0.6, 1.4]
59 | saturation: [0.5, 1.2]
60 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]
61 | val:
62 | name: CocoDataset
63 | img_path: coco/val2017
64 | ann_path: coco/annotations/instances_val2017.json
65 | input_size: [320,320] #[w,h]
66 | keep_ratio: True
67 | pipeline:
68 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]
69 | device:
70 | gpu_ids: [0]
71 | workers_per_gpu: 8
72 | batchsize_per_gpu: 192
73 | schedule:
74 | # resume:
75 | # load_model: YOUR_MODEL_PATH
76 | optimizer:
77 | name: SGD
78 | lr: 0.14
79 | momentum: 0.9
80 | weight_decay: 0.0001
81 | warmup:
82 | name: linear
83 | steps: 300
84 | ratio: 0.1
85 | total_epochs: 280
86 | lr_schedule:
87 | name: MultiStepLR
88 | milestones: [240,260,275]
89 | gamma: 0.1
90 | val_intervals: 10
91 | evaluator:
92 | name: CocoDetectionEvaluator
93 | save_key: mAP
94 |
95 | log:
96 | interval: 10
97 |
98 | class_names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
99 | 'train', 'truck', 'boat', 'traffic_light', 'fire_hydrant',
100 | 'stop_sign', 'parking_meter', 'bench', 'bird', 'cat', 'dog',
101 | 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
102 | 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
103 | 'skis', 'snowboard', 'sports_ball', 'kite', 'baseball_bat',
104 | 'baseball_glove', 'skateboard', 'surfboard', 'tennis_racket',
105 | 'bottle', 'wine_glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
106 | 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
107 | 'hot_dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
108 | 'potted_plant', 'bed', 'dining_table', 'toilet', 'tv', 'laptop',
109 | 'mouse', 'remote', 'keyboard', 'cell_phone', 'microwave',
110 | 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
111 | 'vase', 'scissors', 'teddy_bear', 'hair_drier', 'toothbrush']
112 |
--------------------------------------------------------------------------------
/config/nanodet_custom_xml_dataset.yml:
--------------------------------------------------------------------------------
1 | #Config File example
2 | save_dir: workspace/nanodet_m
3 | model:
4 | arch:
5 | name: OneStageDetector
6 | backbone:
7 | name: ShuffleNetV2
8 | model_size: 1.0x
9 | out_stages: [2,3,4]
10 | activation: LeakyReLU
11 | fpn:
12 | name: PAN
13 | in_channels: [116, 232, 464]
14 | out_channels: 96
15 | start_level: 0
16 | num_outs: 3
17 | head:
18 | name: NanoDetHead
19 | num_classes: 80 #Please fill in the number of categories (not include background category)
20 | input_channel: 96
21 | feat_channels: 96
22 | stacked_convs: 2
23 | share_cls_reg: True
24 | octave_base_scale: 5
25 | scales_per_octave: 1
26 | strides: [8, 16, 32]
27 | reg_max: 7
28 | norm_cfg:
29 | type: BN
30 | loss:
31 | loss_qfl:
32 | name: QualityFocalLoss
33 | use_sigmoid: True
34 | beta: 2.0
35 | loss_weight: 1.0
36 | loss_dfl:
37 | name: DistributionFocalLoss
38 | loss_weight: 0.25
39 | loss_bbox:
40 | name: GIoULoss
41 | loss_weight: 2.0
42 |
43 | class_names: &class_names ['NAME1', 'NAME2', 'NAME3', 'NAME4', '...'] #Please fill in the category names (not include background category)
44 | data:
45 | train:
46 | name: XMLDataset
47 | class_names: *class_names
48 | img_path: TRAIN_IMAGE_FOLDER #Please fill in train image path
49 | ann_path: TRAIN_XML_FOLDER #Please fill in train xml path
50 | input_size: [320,320] #[w,h]
51 | keep_ratio: True
52 | pipeline:
53 | perspective: 0.0
54 | scale: [0.6, 1.4]
55 | stretch: [[1, 1], [1, 1]]
56 | rotation: 0
57 | shear: 0
58 | translate: 0.2
59 | flip: 0.5
60 | brightness: 0.2
61 | contrast: [0.8, 1.2]
62 | saturation: [0.8, 1.2]
63 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]
64 | val:
65 | name: XMLDataset
66 | class_names: *class_names
67 | img_path: VAL_IMAGE_FOLDER #Please fill in val image path
68 | ann_path: VAL_XML_FOLDER #Please fill in val xml path
69 | input_size: [320,320] #[w,h]
70 | keep_ratio: True
71 | pipeline:
72 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]]
73 | device:
74 | gpu_ids: [0]
75 | workers_per_gpu: 12
76 | batchsize_per_gpu: 160
77 | schedule:
78 | # resume:
79 | # load_model: YOUR_MODEL_PATH
80 | optimizer:
81 | name: SGD
82 | lr: 0.14
83 | momentum: 0.9
84 | weight_decay: 0.0001
85 | warmup:
86 | name: linear
87 | steps: 300
88 | ratio: 0.1
89 | total_epochs: 190
90 | lr_schedule:
91 | name: MultiStepLR
92 | milestones: [130,160,175,185]
93 | gamma: 0.1
94 | val_intervals: 10
95 | evaluator:
96 | name: CocoDetectionEvaluator
97 | save_key: mAP
98 |
99 | log:
100 | interval: 10
101 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/demo_android_ncnn/Android_demo.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/Android_demo.jpg
--------------------------------------------------------------------------------
/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 | ## Step3.
17 | Copy the NanoDet ncnn model file (nanodet_m.param and nanodet_m.bin) from models folder into demo_android_ncnn/app/src/main/assets
18 |
19 | * [NanoDet ncnn model download link](https://github.com/RangiLyu/nanodet/releases/download/v0.3.0/nanodet_m_ncnn_model.zip)
20 |
21 | If you want to run yolov4-tiny and yolov5s, download them and also put in demo_android_ncnn/app/src/main/assets.
22 |
23 | * [Yolov4 and v5 ncnn model download link](https://drive.google.com/file/d/1Qk_1fDvOcFmNppDnaMFW-xFpMgLDyeAs/view?usp=sharing)
24 |
25 | ## Step4.
26 | Open demo_android_ncnn folder with Android Studio and then build it.
27 |
28 | # Screenshot
29 | 
30 |
31 |
32 | # Reference
33 |
34 | * [ncnn](https://github.com/tencent/ncnn)
35 | * [YOLOv5_NCNN](https://github.com/WZTENG/YOLOv5_NCNN)
36 |
--------------------------------------------------------------------------------
/demo_android_ncnn/app/.gitignore:
--------------------------------------------------------------------------------
1 | /build
2 |
--------------------------------------------------------------------------------
/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/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 |
--------------------------------------------------------------------------------
/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_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 |
--------------------------------------------------------------------------------
/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-20201218-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 |
--------------------------------------------------------------------------------
/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 |
20 | class NanoDet{
21 | public:
22 | NanoDet(AAssetManager *mgr, const char *param, const char *bin, bool useGPU);
23 |
24 | ~NanoDet();
25 |
26 | std::vector detect(JNIEnv *env, jobject image, float score_threshold, float nms_threshold);
27 | std::vector labels{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
28 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
29 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
30 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
31 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
32 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
33 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
34 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
35 | "hair drier", "toothbrush"};
36 | private:
37 | void preprocess(JNIEnv *env, jobject image, ncnn::Mat& in);
38 | void decode_infer(ncnn::Mat& cls_pred, ncnn::Mat& dis_pred, int stride, float threshold, std::vector>& results, float width_ratio, float height_ratio);
39 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride, float width_ratio, float height_ratio);
40 |
41 | static void nms(std::vector& result, float nms_threshold);
42 |
43 | ncnn::Net *Net;
44 | int input_size = 320;
45 | int num_class = 80;
46 | int reg_max = 7;
47 | std::vector heads_info{
48 | // cls_pred|dis_pred|stride
49 | {"cls_pred_stride_8", "dis_pred_stride_8", 8},
50 | {"cls_pred_stride_16", "dis_pred_stride_16", 16},
51 | {"cls_pred_stride_32", "dis_pred_stride_32", 32},
52 | };
53 |
54 | public:
55 | static NanoDet *detector;
56 | static bool hasGPU;
57 | };
58 |
59 |
60 | #endif //NANODET_H
61 |
--------------------------------------------------------------------------------
/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/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/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/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_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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/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/res/drawable-v24/ic_launcher_foreground.xml:
--------------------------------------------------------------------------------
1 |
7 |
12 |
13 |
19 |
22 |
25 |
26 |
27 |
28 |
34 |
35 |
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/cpu.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/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/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/gpu.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/ncnn_icon.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/ncnn_icon.png
--------------------------------------------------------------------------------
/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/layout/activity_welcome.xml:
--------------------------------------------------------------------------------
1 |
2 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
30 |
31 |
38 |
39 |
47 |
48 |
55 |
56 |
63 |
64 |
71 |
72 |
73 |
74 |
75 |
76 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher_round.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher_round.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher_round.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher_round.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher_round.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher_round.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher_round.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher_round.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher.png
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher_round.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher_round.png
--------------------------------------------------------------------------------
/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/values/strings.xml:
--------------------------------------------------------------------------------
1 |
2 | NanoDet
3 |
4 |
--------------------------------------------------------------------------------
/demo_android_ncnn/app/src/main/res/values/styles.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
10 |
16 |
17 |
18 |
--------------------------------------------------------------------------------
/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_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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/demo_android_ncnn/gradle/wrapper/gradle-wrapper.jar:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_android_ncnn/gradle/wrapper/gradle-wrapper.jar
--------------------------------------------------------------------------------
/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/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 |
--------------------------------------------------------------------------------
/demo_android_ncnn/settings.gradle:
--------------------------------------------------------------------------------
1 | include ':app'
2 | rootProject.name='NanoDet'
3 |
--------------------------------------------------------------------------------
/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_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 |
--------------------------------------------------------------------------------
/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, int stage_idx, 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 };
29 | int input_size_ = 320;
30 | int num_class_ = 80;
31 | int reg_max_ = 7;
32 |
33 | };
34 |
--------------------------------------------------------------------------------
/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_mnn/imgs/000252.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_mnn/imgs/000252.jpg
--------------------------------------------------------------------------------
/demo_mnn/imgs/000258.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_mnn/imgs/000258.jpg
--------------------------------------------------------------------------------
/demo_mnn/nanodet_mnn.hpp:
--------------------------------------------------------------------------------
1 | #ifndef __NanoDet_H__
2 | #define __NanoDet_H__
3 |
4 | #pragma once
5 |
6 | #include "Interpreter.hpp"
7 |
8 | #include "MNNDefine.h"
9 | #include "Tensor.hpp"
10 | #include "ImageProcess.hpp"
11 | #include
12 | #include
13 | #include
14 | #include
15 | #include
16 | #include
17 | #include
18 |
19 |
20 | typedef struct HeadInfo_
21 | {
22 | std::string cls_layer;
23 | std::string dis_layer;
24 | int stride;
25 | } HeadInfo;
26 |
27 | typedef struct BoxInfo_
28 | {
29 | float x1;
30 | float y1;
31 | float x2;
32 | float y2;
33 | float score;
34 | int label;
35 | } BoxInfo;
36 |
37 | class NanoDet {
38 | public:
39 | NanoDet(const std::string &mnn_path,
40 | int input_width, int input_length, int num_thread_ = 4, float score_threshold_ = 0.5, float nms_threshold_ = 0.3);
41 |
42 | ~NanoDet();
43 |
44 | int detect(cv::Mat &img, std::vector &result_list);
45 | std::string get_label_str(int label);
46 |
47 | private:
48 | void decode_infer(MNN::Tensor *cls_pred, MNN::Tensor *dis_pred, int stride, float threshold, std::vector> &results);
49 | BoxInfo disPred2Bbox(const float *&dfl_det, int label, float score, int x, int y, int stride);
50 | void nms(std::vector &input_boxes, float NMS_THRESH);
51 |
52 | private:
53 |
54 | std::shared_ptr NanoDet_interpreter;
55 | MNN::Session *NanoDet_session = nullptr;
56 | MNN::Tensor *input_tensor = nullptr;
57 |
58 | int num_thread;
59 | int image_w;
60 | int image_h;
61 |
62 | int in_w = 320;
63 | int in_h = 320;
64 |
65 | float score_threshold;
66 | float nms_threshold;
67 |
68 | const float mean_vals[3] = { 103.53f, 116.28f, 123.675f };
69 | const float norm_vals[3] = { 0.017429f, 0.017507f, 0.017125f };
70 |
71 | const int num_class = 80;
72 | const int reg_max = 7;
73 |
74 | std::vector heads_info{
75 | // cls_pred|dis_pred|stride
76 | {"cls_pred_stride_8", "dis_pred_stride_8", 8},
77 | {"cls_pred_stride_16", "dis_pred_stride_16", 16},
78 | {"cls_pred_stride_32", "dis_pred_stride_32", 32},
79 | };
80 |
81 | std::vector
82 | labels{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
83 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
84 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
85 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
86 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
87 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
88 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
89 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
90 | "hair drier", "toothbrush"};
91 | };
92 |
93 | template
94 | int activation_function_softmax(const _Tp *src, _Tp *dst, int length);
95 |
96 | inline float fast_exp(float x);
97 | inline float sigmoid(float x);
98 |
99 | #endif // __NanoDet_H__
100 |
--------------------------------------------------------------------------------
/demo_mnn/results/000252.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_mnn/results/000252.jpg
--------------------------------------------------------------------------------
/demo_mnn/results/000258.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_mnn/results/000258.jpg
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/demo_ncnn/README.md:
--------------------------------------------------------------------------------
1 | # NanoDet NCNN Demo
2 |
3 | This project provides NanoDet image inference, webcam inference and benchmark using
4 | [Tencent's NCNN framework](https://github.com/Tencent/ncnn).
5 |
6 | # How to build
7 |
8 | ## Windows
9 | ### Step1.
10 | Download and Install Visual Studio from https://visualstudio.microsoft.com/vs/community/
11 |
12 | ### Step2.
13 | Download and install OpenCV from https://github.com/opencv/opencv/releases
14 |
15 | ### Step3(Optional).
16 | Download and install Vulkan SDK from https://vulkan.lunarg.com/sdk/home
17 |
18 | ### Step4.
19 | Clone NCNN repository
20 |
21 | ``` shell script
22 | git clone --recursive https://github.com/Tencent/ncnn.git
23 | ```
24 | Build NCNN following this tutorial: [Build for Windows x64 using VS2017](https://github.com/Tencent/ncnn/wiki/how-to-build#build-for-windows-x64-using-visual-studio-community-2017)
25 |
26 | ### Step5.
27 |
28 | Add `ncnn_DIR` = `YOUR_NCNN_PATH/build/install/lib/cmake/ncnn` to system environment variables.
29 |
30 | Build project: Open x64 Native Tools Command Prompt for VS 2019 or 2017
31 |
32 | ``` cmd
33 | cd
34 | mkdir -p build
35 | cd build
36 | cmake ..
37 | msbuild nanodet_demo.vcxproj /p:configuration=release /p:platform=x64
38 | ```
39 |
40 | ## Linux
41 |
42 | ### Step1.
43 | Build and install OpenCV from https://github.com/opencv/opencv
44 |
45 | ### Step2(Optional).
46 | Download Vulkan SDK from https://vulkan.lunarg.com/sdk/home
47 |
48 | ### Step3.
49 | Clone NCNN repository
50 |
51 | ``` shell script
52 | git clone --recursive https://github.com/Tencent/ncnn.git
53 | ```
54 |
55 | Build NCNN following this tutorial: [Build for Linux / NVIDIA Jetson / Raspberry Pi](https://github.com/Tencent/ncnn/wiki/how-to-build#build-for-linux)
56 |
57 | ### Step4.
58 |
59 | Set environment variables. Run:
60 |
61 | ``` shell script
62 | export ncnn_DIR=YOUR_NCNN_PATH/build/install/lib/cmake/ncnn
63 | ```
64 |
65 | Build project
66 |
67 | ``` shell script
68 | cd
69 | mkdir build
70 | cd build
71 | cmake ..
72 | make
73 | ```
74 |
75 | # Run demo
76 |
77 | Download NanoDet ncnn model.
78 | * [NanoDet ncnn model download link](https://github.com/RangiLyu/nanodet/releases/download/v0.3.0/nanodet_m_ncnn_model.zip)
79 |
80 | Copy nanodet_m.param and nanodet_m.bin to demo program folder.
81 |
82 | ## Webcam
83 |
84 | ```shell script
85 | nanodet_demo 0 0
86 | ```
87 |
88 | ## Inference images
89 |
90 | ```shell script
91 | nanodet_demo 1 IMAGE_FOLDER/*.jpg
92 | ```
93 |
94 | ## Inference video
95 |
96 | ```shell script
97 | nanodet_demo 2 VIDEO_PATH
98 | ```
99 |
100 | ## Benchmark
101 |
102 | ```shell script
103 | nanodet_demo 3 0
104 | ```
105 | 
106 | ****
107 |
108 | Notice:
109 |
110 | If benchmark speed is slow, try to limit omp thread num.
111 |
112 | Linux:
113 |
114 | ```shell script
115 | export OMP_THREAD_LIMIT=4
116 | ```
117 |
--------------------------------------------------------------------------------
/demo_ncnn/benchmark.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/demo_ncnn/benchmark.jpg
--------------------------------------------------------------------------------
/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 | typedef struct BoxInfo
20 | {
21 | float x1;
22 | float y1;
23 | float x2;
24 | float y2;
25 | float score;
26 | int label;
27 | } BoxInfo;
28 |
29 | class NanoDet
30 | {
31 | public:
32 | NanoDet(const char* param, const char* bin, bool useGPU);
33 |
34 | ~NanoDet();
35 |
36 | static NanoDet* detector;
37 | ncnn::Net* Net;
38 | static bool hasGPU;
39 |
40 | std::vector heads_info{
41 | // cls_pred|dis_pred|stride
42 | {"cls_pred_stride_8", "dis_pred_stride_8", 8},
43 | {"cls_pred_stride_16", "dis_pred_stride_16", 16},
44 | {"cls_pred_stride_32", "dis_pred_stride_32", 32},
45 | };
46 |
47 | std::vector detect(cv::Mat image, float score_threshold, float nms_threshold);
48 |
49 | std::vector labels{ "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
50 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
51 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
52 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
53 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
54 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
55 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
56 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
57 | "hair drier", "toothbrush" };
58 | private:
59 | void preprocess(cv::Mat& image, ncnn::Mat& in);
60 | void decode_infer(ncnn::Mat& cls_pred, ncnn::Mat& dis_pred, int stride, float threshold, std::vector>& results);
61 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride);
62 | static void nms(std::vector& result, float nms_threshold);
63 | int input_size[2] = {320, 320};
64 | int num_class = 80;
65 | int reg_max = 7;
66 |
67 | };
68 |
69 |
70 | #endif //NANODET_H
71 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/demo_openvino/README.md:
--------------------------------------------------------------------------------
1 | # NanoDet OpenVINO Demo
2 |
3 | This fold provides NanoDet inference code using
4 | [Intel's OpenVINO Toolkit](https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit.html). Most of the implements in this fold are same as *demo_ncnn*.
5 |
6 | ## Install OpenVINO Toolkit
7 |
8 | Go to [OpenVINO HomePage](https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit.html)
9 |
10 | Download a suitable version and install.
11 |
12 | Follow the official Get Started Guides: https://docs.openvinotoolkit.org/latest/get_started_guides.html
13 |
14 | ## Set the Environment Variables
15 |
16 | ### Windows:
17 |
18 | Run this command in cmd. (Every time before using OpenVINO)
19 | ```cmd
20 | \openvino_2021\bin\setupvars.bat
21 | ```
22 |
23 |
24 | Or set the system environment variables once for all:
25 |
26 | Name |Value
27 | :--------------------:|:--------:
28 | INTEL_OPENVINO_DIR | \openvino_2021
29 | INTEL_CVSDK_DIR | %INTEL_OPENVINO_DIR%
30 | InferenceEngine_DIR | %INTEL_OPENVINO_DIR%\deployment_tools\inference_engine\share
31 | HDDL_INSTALL_DIR | %INTEL_OPENVINO_DIR%\deployment_tools\inference_engine\external\hddl
32 | ngraph_DIR | %INTEL_OPENVINO_DIR%\deployment_tools\ngraph\cmake
33 |
34 | And add this to ```Path```
35 | ```
36 | %INTEL_OPENVINO_DIR%\deployment_tools\inference_engine\bin\intel64\Debug;%INTEL_OPENVINO_DIR%\deployment_tools\inference_engine\bin\intel64\Release;%HDDL_INSTALL_DIR%\bin;%INTEL_OPENVINO_DIR%\deployment_tools\inference_engine\external\tbb\bin;%INTEL_OPENVINO_DIR%\deployment_tools\ngraph\lib
37 | ```
38 |
39 | ### Linux
40 |
41 | Run this command in shell. (Every time before using OpenVINO)
42 |
43 | ```shell
44 | source /opt/intel/openvino_2021/bin/setupvars.sh
45 | ```
46 |
47 | Or edit .bashrc
48 |
49 | ```shell
50 | vi ~/.bashrc
51 | ```
52 |
53 | Add this line to the end of the file
54 |
55 | ```shell
56 | source /opt/intel/openvino_2021/bin/setupvars.sh
57 | ```
58 |
59 | ## Convert model
60 |
61 | 1. Export ONNX model
62 |
63 | ```shell
64 | python ./tools/export_onnx.py --cfg_path ${CONFIG_PATH} --model_path ${PYTORCH_MODEL_PATH}
65 | ```
66 |
67 | 2. Use *onnx-simplifier* to simplify it
68 |
69 | ``` shell
70 | python -m onnxsim ${INPUT_ONNX_MODEL} ${OUTPUT_ONNX_MODEL}
71 | ```
72 |
73 | 3. Convert to OpenVINO
74 |
75 | ``` shell
76 | cd /openvino_2021/deployment_tools/model_optimizer
77 | ```
78 |
79 | Install requirements for convert tool
80 |
81 | ```shell
82 | sudo ./install_prerequisites/install_prerequisites_onnx.sh
83 | ```
84 |
85 | Then convert model. Notice: mean_values and scale_values should be the same with your training settings in YAML config file.
86 | ```shell
87 | python3 mo_onnx.py --input_model --mean_values [103.53,116.28,123.675] --scale_values [57.375,57.12,58.395]
88 | ```
89 |
90 | ## Build
91 |
92 | ### Windows
93 |
94 | ```cmd
95 | \openvino_2021\bin\setupvars.bat
96 | mkdir -p build
97 | cd build
98 | cmake ..
99 | msbuild nanodet_demo.vcxproj /p:configuration=release /p:platform=x64
100 | ```
101 |
102 | ### Linux
103 | ```shell
104 | source /opt/intel/openvino_2021/bin/setupvars.sh
105 | mkdir build
106 | cd build
107 | cmake ..
108 | make
109 | ```
110 |
111 |
112 | ## Run demo
113 |
114 | First, move nanodet openvino model files to the demo's folder. Then run these commands:
115 |
116 | ### Webcam
117 |
118 | ```shell
119 | nanodet_demo 0 0
120 | ```
121 |
122 | ### Inference images
123 |
124 | ```shell
125 | nanodet_demo 1 IMAGE_FOLDER/*.jpg
126 | ```
127 |
128 | ### Inference video
129 |
130 | ```shell
131 | nanodet_demo 2 VIDEO_PATH
132 | ```
133 |
134 | ### Benchmark
135 |
136 | ```shell
137 | nanodet_demo 3 0
138 | ```
139 |
--------------------------------------------------------------------------------
/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 | typedef struct BoxInfo
22 | {
23 | float x1;
24 | float y1;
25 | float x2;
26 | float y2;
27 | float score;
28 | int label;
29 | } BoxInfo;
30 |
31 | class NanoDet
32 | {
33 | public:
34 | NanoDet(const char* param);
35 |
36 | ~NanoDet();
37 |
38 | InferenceEngine::ExecutableNetwork network_;
39 | InferenceEngine::InferRequest infer_request_;
40 | // static bool hasGPU;
41 |
42 | std::vector heads_info_{
43 | // cls_pred|dis_pred|stride
44 | {"cls_pred_stride_8", "dis_pred_stride_8", 8},
45 | {"cls_pred_stride_16", "dis_pred_stride_16", 16},
46 | {"cls_pred_stride_32", "dis_pred_stride_32", 32},
47 | };
48 |
49 | std::vector detect(cv::Mat image, float score_threshold, float nms_threshold);
50 |
51 | private:
52 | void preprocess(cv::Mat& image, InferenceEngine::Blob::Ptr& blob);
53 | void decode_infer(const float*& cls_pred, const float*& dis_pred, int stride, float threshold, std::vector>& results);
54 | BoxInfo disPred2Bbox(const float*& dfl_det, int label, float score, int x, int y, int stride);
55 | static void nms(std::vector& result, float nms_threshold);
56 | std::string input_name_;
57 | int input_size_ = 320;
58 | int num_class_ = 80;
59 | int reg_max_ = 7;
60 |
61 | };
62 |
63 |
64 | #endif //_NANODE_TOPENVINO_H_
65 |
--------------------------------------------------------------------------------
/docs/imgs/Android_demo.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/docs/imgs/Android_demo.jpg
--------------------------------------------------------------------------------
/docs/imgs/Model_arch.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/docs/imgs/Model_arch.png
--------------------------------------------------------------------------------
/docs/imgs/Title.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/docs/imgs/Title.jpg
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/nanodet/__about__.py:
--------------------------------------------------------------------------------
1 | import time
2 |
3 | _this_year = time.strftime("%Y")
4 | __version__ = "0.4.2"
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 |
--------------------------------------------------------------------------------
/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/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/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 |
--------------------------------------------------------------------------------
/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 |
21 |
22 | def build_dataset(cfg, mode):
23 | dataset_cfg = copy.deepcopy(cfg)
24 | name = dataset_cfg.pop("name")
25 | if name == "coco":
26 | warnings.warn(
27 | "Dataset name coco has been deprecated. Please use CocoDataset instead."
28 | )
29 | return CocoDataset(mode=mode, **dataset_cfg)
30 | elif name == "xml_dataset":
31 | warnings.warn(
32 | "Dataset name xml_dataset has been deprecated. "
33 | "Please use XMLDataset instead."
34 | )
35 | return XMLDataset(mode=mode, **dataset_cfg)
36 | elif name == "CocoDataset":
37 | return CocoDataset(mode=mode, **dataset_cfg)
38 | elif name == "XMLDataset":
39 | return XMLDataset(mode=mode, **dataset_cfg)
40 | else:
41 | raise NotImplementedError("Unknown dataset type!")
42 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/nanodet/data/transform/mosaic.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/nanodet/data/transform/mosaic.py
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/evaluator/coco_detection.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 json
17 | import os
18 | import warnings
19 |
20 | from pycocotools.cocoeval import COCOeval
21 |
22 |
23 | def xyxy2xywh(bbox):
24 | """
25 | change bbox to coco format
26 | :param bbox: [x1, y1, x2, y2]
27 | :return: [x, y, w, h]
28 | """
29 | return [
30 | bbox[0],
31 | bbox[1],
32 | bbox[2] - bbox[0],
33 | bbox[3] - bbox[1],
34 | ]
35 |
36 |
37 | class CocoDetectionEvaluator:
38 | def __init__(self, dataset):
39 | assert hasattr(dataset, "coco_api")
40 | self.coco_api = dataset.coco_api
41 | self.cat_ids = dataset.cat_ids
42 | self.metric_names = ["mAP", "AP_50", "AP_75", "AP_small", "AP_m", "AP_l"]
43 |
44 | def results2json(self, results):
45 | """
46 | results: {image_id: {label: [bboxes...] } }
47 | :return coco json format: {image_id:
48 | category_id:
49 | bbox:
50 | score: }
51 | """
52 | json_results = []
53 | for image_id, dets in results.items():
54 | for label, bboxes in dets.items():
55 | category_id = self.cat_ids[label]
56 | for bbox in bboxes:
57 | score = float(bbox[4])
58 | detection = dict(
59 | image_id=int(image_id),
60 | category_id=int(category_id),
61 | bbox=xyxy2xywh(bbox),
62 | score=score,
63 | )
64 | json_results.append(detection)
65 | return json_results
66 |
67 | def evaluate(self, results, save_dir, rank=-1):
68 | results_json = self.results2json(results)
69 | if len(results_json) == 0:
70 | warnings.warn(
71 | "Detection result is empty! Please check whether "
72 | "training set is too small (need to increase val_interval "
73 | "in config and train more epochs). Or check annotation "
74 | "correctness."
75 | )
76 | empty_eval_results = {}
77 | for key in self.metric_names:
78 | empty_eval_results[key] = 0
79 | return empty_eval_results
80 | json_path = os.path.join(save_dir, "results{}.json".format(rank))
81 | json.dump(results_json, open(json_path, "w"))
82 | coco_dets = self.coco_api.loadRes(json_path)
83 | coco_eval = COCOeval(
84 | copy.deepcopy(self.coco_api), copy.deepcopy(coco_dets), "bbox"
85 | )
86 | coco_eval.evaluate()
87 | coco_eval.accumulate()
88 | coco_eval.summarize()
89 | aps = coco_eval.stats[:6]
90 | eval_results = {}
91 | for k, v in zip(self.metric_names, aps):
92 | eval_results[k] = v
93 | return eval_results
94 |
--------------------------------------------------------------------------------
/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 warnings
16 |
17 | from .one_stage_detector import OneStageDetector
18 |
19 |
20 | def build_model(model_cfg):
21 | if model_cfg.arch.name == "GFL":
22 | warnings.warn(
23 | "Model architecture name is changed to 'OneStageDetector'. "
24 | "The name 'GFL' is deprecated, please change the model->arch->name "
25 | "in your YAML config file to OneStageDetector."
26 | )
27 | model = OneStageDetector(
28 | model_cfg.arch.backbone, model_cfg.arch.fpn, model_cfg.arch.head
29 | )
30 | elif model_cfg.arch.name == "OneStageDetector":
31 | model = OneStageDetector(
32 | model_cfg.arch.backbone, model_cfg.arch.fpn, model_cfg.arch.head
33 | )
34 | else:
35 | raise NotImplementedError
36 | return model
37 |
--------------------------------------------------------------------------------
/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 |
39 | def forward(self, x):
40 | x = self.backbone(x)
41 | if hasattr(self, "fpn"):
42 | x = self.fpn(x)
43 | if hasattr(self, "head"):
44 | x = self.head(x)
45 | return x
46 |
47 | def inference(self, meta):
48 | with torch.no_grad():
49 | torch.cuda.synchronize()
50 | time1 = time.time()
51 | preds = self(meta["img"])
52 | torch.cuda.synchronize()
53 | time2 = time.time()
54 | print("forward time: {:.3f}s".format((time2 - time1)), end=" | ")
55 | results = self.head.post_process(preds, meta)
56 | torch.cuda.synchronize()
57 | print("decode time: {:.3f}s".format((time.time() - time2)), end=" | ")
58 | return results
59 |
60 | def forward_train(self, gt_meta):
61 | preds = self(gt_meta["img"])
62 | loss, loss_states = self.head.loss(preds, gt_meta)
63 |
64 | return preds, loss, loss_states
65 |
--------------------------------------------------------------------------------
/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 |
25 |
26 | def build_backbone(cfg):
27 | backbone_cfg = copy.deepcopy(cfg)
28 | name = backbone_cfg.pop("name")
29 | if name == "ResNet":
30 | return ResNet(**backbone_cfg)
31 | elif name == "ShuffleNetV2":
32 | return ShuffleNetV2(**backbone_cfg)
33 | elif name == "GhostNet":
34 | return GhostNet(**backbone_cfg)
35 | elif name == "MobileNetV2":
36 | return MobileNetV2(**backbone_cfg)
37 | elif name == "EfficientNetLite":
38 | return EfficientNetLite(**backbone_cfg)
39 | elif name == "CustomCspNet":
40 | return CustomCspNet(**backbone_cfg)
41 | elif name == "RepVGG":
42 | return RepVGG(**backbone_cfg)
43 | else:
44 | raise NotImplementedError
45 |
--------------------------------------------------------------------------------
/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 .pan import PAN
19 | from .tan import TAN
20 |
21 |
22 | def build_fpn(cfg):
23 | fpn_cfg = copy.deepcopy(cfg)
24 | name = fpn_cfg.pop("name")
25 | if name == "FPN":
26 | return FPN(**fpn_cfg)
27 | elif name == "PAN":
28 | return PAN(**fpn_cfg)
29 | elif name == "TAN":
30 | return TAN(**fpn_cfg)
31 | else:
32 | raise NotImplementedError
33 |
--------------------------------------------------------------------------------
/nanodet/model/fpn/fpn.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 | import torch.nn as nn
17 | import torch.nn.functional as F
18 |
19 | from ..module.conv import ConvModule
20 | from ..module.init_weights import xavier_init
21 |
22 |
23 | class FPN(nn.Module):
24 | def __init__(
25 | self,
26 | in_channels,
27 | out_channels,
28 | num_outs,
29 | start_level=0,
30 | end_level=-1,
31 | conv_cfg=None,
32 | norm_cfg=None,
33 | activation=None,
34 | ):
35 | super(FPN, self).__init__()
36 | assert isinstance(in_channels, list)
37 | self.in_channels = in_channels
38 | self.out_channels = out_channels
39 | self.num_ins = len(in_channels)
40 | self.num_outs = num_outs
41 | self.fp16_enabled = False
42 |
43 | if end_level == -1:
44 | self.backbone_end_level = self.num_ins
45 | assert num_outs >= self.num_ins - start_level
46 | else:
47 | # if end_level < inputs, no extra level is allowed
48 | self.backbone_end_level = end_level
49 | assert end_level <= len(in_channels)
50 | assert num_outs == end_level - start_level
51 | self.start_level = start_level
52 | self.end_level = end_level
53 | self.lateral_convs = nn.ModuleList()
54 |
55 | for i in range(self.start_level, self.backbone_end_level):
56 | l_conv = ConvModule(
57 | in_channels[i],
58 | out_channels,
59 | 1,
60 | conv_cfg=conv_cfg,
61 | norm_cfg=norm_cfg,
62 | activation=activation,
63 | inplace=False,
64 | )
65 |
66 | self.lateral_convs.append(l_conv)
67 | self.init_weights()
68 |
69 | # default init_weights for conv(msra) and norm in ConvModule
70 | def init_weights(self):
71 | for m in self.modules():
72 | if isinstance(m, nn.Conv2d):
73 | xavier_init(m, distribution="uniform")
74 |
75 | def forward(self, inputs):
76 | assert len(inputs) == len(self.in_channels)
77 |
78 | # build laterals
79 | laterals = [
80 | lateral_conv(inputs[i + self.start_level])
81 | for i, lateral_conv in enumerate(self.lateral_convs)
82 | ]
83 |
84 | # build top-down path
85 | used_backbone_levels = len(laterals)
86 | for i in range(used_backbone_levels - 1, 0, -1):
87 | laterals[i - 1] += F.interpolate(
88 | laterals[i], scale_factor=2, mode="bilinear"
89 | )
90 |
91 | # build outputs
92 | outs = [
93 | # self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
94 | laterals[i]
95 | for i in range(used_backbone_levels)
96 | ]
97 | return tuple(outs)
98 |
99 |
100 | # if __name__ == '__main__':
101 |
--------------------------------------------------------------------------------
/nanodet/model/fpn/pan.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 | import torch.nn.functional as F
17 |
18 | from .fpn import FPN
19 |
20 |
21 | class PAN(FPN):
22 | """Path Aggregation Network for Instance Segmentation.
23 |
24 | This is an implementation of the `PAN in Path Aggregation Network
25 | `_.
26 |
27 | Args:
28 | in_channels (List[int]): Number of input channels per scale.
29 | out_channels (int): Number of output channels (used at each scale)
30 | num_outs (int): Number of output scales.
31 | start_level (int): Index of the start input backbone level used to
32 | build the feature pyramid. Default: 0.
33 | end_level (int): Index of the end input backbone level (exclusive) to
34 | build the feature pyramid. Default: -1, which means the last level.
35 | conv_cfg (dict): Config dict for convolution layer. Default: None.
36 | norm_cfg (dict): Config dict for normalization layer. Default: None.
37 | activation (str): Config dict for activation layer in ConvModule.
38 | Default: None.
39 | """
40 |
41 | def __init__(
42 | self,
43 | in_channels,
44 | out_channels,
45 | num_outs,
46 | start_level=0,
47 | end_level=-1,
48 | conv_cfg=None,
49 | norm_cfg=None,
50 | activation=None,
51 | ):
52 | super(PAN, self).__init__(
53 | in_channels,
54 | out_channels,
55 | num_outs,
56 | start_level,
57 | end_level,
58 | conv_cfg,
59 | norm_cfg,
60 | activation,
61 | )
62 | self.init_weights()
63 |
64 | def forward(self, inputs):
65 | """Forward function."""
66 | assert len(inputs) == len(self.in_channels)
67 |
68 | # build laterals
69 | laterals = [
70 | lateral_conv(inputs[i + self.start_level])
71 | for i, lateral_conv in enumerate(self.lateral_convs)
72 | ]
73 |
74 | # build top-down path
75 | used_backbone_levels = len(laterals)
76 | for i in range(used_backbone_levels - 1, 0, -1):
77 | laterals[i - 1] += F.interpolate(
78 | laterals[i], scale_factor=2, mode="bilinear"
79 | )
80 |
81 | # build outputs
82 | # part 1: from original levels
83 | inter_outs = [laterals[i] for i in range(used_backbone_levels)]
84 |
85 | # part 2: add bottom-up path
86 | for i in range(0, used_backbone_levels - 1):
87 | inter_outs[i + 1] += F.interpolate(
88 | inter_outs[i], scale_factor=0.5, mode="bilinear"
89 | )
90 |
91 | outs = []
92 | outs.append(inter_outs[0])
93 | outs.extend([inter_outs[i] for i in range(1, used_backbone_levels)])
94 | return tuple(outs)
95 |
--------------------------------------------------------------------------------
/nanodet/model/head/__init__.py:
--------------------------------------------------------------------------------
1 | import copy
2 |
3 | from .gfl_head import GFLHead
4 | from .nanodet_head import NanoDetHead
5 |
6 |
7 | def build_head(cfg):
8 | head_cfg = copy.deepcopy(cfg)
9 | name = head_cfg.pop("name")
10 | if name == "GFLHead":
11 | return GFLHead(**head_cfg)
12 | elif name == "NanoDetHead":
13 | return NanoDetHead(**head_cfg)
14 | else:
15 | raise NotImplementedError
16 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/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 | None: nn.Identity,
26 | }
27 |
28 |
29 | def act_layers(name):
30 | assert name in activations.keys()
31 | if name == "LeakyReLU":
32 | return nn.LeakyReLU(negative_slope=0.1, inplace=True)
33 | elif name == "GELU":
34 | return nn.GELU()
35 | elif name == "PReLU":
36 | return nn.PReLU()
37 | else:
38 | return activations[name](inplace=True)
39 |
--------------------------------------------------------------------------------
/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/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
15 | import torch
16 |
17 | from .dist_trainer import DistTrainer
18 | from .trainer import Trainer
19 |
20 |
21 | def build_trainer(rank, cfg, model, logger, device="cuda"):
22 | if len(cfg.device.gpu_ids) > 1:
23 | trainer = DistTrainer(rank, cfg, model, logger)
24 | trainer.set_device(
25 | cfg.device.batchsize_per_gpu, rank, device=device
26 | ) # TODO: device
27 | else:
28 | trainer = Trainer(rank, cfg, model, logger)
29 | trainer.set_device(
30 | cfg.device.batchsize_per_gpu,
31 | cfg.device.gpu_ids,
32 | device=device,
33 | )
34 | return trainer
35 |
--------------------------------------------------------------------------------
/nanodet/trainer/dist_trainer.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.distributed as dist
16 |
17 | from ..util import DDP
18 | from .trainer import Trainer
19 |
20 |
21 | def average_gradients(model):
22 | """Gradient averaging."""
23 | size = float(dist.get_world_size())
24 | for param in model.parameters():
25 | if param.grad is not None:
26 | dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
27 | param.grad.data /= size
28 |
29 |
30 | class DistTrainer(Trainer):
31 | """
32 | Distributed trainer for multi-gpu training. (not finish yet)
33 | """
34 |
35 | def run_step(self, model, batch, mode="train"):
36 | batch = self._preprocess_batch_input(batch)
37 | output, loss, loss_stats = model.module.forward_train(batch)
38 | loss = loss.mean()
39 | if mode == "train":
40 | self.optimizer.zero_grad()
41 | loss.backward()
42 | average_gradients(model)
43 | self.optimizer.step()
44 | return output, loss, loss_stats
45 |
46 | def set_device(self, batch_per_gpu, rank, device):
47 | """
48 | Set model device for Distributed-Data-Parallel
49 | :param batch_per_gpu: batch size of each gpu
50 | :param rank: distributed training process rank
51 | :param device: cuda
52 | """
53 | assert device == "cuda"
54 | self.rank = rank
55 | self.model = DDP(
56 | batch_per_gpu,
57 | module=self.model.cuda(),
58 | device_ids=[rank],
59 | output_device=rank,
60 | )
61 | self.device = device
62 |
--------------------------------------------------------------------------------
/nanodet/util/__init__.py:
--------------------------------------------------------------------------------
1 | from .box_transform import bbox2distance, distance2bbox
2 | from .check_point import convert_old_model, load_model_weight, save_model
3 | from .config import cfg, load_config
4 | from .data_parallel import DataParallel
5 | from .distributed_data_parallel import DDP
6 | from .flops_counter import get_model_complexity_info
7 | from .logger import AverageMeter, Logger, MovingAverage, NanoDetLightningLogger
8 | from .misc import images_to_levels, multi_apply, unmap
9 | from .path import collect_files, mkdir
10 | from .rank_filter import rank_filter
11 | from .scatter_gather import gather_results, scatter_kwargs
12 | from .util_mixins import NiceRepr
13 | from .visualization import Visualizer, overlay_bbox_cv
14 |
15 | __all__ = [
16 | "distance2bbox",
17 | "bbox2distance",
18 | "convert_old_model",
19 | "load_model_weight",
20 | "save_model",
21 | "cfg",
22 | "load_config",
23 | "DataParallel",
24 | "DDP",
25 | "get_model_complexity_info",
26 | "AverageMeter",
27 | "Logger",
28 | "MovingAverage",
29 | "images_to_levels",
30 | "multi_apply",
31 | "unmap",
32 | "mkdir",
33 | "rank_filter",
34 | "gather_results",
35 | "scatter_kwargs",
36 | "NiceRepr",
37 | "Visualizer",
38 | "overlay_bbox_cv",
39 | "collect_files",
40 | "NanoDetLightningLogger",
41 | ]
42 |
--------------------------------------------------------------------------------
/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/check_point.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 collections import OrderedDict
16 |
17 | import pytorch_lightning as pl
18 | import torch
19 |
20 | from .rank_filter import rank_filter
21 |
22 |
23 | def load_model_weight(model, checkpoint, logger):
24 | state_dict = checkpoint["state_dict"]
25 | # strip prefix of state_dict
26 | if list(state_dict.keys())[0].startswith("module."):
27 | state_dict = {k[7:]: v for k, v in checkpoint["state_dict"].items()}
28 | if list(state_dict.keys())[0].startswith("model."):
29 | state_dict = {k[6:]: v for k, v in checkpoint["state_dict"].items()}
30 |
31 | model_state_dict = (
32 | model.module.state_dict() if hasattr(model, "module") else model.state_dict()
33 | )
34 |
35 | # check loaded parameters and created model parameters
36 | for k in state_dict:
37 | if k in model_state_dict:
38 | if state_dict[k].shape != model_state_dict[k].shape:
39 | logger.log(
40 | "Skip loading parameter {}, required shape{}, "
41 | "loaded shape{}.".format(
42 | k, model_state_dict[k].shape, state_dict[k].shape
43 | )
44 | )
45 | state_dict[k] = model_state_dict[k]
46 | else:
47 | logger.log("Drop parameter {}.".format(k))
48 | for k in model_state_dict:
49 | if not (k in state_dict):
50 | logger.log("No param {}.".format(k))
51 | state_dict[k] = model_state_dict[k]
52 | model.load_state_dict(state_dict, strict=False)
53 |
54 |
55 | @rank_filter
56 | def save_model(model, path, epoch, iter, optimizer=None):
57 | model_state_dict = (
58 | model.module.state_dict() if hasattr(model, "module") else model.state_dict()
59 | )
60 | data = {"epoch": epoch, "state_dict": model_state_dict, "iter": iter}
61 | if optimizer is not None:
62 | data["optimizer"] = optimizer.state_dict()
63 |
64 | torch.save(data, path)
65 |
66 |
67 | def convert_old_model(old_model_dict):
68 | if "pytorch-lightning_version" in old_model_dict:
69 | raise ValueError("This model is not old format. No need to convert!")
70 | version = pl.__version__
71 | epoch = old_model_dict["epoch"]
72 | global_step = old_model_dict["iter"]
73 | state_dict = old_model_dict["state_dict"]
74 | new_state_dict = OrderedDict()
75 | for name, value in state_dict.items():
76 | new_state_dict["model." + name] = value
77 |
78 | new_checkpoint = {
79 | "epoch": epoch,
80 | "global_step": global_step,
81 | "pytorch-lightning_version": version,
82 | "state_dict": new_state_dict,
83 | "lr_schedulers": [],
84 | }
85 |
86 | if "optimizer" in old_model_dict:
87 | optimizer_states = [old_model_dict["optimizer"]]
88 | new_checkpoint["optimizer_states"] = optimizer_states
89 |
90 | return new_checkpoint
91 |
--------------------------------------------------------------------------------
/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()
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 | # train
18 | cfg.schedule = CfgNode(new_allowed=True)
19 |
20 | # logger
21 | cfg.log = CfgNode()
22 | cfg.log.interval = 50
23 |
24 | # testing
25 | cfg.test = CfgNode()
26 | # size of images for each device
27 |
28 |
29 | def load_config(cfg, args_cfg):
30 | cfg.defrost()
31 | cfg.merge_from_file(args_cfg)
32 | cfg.freeze()
33 |
34 |
35 | if __name__ == "__main__":
36 | import sys
37 |
38 | with open(sys.argv[1], "w") as f:
39 | print(cfg, file=f)
40 |
--------------------------------------------------------------------------------
/nanodet/util/distributed_data_parallel.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 torch.nn.parallel import DistributedDataParallel
16 |
17 | from .scatter_gather import scatter_kwargs
18 |
19 |
20 | class DDP(DistributedDataParallel):
21 | def __init__(self, batchsize, **kwargs):
22 | self.batchsize = batchsize
23 | super(DDP, self).__init__(**kwargs)
24 |
25 | def scatter(self, inputs, kwargs, device_ids):
26 | return scatter_kwargs(
27 | inputs, kwargs, device_ids, dim=self.dim, chunk_sizes=[self.batchsize]
28 | )
29 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | Cython
2 | matplotlib
3 | numpy
4 | opencv-python
5 | pyaml
6 | pycocotools
7 | pytorch-lightning>=1.2.5
8 | tensorboard
9 | termcolor
10 | torch>=1.6
11 | torchmetrics
12 | torchvision
13 | tqdm
14 |
--------------------------------------------------------------------------------
/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/data/batched_nms_data.pkl:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/tests/data/batched_nms_data.pkl
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/tests/data/test_img.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/nihui/nanodet/0f32d4a7d2817af30fe1a7139477fbadab165fa3/tests/data/test_img.jpg
--------------------------------------------------------------------------------
/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 |
26 |
38 |
50 |
62 |
63 |
--------------------------------------------------------------------------------
/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_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_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_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 |
--------------------------------------------------------------------------------
/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_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 |
--------------------------------------------------------------------------------
/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_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_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_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 |
--------------------------------------------------------------------------------
/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_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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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_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_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 |
--------------------------------------------------------------------------------
/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(gt_bboxes=[np.random.random((0, 4))], gt_labels=[np.array([])])
32 | loss, empty_gt_losses = head.loss(preds, meta)
33 | # When there is no truth, the cls loss should be nonzero but there should
34 | # be no box loss.
35 | empty_qfl_loss = empty_gt_losses["loss_qfl"]
36 | empty_box_loss = empty_gt_losses["loss_bbox"]
37 | empty_dfl_loss = empty_gt_losses["loss_dfl"]
38 | assert empty_qfl_loss.item() == 0
39 | assert (
40 | empty_box_loss.item() == 0
41 | ), "there should be no box loss when there are no true boxes"
42 | assert (
43 | empty_dfl_loss.item() == 0
44 | ), "there should be no dfl loss when there are no true boxes"
45 |
46 | # When truth is non-empty then both cls and box loss should be nonzero for
47 | # random inputs
48 | gt_bboxes = [
49 | np.array([[23.6667, 23.8757, 238.6326, 151.8874]], dtype=np.float32),
50 | ]
51 | gt_labels = [np.array([2])]
52 | meta = dict(gt_bboxes=gt_bboxes, gt_labels=gt_labels)
53 | loss, one_gt_losses = head.loss(preds, meta)
54 | onegt_qfl_loss = one_gt_losses["loss_qfl"]
55 | onegt_box_loss = one_gt_losses["loss_bbox"]
56 | onegt_dfl_loss = one_gt_losses["loss_dfl"]
57 | assert onegt_qfl_loss.item() > 0, "qfl loss should be non-zero"
58 | assert onegt_box_loss.item() > 0, "box loss should be non-zero"
59 | assert onegt_dfl_loss.item() > 0, "dfl loss should be non-zero"
60 |
--------------------------------------------------------------------------------
/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 | cls_preds, reg_preds = head.forward(feat)
31 | for cls, reg, stride in zip(cls_preds, reg_preds, [8, 16, 32]):
32 | assert cls.shape == (1, 80, 320 // stride, 320 // stride)
33 | assert reg.shape == (1, (8 + 1) * 4, 320 // stride, 320 // stride)
34 |
35 | head_cfg = dict(
36 | name="NanoDetHead",
37 | num_classes=20,
38 | input_channel=1,
39 | feat_channels=96,
40 | stacked_convs=2,
41 | conv_type="Conv",
42 | reg_max=5,
43 | share_cls_reg=False,
44 | strides=[8, 16, 32],
45 | loss=dict(
46 | loss_qfl=dict(
47 | name="QualityFocalLoss", use_sigmoid=True, beta=2.0, loss_weight=1.0
48 | ),
49 | loss_dfl=dict(name="DistributionFocalLoss", loss_weight=0.25),
50 | loss_bbox=dict(name="GIoULoss", loss_weight=2.0),
51 | ),
52 | )
53 | cfg = CfgNode(head_cfg)
54 | head = build_head(cfg)
55 |
56 | cls_preds, reg_preds = head.forward(feat)
57 | for cls, reg, stride in zip(cls_preds, reg_preds, [8, 16, 32]):
58 | assert cls.shape == (1, 20, 320 // stride, 320 // stride)
59 | assert reg.shape == (1, (5 + 1) * 4, 320 // stride, 320 // stride)
60 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/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/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_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_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, 3, 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 == 3
54 | assert layer.num_groups == cfg["num_groups"]
55 | elif type_name != "LN":
56 | assert layer.num_features == 3
57 |
--------------------------------------------------------------------------------
/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_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_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 |
--------------------------------------------------------------------------------
/tests/test_trainer/test_lightning_task.py:
--------------------------------------------------------------------------------
1 | import tempfile
2 |
3 | import numpy as np
4 | import torch
5 | import torch.nn as nn
6 |
7 | from nanodet.trainer.task import TrainingTask
8 | from nanodet.util import NanoDetLightningLogger, cfg, load_config
9 |
10 |
11 | class DummyTrainer(nn.Module):
12 | current_epoch = 0
13 | global_step = 0
14 | local_rank = 0
15 | use_ddp = False
16 | logger = NanoDetLightningLogger(tempfile.TemporaryDirectory().name)
17 |
18 | def save_checkpoint(self, *args, **kwargs):
19 | pass
20 |
21 |
22 | class DummyRunner:
23 | def __init__(self, task):
24 | self.task = task
25 |
26 | def test(self):
27 | self.task.trainer = DummyTrainer()
28 |
29 | optimizer = self.task.configure_optimizers()
30 |
31 | def optimizers():
32 | return optimizer
33 |
34 | self.task.optimizers = optimizers
35 |
36 | self.task.on_train_start()
37 | assert self.task.current_epoch == 0
38 | assert self.task.lr_scheduler.last_epoch == 0
39 |
40 | dummy_batch = {
41 | "img": torch.randn((2, 3, 32, 32)),
42 | "img_info": {
43 | "height": torch.randn(2),
44 | "width": torch.randn(2),
45 | "id": torch.from_numpy(np.array([0, 1])),
46 | },
47 | "gt_bboxes": [
48 | np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32),
49 | np.array(
50 | [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]], dtype=np.float32
51 | ),
52 | ],
53 | "gt_labels": [np.array([1]), np.array([1, 2])],
54 | "warp_matrix": [np.eye(3), np.eye(3)],
55 | }
56 |
57 | def func(*args, **kwargs):
58 | pass
59 |
60 | self.task.scalar_summary = func
61 | self.task.training_step(dummy_batch, 0)
62 |
63 | self.task.optimizer_step(optimizer=optimizer)
64 | self.task.training_epoch_end([])
65 | assert self.task.lr_scheduler.last_epoch == 1
66 |
67 | self.task.validation_step(dummy_batch, 0)
68 | self.task.validation_epoch_end([])
69 |
70 | self.task.test_step(dummy_batch, 0)
71 | self.task.test_epoch_end([])
72 |
73 |
74 | def test_lightning_training_task():
75 | load_config(cfg, "./config/nanodet-m.yml")
76 | task = TrainingTask(cfg)
77 | runner = DummyRunner(task)
78 | runner.test()
79 |
--------------------------------------------------------------------------------
/tests/test_trainer/test_trainer.py:
--------------------------------------------------------------------------------
1 | import tempfile
2 |
3 | import numpy as np
4 | import torch
5 | from torch.utils.data import DataLoader, Dataset
6 |
7 | from nanodet.data.collate import collate_function
8 | from nanodet.model.arch import build_model
9 | from nanodet.trainer import build_trainer
10 | from nanodet.util import Logger, cfg, load_config
11 |
12 |
13 | class DummyDataset(Dataset):
14 | def __len__(self):
15 | return 10
16 |
17 | def __getitem__(self, idx):
18 | data = {
19 | "img": torch.randn((3, 32, 32)),
20 | "img_info": {
21 | "file_name": "dummy_data.jpg",
22 | "height": 500,
23 | "width": 500,
24 | "id": 1,
25 | },
26 | "gt_bboxes": np.array(
27 | [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 3.0, 4.0]], dtype=np.float32
28 | ),
29 | "gt_labels": np.array([1, 2]),
30 | "warp_matrix": [np.eye(3), np.eye(3)],
31 | }
32 | return data
33 |
34 |
35 | class DummyEvaluator:
36 | metric_names = ["mAP"]
37 |
38 | def evaluate(self, results, save_dir, rank=-1):
39 | return {"mAP": 0.5}
40 |
41 |
42 | def test_trainer():
43 | tmp_dir = tempfile.TemporaryDirectory()
44 |
45 | load_config(cfg, "./config/nanodet-m.yml")
46 | cfg.defrost()
47 | cfg.model.arch.backbone.pretrain = False
48 | cfg.schedule.total_epochs = 4
49 | cfg.schedule.val_intervals = 1
50 | cfg.schedule.warmup.steps = 2
51 | cfg.save_dir = tmp_dir.name
52 | dummy_dataset = DummyDataset()
53 | train_loader = DataLoader(
54 | dummy_dataset,
55 | batch_size=2,
56 | num_workers=1,
57 | pin_memory=True,
58 | collate_fn=collate_function,
59 | drop_last=True,
60 | )
61 | val_loader = DataLoader(
62 | dummy_dataset,
63 | batch_size=2,
64 | num_workers=1,
65 | pin_memory=True,
66 | collate_fn=collate_function,
67 | shuffle=False,
68 | drop_last=False,
69 | )
70 |
71 | model = build_model(cfg.model)
72 | logger = Logger(-1, tmp_dir.name)
73 | device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
74 |
75 | trainer = build_trainer(rank=-1, cfg=cfg, model=model, logger=logger, device=device)
76 | trainer.run(train_loader, val_loader, DummyEvaluator())
77 |
--------------------------------------------------------------------------------
/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/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 |
--------------------------------------------------------------------------------
/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/nanodet-m.yml")
22 | logger.dump_cfg(cfg)
23 |
24 | logger.finalize(None)
25 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/tools/export_onnx.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 | import os
17 |
18 | import torch
19 |
20 | from nanodet.model.arch import build_model
21 | from nanodet.util import Logger, cfg, load_config, load_model_weight
22 |
23 |
24 | def generate_ouput_names(head_cfg):
25 | cls_names, dis_names = [], []
26 | for stride in head_cfg.strides:
27 | cls_names.append("cls_pred_stride_{}".format(stride))
28 | dis_names.append("dis_pred_stride_{}".format(stride))
29 | return cls_names + dis_names
30 |
31 |
32 | def main(config, model_path, output_path, input_shape=(320, 320)):
33 | logger = Logger(-1, config.save_dir, False)
34 | model = build_model(config.model)
35 | checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
36 | load_model_weight(model, checkpoint, logger)
37 | if config.model.arch.backbone.name == "RepVGG":
38 | deploy_config = config.model
39 | deploy_config.arch.backbone.update({"deploy": True})
40 | deploy_model = build_model(deploy_config)
41 | from nanodet.model.backbone.repvgg import repvgg_det_model_convert
42 |
43 | model = repvgg_det_model_convert(model, deploy_model)
44 | dummy_input = torch.autograd.Variable(
45 | torch.randn(1, 3, input_shape[0], input_shape[1])
46 | )
47 | output_names = None
48 | if config.model.arch.head.name == "NanoDetHead":
49 | output_names = generate_ouput_names(config.model.arch.head)
50 | torch.onnx.export(
51 | model,
52 | dummy_input,
53 | output_path,
54 | verbose=True,
55 | keep_initializers_as_inputs=True,
56 | opset_version=11,
57 | output_names=output_names,
58 | )
59 | logger.log("finished exporting onnx ")
60 |
61 |
62 | def parse_args():
63 | parser = argparse.ArgumentParser(
64 | formatter_class=argparse.ArgumentDefaultsHelpFormatter,
65 | description="Convert .pth model to onnx.",
66 | )
67 | parser.add_argument("--cfg_path", type=str, help="Path to .yml config file.")
68 | parser.add_argument(
69 | "--model_path", type=str, default=None, help="Path to .ckpt model."
70 | )
71 | parser.add_argument(
72 | "--out_path", type=str, default="nanodet.onnx", help="Onnx model output path."
73 | )
74 | parser.add_argument(
75 | "--input_shape", type=str, default=None, help="Model intput shape."
76 | )
77 | return parser.parse_args()
78 |
79 |
80 | if __name__ == "__main__":
81 | args = parse_args()
82 | cfg_path = args.cfg_path
83 | model_path = args.model_path
84 | out_path = args.out_path
85 | input_shape = args.input_shape
86 | load_config(cfg, cfg_path)
87 | if input_shape is None:
88 | input_shape = cfg.data.train.input_size
89 | else:
90 | input_shape = tuple(map(int, input_shape.split(",")))
91 | assert len(input_shape) == 2
92 | if model_path is None:
93 | model_path = os.path.join(cfg.save_dir, "model_best/model_best.ckpt")
94 | main(cfg, model_path, out_path, input_shape)
95 | print("Model saved to:", out_path)
96 |
--------------------------------------------------------------------------------
/tools/export_torchscript.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 | import os
17 |
18 | import torch
19 |
20 | from nanodet.model.arch import build_model
21 | from nanodet.util import Logger, cfg, load_config, load_model_weight
22 |
23 |
24 | def main(config, model_path: str, output_path: str, input_shape=(320, 320)):
25 | logger = Logger(local_rank=-1, save_dir=config.save_dir, use_tensorboard=False)
26 |
27 | # Create model and load weights
28 | model = build_model(config.model)
29 | checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
30 | load_model_weight(model, checkpoint, logger)
31 |
32 | # Convert backbone weights for RepVGG models
33 | if config.model.arch.backbone.name == "RepVGG":
34 | deploy_config = config.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 |
41 | # TorchScript: tracing the model with dummy inputs
42 | with torch.no_grad():
43 | dummy_input = torch.zeros(
44 | 1, 3, input_shape[0], input_shape[1]
45 | ) # Batch size = 1
46 | model.eval().cpu()
47 | model_traced = torch.jit.trace(model, example_inputs=dummy_input).eval()
48 | model_traced.save(output_path)
49 | print("Finished export to TorchScript")
50 |
51 |
52 | def parse_args():
53 | parser = argparse.ArgumentParser(
54 | formatter_class=argparse.ArgumentDefaultsHelpFormatter,
55 | description="Convert .pth model weights to TorchScript.",
56 | )
57 | parser.add_argument("--cfg_path", type=str, help="Path to .yml config file.")
58 | parser.add_argument(
59 | "--model_path", type=str, default=None, help="Path to .ckpt model."
60 | )
61 | parser.add_argument(
62 | "--out_path",
63 | type=str,
64 | default="nanodet.torchscript.pth",
65 | help="TorchScript model output path.",
66 | )
67 | parser.add_argument(
68 | "--input_shape", type=str, default=None, help="Model input shape."
69 | )
70 | return parser.parse_args()
71 |
72 |
73 | if __name__ == "__main__":
74 | args = parse_args()
75 | cfg_path = args.cfg_path
76 | model_path = args.model_path
77 | out_path = args.out_path
78 | input_shape = args.input_shape
79 | load_config(cfg, cfg_path)
80 | if input_shape is None:
81 | input_shape = cfg.data.train.input_size
82 | else:
83 | input_shape = tuple(map(int, input_shape.split(",")))
84 | assert len(input_shape) == 2
85 | if model_path is None:
86 | model_path = os.path.join(cfg.save_dir, "model_best/model_best.ckpt")
87 | main(cfg, model_path, out_path, input_shape)
88 | print("Model saved to:", out_path)
89 |
--------------------------------------------------------------------------------
/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 | from nanodet.model.arch import build_model
16 | from nanodet.util import cfg, get_model_complexity_info, load_config
17 |
18 |
19 | def main(config, input_shape=(3, 320, 320)):
20 | model = build_model(config.model)
21 | flops, params = get_model_complexity_info(model, input_shape)
22 | split_line = "=" * 30
23 | print(
24 | f"{split_line}\nInput shape: {input_shape}\n"
25 | f"Flops: {flops}\nParams: {params}\n{split_line}"
26 | )
27 |
28 |
29 | if __name__ == "__main__":
30 | cfg_path = r"config/nanodet-m.yml"
31 | load_config(cfg, cfg_path)
32 | main(config=cfg, input_shape=(3, 320, 320))
33 |
--------------------------------------------------------------------------------
/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 |
--------------------------------------------------------------------------------
/tools/test.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 | import datetime
17 | import os
18 | import warnings
19 |
20 | import pytorch_lightning as pl
21 | import torch
22 |
23 | from nanodet.data.collate import naive_collate
24 | from nanodet.data.dataset import build_dataset
25 | from nanodet.evaluator import build_evaluator
26 | from nanodet.trainer.task import TrainingTask
27 | from nanodet.util import (
28 | NanoDetLightningLogger,
29 | cfg,
30 | convert_old_model,
31 | load_config,
32 | mkdir,
33 | )
34 |
35 |
36 | def parse_args():
37 | parser = argparse.ArgumentParser()
38 | parser.add_argument(
39 | "--task", type=str, default="val", help="task to run, test or val"
40 | )
41 | parser.add_argument("--config", type=str, help="model config file(.yml) path")
42 | parser.add_argument("--model", type=str, help="ckeckpoint file(.ckpt) path")
43 | args = parser.parse_args()
44 | return args
45 |
46 |
47 | def main(args):
48 | load_config(cfg, args.config)
49 | local_rank = -1
50 | torch.backends.cudnn.enabled = True
51 | torch.backends.cudnn.benchmark = True
52 | cfg.defrost()
53 | timestr = datetime.datetime.now().__format__("%Y%m%d%H%M%S")
54 | cfg.save_dir = os.path.join(cfg.save_dir, timestr)
55 | mkdir(local_rank, cfg.save_dir)
56 | logger = NanoDetLightningLogger(cfg.save_dir)
57 |
58 | assert args.task in ["val", "test"]
59 | cfg.update({"test_mode": args.task})
60 |
61 | logger.info("Setting up data...")
62 | val_dataset = build_dataset(cfg.data.val, args.task)
63 | val_dataloader = torch.utils.data.DataLoader(
64 | val_dataset,
65 | batch_size=cfg.device.batchsize_per_gpu,
66 | shuffle=False,
67 | num_workers=cfg.device.workers_per_gpu,
68 | pin_memory=True,
69 | collate_fn=naive_collate,
70 | drop_last=False,
71 | )
72 | evaluator = build_evaluator(cfg.evaluator, val_dataset)
73 |
74 | logger.info("Creating model...")
75 | task = TrainingTask(cfg, evaluator)
76 |
77 | ckpt = torch.load(args.model)
78 | if "pytorch-lightning_version" not in ckpt:
79 | warnings.warn(
80 | "Warning! Old .pth checkpoint is deprecated. "
81 | "Convert the checkpoint with tools/convert_old_checkpoint.py "
82 | )
83 | ckpt = convert_old_model(ckpt)
84 | task.load_state_dict(ckpt["state_dict"])
85 |
86 | trainer = pl.Trainer(
87 | default_root_dir=cfg.save_dir,
88 | gpus=cfg.device.gpu_ids,
89 | accelerator="ddp",
90 | log_every_n_steps=cfg.log.interval,
91 | num_sanity_val_steps=0,
92 | logger=logger,
93 | )
94 | logger.info("Starting testing...")
95 | trainer.test(task, val_dataloader)
96 |
97 |
98 | if __name__ == "__main__":
99 | args = parse_args()
100 | main(args)
101 |
--------------------------------------------------------------------------------