├── demo_android_ncnn ├── app │ ├── .gitignore │ ├── src │ │ ├── main │ │ │ ├── res │ │ │ │ ├── values │ │ │ │ │ ├── strings.xml │ │ │ │ │ ├── colors.xml │ │ │ │ │ └── styles.xml │ │ │ │ ├── drawable-xxhdpi │ │ │ │ │ ├── cpu.png │ │ │ │ │ ├── gpu.png │ │ │ │ │ └── ncnn_icon.png │ │ │ │ ├── mipmap-hdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-mdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-xhdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-xxhdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── mipmap-xxxhdpi │ │ │ │ │ ├── ic_launcher.png │ │ │ │ │ └── ic_launcher_round.png │ │ │ │ ├── drawable │ │ │ │ │ └── cpu_gpu_bg.xml │ │ │ │ ├── mipmap-anydpi-v26 │ │ │ │ │ ├── ic_launcher.xml │ │ │ │ │ └── ic_launcher_round.xml │ │ │ │ ├── drawable-v24 │ │ │ │ │ └── ic_launcher_foreground.xml │ │ │ │ └── layout │ │ │ │ │ ├── activity_welcome.xml │ │ │ │ │ └── activity_main.xml │ │ │ ├── cpp │ │ │ │ ├── CMakeLists.txt │ │ │ │ ├── YoloV4.h │ │ │ │ ├── YoloV4.cpp │ │ │ │ ├── NanoDet.h │ │ │ │ ├── YoloV5.h │ │ │ │ └── jni_interface.cpp │ │ │ ├── java │ │ │ │ └── com │ │ │ │ │ └── rangi │ │ │ │ │ └── nanodet │ │ │ │ │ ├── YOLOv5.java │ │ │ │ │ ├── NanoDet.java │ │ │ │ │ ├── YOLOv4.java │ │ │ │ │ ├── AppCrashHandler.java │ │ │ │ │ ├── Box.java │ │ │ │ │ ├── NcnnApp.java │ │ │ │ │ └── WelcomeActivity.java │ │ │ └── AndroidManifest.xml │ │ ├── test │ │ │ └── java │ │ │ │ └── com │ │ │ │ └── rangi │ │ │ │ └── nanodet │ │ │ │ └── ExampleUnitTest.java │ │ └── androidTest │ │ │ └── java │ │ │ └── com │ │ │ └── rangi │ │ │ └── nanodet │ │ │ └── ExampleInstrumentedTest.java │ ├── proguard-rules.pro │ └── build.gradle ├── settings.gradle ├── Android_demo.jpg ├── gradle │ └── wrapper │ │ ├── gradle-wrapper.jar │ │ └── gradle-wrapper.properties ├── .gitignore ├── build.gradle ├── gradle.properties ├── README.md └── gradlew.bat ├── nanodet ├── data │ ├── transform │ │ ├── __init__.py │ │ ├── pipeline.py │ │ ├── mosaic.py │ │ └── color.py │ ├── dataset │ │ └── __init__.py │ └── collate.py ├── util │ ├── path.py │ ├── rank_filter.py │ ├── distributed_data_parallel.py │ ├── __init__.py │ ├── config.py │ ├── box_transform.py │ ├── misc.py │ ├── scatter_gather.py │ ├── check_point.py │ ├── logger.py │ └── util_mixins.py ├── model │ ├── head │ │ ├── assigner │ │ │ └── base_assigner.py │ │ └── __init__.py │ ├── module │ │ ├── scale.py │ │ ├── activation.py │ │ ├── init_weights.py │ │ ├── norm.py │ │ └── transformer.py │ ├── fpn │ │ ├── __init__.py │ │ ├── pan.py │ │ ├── fpn.py │ │ └── tan.py │ ├── arch │ │ ├── __init__.py │ │ └── one_stage_detector.py │ ├── backbone │ │ ├── __init__.py │ │ ├── custom_csp.py │ │ └── mobilenetv2.py │ └── loss │ │ └── utils.py ├── evaluator │ ├── __init__.py │ └── coco_detection.py └── trainer │ ├── __init__.py │ └── dist_trainer.py ├── docs └── imgs │ ├── Title.jpg │ ├── Model_arch.png │ └── Android_demo.jpg ├── demo_mnn ├── imgs │ ├── 000252.jpg │ └── 000258.jpg ├── results │ ├── 000252.jpg │ └── 000258.jpg ├── CMakeLists.txt ├── nanodet_mnn.hpp └── README.md ├── demo_ncnn ├── benchmark.jpg ├── CMakeLists.txt ├── nanodet.h └── README.md ├── requirements.txt ├── demo_openvino ├── CMakeLists.txt ├── nanodet_openvino.h └── README.md ├── tools ├── flops.py ├── convert_old_checkpoint.py ├── inference.py ├── show_dataset.py ├── export_torchscript.py ├── deprecated │ ├── test.py │ └── train.py ├── export_onnx.py ├── test.py └── train.py ├── demo_libtorch ├── CMakeLists.txt ├── README.md └── nanodet_libtorch.h ├── setup.py ├── .gitignore └── config ├── nanodet-m-640-mod.yml ├── nanodet_custom_xml_dataset.yml ├── nanodet-g-mod.yml ├── nanodet-m.yml ├── RepVGG └── nanodet-RepVGG-A0_416.yml ├── nanodet-m-416.yml ├── EfficientNet-Lite ├── nanodet-EfficientNet-Lite0_320.yml ├── nanodet-EfficientNet-Lite1_416.yml └── nanodet-EfficientNet-Lite2_512.yml └── Transformer └── nanodet-t.yml /demo_android_ncnn/app/.gitignore: -------------------------------------------------------------------------------- 1 | /build 2 | -------------------------------------------------------------------------------- /nanodet/data/transform/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from .pipeline import Pipeline -------------------------------------------------------------------------------- /demo_android_ncnn/settings.gradle: -------------------------------------------------------------------------------- 1 | include ':app' 2 | rootProject.name='NanoDet' 3 | -------------------------------------------------------------------------------- /docs/imgs/Title.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/docs/imgs/Title.jpg -------------------------------------------------------------------------------- /demo_mnn/imgs/000252.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_mnn/imgs/000252.jpg -------------------------------------------------------------------------------- /demo_mnn/imgs/000258.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_mnn/imgs/000258.jpg -------------------------------------------------------------------------------- /demo_ncnn/benchmark.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_ncnn/benchmark.jpg -------------------------------------------------------------------------------- /docs/imgs/Model_arch.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/docs/imgs/Model_arch.png -------------------------------------------------------------------------------- /docs/imgs/Android_demo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/docs/imgs/Android_demo.jpg -------------------------------------------------------------------------------- /demo_mnn/results/000252.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_mnn/results/000252.jpg -------------------------------------------------------------------------------- /demo_mnn/results/000258.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_mnn/results/000258.jpg -------------------------------------------------------------------------------- /demo_android_ncnn/Android_demo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/Android_demo.jpg -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/values/strings.xml: -------------------------------------------------------------------------------- 1 | 2 | NanoDet 3 | 4 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradle/wrapper/gradle-wrapper.jar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/gradle/wrapper/gradle-wrapper.jar -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-xxhdpi/cpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/cpu.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-xxhdpi/gpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/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/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/drawable-xxhdpi/ncnn_icon.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher.png -------------------------------------------------------------------------------- /nanodet/util/path.py: -------------------------------------------------------------------------------- 1 | import os 2 | from .rank_filter import rank_filter 3 | 4 | 5 | @rank_filter 6 | def mkdir(path): 7 | if not os.path.exists(path): 8 | os.makedirs(path) 9 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-hdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-mdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xhdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxhdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher_round.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/cxnaive/nanodet-TensorRT/HEAD/demo_android_ncnn/app/src/main/res/mipmap-xxxhdpi/ic_launcher_round.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | Cython 2 | termcolor 3 | numpy 4 | torch>=1.6 5 | torchvision 6 | tensorboard 7 | pycocotools 8 | matplotlib 9 | pyaml 10 | opencv-python 11 | tqdm 12 | pytorch-lightning>=1.2.5 13 | torchmetrics -------------------------------------------------------------------------------- /nanodet/util/rank_filter.py: -------------------------------------------------------------------------------- 1 | 2 | def rank_filter(func): 3 | def func_filter(local_rank=-1, *args, **kwargs): 4 | if local_rank < 1: 5 | return func(*args, **kwargs) 6 | else: 7 | pass 8 | return func_filter 9 | -------------------------------------------------------------------------------- /nanodet/model/head/assigner/base_assigner.py: -------------------------------------------------------------------------------- 1 | from abc import ABCMeta, abstractmethod 2 | 3 | 4 | class BaseAssigner(metaclass=ABCMeta): 5 | 6 | @abstractmethod 7 | def assign(self, bboxes, gt_bboxes, gt_bboxes_ignore=None, gt_labels=None): 8 | pass -------------------------------------------------------------------------------- /nanodet/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | from .coco_detection import CocoDetectionEvaluator 2 | 3 | 4 | def build_evaluator(cfg, dataset): 5 | if cfg.evaluator.name == 'CocoDetectionEvaluator': 6 | return CocoDetectionEvaluator(dataset) 7 | else: 8 | raise NotImplementedError 9 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradle/wrapper/gradle-wrapper.properties: -------------------------------------------------------------------------------- 1 | #Mon Jun 15 13:17:33 HKT 2020 2 | distributionBase=GRADLE_USER_HOME 3 | distributionPath=wrapper/dists 4 | zipStoreBase=GRADLE_USER_HOME 5 | zipStorePath=wrapper/dists 6 | distributionUrl=https\://services.gradle.org/distributions/gradle-6.1.1-all.zip 7 | -------------------------------------------------------------------------------- /demo_android_ncnn/.gitignore: -------------------------------------------------------------------------------- 1 | *.iml 2 | .gradle 3 | /local.properties 4 | /.idea 5 | /.idea/caches 6 | /.idea/libraries 7 | /.idea/modules.xml 8 | /.idea/workspace.xml 9 | /.idea/navEditor.xml 10 | /.idea/assetWizardSettings.xml 11 | .DS_Store 12 | /build 13 | /captures 14 | .externalNativeBuild 15 | .cxx 16 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/values/colors.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | #123F70 4 | #122F49 5 | #D81B60 6 | 7 | #DDDDDD 8 | 9 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable/cpu_gpu_bg.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-anydpi-v26/ic_launcher.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/mipmap-anydpi-v26/ic_launcher_round.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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/model/head/__init__.py: -------------------------------------------------------------------------------- 1 | import copy 2 | from .gfl_head import GFLHead 3 | from .nanodet_head import NanoDetHead 4 | 5 | 6 | def build_head(cfg): 7 | head_cfg = copy.deepcopy(cfg) 8 | name = head_cfg.pop('name') 9 | if name == 'GFLHead': 10 | return GFLHead(**head_cfg) 11 | elif name == 'NanoDetHead': 12 | return NanoDetHead(**head_cfg) 13 | else: 14 | raise NotImplementedError -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /nanodet/model/fpn/__init__.py: -------------------------------------------------------------------------------- 1 | import copy 2 | from .fpn import FPN 3 | from .pan import PAN 4 | from .tan import TAN 5 | 6 | 7 | def build_fpn(cfg): 8 | fpn_cfg = copy.deepcopy(cfg) 9 | name = fpn_cfg.pop('name') 10 | if name == 'FPN': 11 | return FPN(**fpn_cfg) 12 | elif name == 'PAN': 13 | return PAN(**fpn_cfg) 14 | elif name == 'TAN': 15 | return TAN(**fpn_cfg) 16 | else: 17 | raise NotImplementedError 18 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/NanoDet.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.res.AssetManager; 4 | import android.graphics.Bitmap; 5 | 6 | public class NanoDet { 7 | static { 8 | System.loadLibrary("yolov5"); 9 | } 10 | 11 | public static native void init(AssetManager manager, boolean useGPU); 12 | public static native Box[] detect(Bitmap bitmap, double threshold, double nms_threshold); 13 | } 14 | -------------------------------------------------------------------------------- /demo_mnn/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.9) 2 | project(nanodet-mnn) 3 | 4 | set(CMAKE_CXX_STANDARD 17) 5 | 6 | # find_package(OpenCV REQUIRED PATHS "/work/dependence/opencv/opencv-3.4.3/build") 7 | find_package(OpenCV REQUIRED) 8 | include_directories( 9 | mnn/include 10 | . 11 | ) 12 | 13 | link_directories(mnn/lib) 14 | 15 | add_executable(nanodet-mnn main.cpp nanodet_mnn.cpp) 16 | target_link_libraries(nanodet-mnn MNN ${OpenCV_LIBS}) 17 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/YOLOv4.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.res.AssetManager; 4 | import android.graphics.Bitmap; 5 | 6 | public class YOLOv4 { 7 | static { 8 | System.loadLibrary("yolov5"); // 存放在yolov5.so中 9 | } 10 | 11 | public static native void init(AssetManager manager, boolean useGPU); 12 | public static native Box[] detect(Bitmap bitmap, double threshold, double nms_threshold); 13 | } 14 | -------------------------------------------------------------------------------- /nanodet/data/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | import copy 2 | from .coco import CocoDataset 3 | from .xml_dataset import XMLDataset 4 | 5 | 6 | def build_dataset(cfg, mode): 7 | dataset_cfg = copy.deepcopy(cfg) 8 | name = dataset_cfg.pop('name') 9 | if name == 'coco': 10 | return CocoDataset(mode=mode, **dataset_cfg) 11 | if name == 'xml_dataset': 12 | return XMLDataset(mode=mode, **dataset_cfg) 13 | else: 14 | raise NotImplementedError('Unknown dataset type!') 15 | -------------------------------------------------------------------------------- /nanodet/util/distributed_data_parallel.py: -------------------------------------------------------------------------------- 1 | from torch.nn.parallel import DistributedDataParallel 2 | from .scatter_gather import scatter_kwargs 3 | 4 | 5 | class DDP(DistributedDataParallel): 6 | 7 | def __init__(self, batchsize, **kwargs): 8 | self.batchsize = batchsize 9 | super(DDP, self).__init__(**kwargs) 10 | 11 | def scatter(self, inputs, kwargs, device_ids): 12 | return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim, chunk_sizes=[self.batchsize]) -------------------------------------------------------------------------------- /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 | } -------------------------------------------------------------------------------- /nanodet/trainer/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .trainer import Trainer 3 | from .dist_trainer import DistTrainer 4 | 5 | 6 | def build_trainer(rank, cfg, model, logger): 7 | if len(cfg.device.gpu_ids) > 1: 8 | trainer = DistTrainer(rank, cfg, model, logger) 9 | trainer.set_device(cfg.device.batchsize_per_gpu, rank, device=torch.device('cuda')) # TODO: device 10 | else: 11 | trainer = Trainer(rank, cfg, model, logger) 12 | trainer.set_device(cfg.device.batchsize_per_gpu, cfg.device.gpu_ids, device=torch.device('cuda')) 13 | return trainer 14 | 15 | -------------------------------------------------------------------------------- /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 | ) -------------------------------------------------------------------------------- /nanodet/util/__init__.py: -------------------------------------------------------------------------------- 1 | from .rank_filter import rank_filter 2 | from .path import mkdir 3 | from .logger import Logger, MovingAverage, AverageMeter 4 | from .data_parallel import DataParallel 5 | from .distributed_data_parallel import DDP 6 | from .check_point import load_model_weight, save_model, convert_old_model 7 | from .config import cfg, load_config 8 | from .box_transform import * 9 | from .util_mixins import NiceRepr 10 | from .visualization import Visualizer, overlay_bbox_cv 11 | from .flops_counter import get_model_complexity_info 12 | from .misc import multi_apply, images_to_levels, unmap 13 | -------------------------------------------------------------------------------- /nanodet/model/module/activation.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | activations = {'ReLU': nn.ReLU, 4 | 'LeakyReLU': nn.LeakyReLU, 5 | 'ReLU6': nn.ReLU6, 6 | 'SELU': nn.SELU, 7 | 'ELU': nn.ELU, 8 | 'GELU': nn.GELU, 9 | None: nn.Identity 10 | } 11 | 12 | 13 | def act_layers(name): 14 | assert name in activations.keys() 15 | if name == 'LeakyReLU': 16 | return nn.LeakyReLU(negative_slope=0.1, inplace=True) 17 | elif name == 'GELU': 18 | return nn.GELU() 19 | else: 20 | return activations[name](inplace=True) 21 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /tools/flops.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from nanodet.model.arch import build_model 3 | from nanodet.util import cfg, load_config, get_model_complexity_info 4 | 5 | 6 | def main(config, input_shape=(3, 320, 320)): 7 | model = build_model(config.model) 8 | flops, params = get_model_complexity_info(model, input_shape) 9 | split_line = '=' * 30 10 | print(f'{split_line}\nInput shape: {input_shape}\n' 11 | f'Flops: {flops}\nParams: {params}\n{split_line}') 12 | 13 | 14 | if __name__ == '__main__': 15 | cfg_path = r"config/nanodet-m.yml" 16 | load_config(cfg, cfg_path) 17 | main(config=cfg, 18 | input_shape=(3, 320, 320) 19 | ) 20 | -------------------------------------------------------------------------------- /nanodet/data/transform/pipeline.py: -------------------------------------------------------------------------------- 1 | from .warp import warp_and_resize 2 | from .color import color_aug_and_norm 3 | import functools 4 | 5 | 6 | class Pipeline: 7 | def __init__(self, 8 | cfg, 9 | keep_ratio): 10 | self.warp = functools.partial(warp_and_resize, 11 | warp_kwargs=cfg, 12 | keep_ratio=keep_ratio) 13 | self.color = functools.partial(color_aug_and_norm, 14 | kwargs=cfg) 15 | 16 | def __call__(self, meta, dst_shape): 17 | meta = self.warp(meta=meta, dst_shape=dst_shape) 18 | meta = self.color(meta=meta) 19 | return meta 20 | -------------------------------------------------------------------------------- /demo_android_ncnn/build.gradle: -------------------------------------------------------------------------------- 1 | // Top-level build file where you can add configuration options common to all sub-projects/modules. 2 | 3 | buildscript { 4 | repositories { 5 | google() 6 | jcenter() 7 | 8 | } 9 | dependencies { 10 | classpath 'com.android.tools.build:gradle:4.0.0' 11 | 12 | // NOTE: Do not place your application dependencies here; they belong 13 | // in the individual module build.gradle files 14 | } 15 | } 16 | 17 | allprojects { 18 | repositories { 19 | google() 20 | jcenter() 21 | maven { url "https://jitpack.io" } 22 | } 23 | } 24 | 25 | task clean(type: Delete) { 26 | delete rootProject.buildDir 27 | } 28 | -------------------------------------------------------------------------------- /nanodet/model/arch/__init__.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | 3 | from .one_stage_detector import OneStageDetector 4 | 5 | 6 | def build_model(model_cfg): 7 | if model_cfg.arch.name == 'GFL': 8 | warnings.warn("Model architecture name is changed to 'OneStageDetector'. " 9 | "The name 'GFL' is deprecated, please change the model->arch->name " 10 | "in your YAML config file to OneStageDetector. ") 11 | model = OneStageDetector(model_cfg.arch.backbone, model_cfg.arch.fpn, model_cfg.arch.head) 12 | elif model_cfg.arch.name == 'OneStageDetector': 13 | model = OneStageDetector(model_cfg.arch.backbone, model_cfg.arch.fpn, model_cfg.arch.head) 14 | else: 15 | raise NotImplementedError 16 | return model 17 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/values/styles.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 10 | 16 | 17 | 18 | -------------------------------------------------------------------------------- /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_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) -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/androidTest/java/com/rangi/nanodet/ExampleInstrumentedTest.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.content.Context; 4 | 5 | import androidx.test.platform.app.InstrumentationRegistry; 6 | import androidx.test.ext.junit.runners.AndroidJUnit4; 7 | 8 | import org.junit.Test; 9 | import org.junit.runner.RunWith; 10 | 11 | import static org.junit.Assert.*; 12 | 13 | /** 14 | * Instrumented test, which will execute on an Android device. 15 | * 16 | * @see Testing documentation 17 | */ 18 | @RunWith(AndroidJUnit4.class) 19 | public class ExampleInstrumentedTest { 20 | @Test 21 | public void useAppContext() { 22 | // Context of the app under test. 23 | Context appContext = InstrumentationRegistry.getInstrumentation().getTargetContext(); 24 | 25 | assertEquals("gd.hq.yolov5", appContext.getPackageName()); 26 | } 27 | } 28 | -------------------------------------------------------------------------------- /demo_libtorch/README.md: -------------------------------------------------------------------------------- 1 | # NanoDet TorchScript / LibTorch Demo 2 | 3 | This folder provides NanoDet inference code using for LibTorch. 4 | 5 | ## Install dependencies 6 | 7 | This project needs OpenCV and CMake to work. 8 | 9 | Install CMake using a package manager of your choice. For example, the following command will install CMake on Ubuntu: 10 | 11 | ```bash 12 | sudo apt install cmake libopencv-dev 13 | ``` 14 | 15 | Also, you'll need to download LibTorch. Refer to [this page](https://pytorch.org/cppdocs/installing.html) for more info. 16 | 17 | ## Convert model 18 | 19 | Export TorchScript model using `tools/export_torchscript.py`: 20 | 21 | ```shell 22 | python ./tools/export_torchscript.py --cfg_path ${CONFIG_PATH} --model_path ${PYTORCH_MODEL_PATH} --input_shape ${MO} 23 | ``` 24 | ## Build 25 | 26 | ### Linux 27 | ```shell 28 | mkdir build 29 | cd build 30 | cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch .. 31 | make 32 | ``` -------------------------------------------------------------------------------- /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 | ) -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | from setuptools import find_packages, setup 3 | 4 | __version__ = "0.3.0" 5 | 6 | if __name__ == '__main__': 7 | setup( 8 | name='nanodet', 9 | version=__version__, 10 | description='Deep Learning Object Detection Toolbox', 11 | url='https://github.com/RangiLyu/nanodet', 12 | author='RangiLyu', 13 | author_email='lyuchqi@gmail.com', 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 | -------------------------------------------------------------------------------- /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 | }; -------------------------------------------------------------------------------- /nanodet/model/backbone/__init__.py: -------------------------------------------------------------------------------- 1 | import copy 2 | from .resnet import ResNet 3 | from .ghostnet import GhostNet 4 | from .shufflenetv2 import ShuffleNetV2 5 | from .mobilenetv2 import MobileNetV2 6 | from .efficientnet_lite import EfficientNetLite 7 | from .custom_csp import CustomCspNet 8 | from .repvgg import RepVGG 9 | 10 | 11 | def build_backbone(cfg): 12 | backbone_cfg = copy.deepcopy(cfg) 13 | name = backbone_cfg.pop('name') 14 | if name == 'ResNet': 15 | return ResNet(**backbone_cfg) 16 | elif name == 'ShuffleNetV2': 17 | return ShuffleNetV2(**backbone_cfg) 18 | elif name == 'GhostNet': 19 | return GhostNet(**backbone_cfg) 20 | elif name == 'MobileNetV2': 21 | return MobileNetV2(**backbone_cfg) 22 | elif name == 'EfficientNetLite': 23 | return EfficientNetLite(**backbone_cfg) 24 | elif name == 'CustomCspNet': 25 | return CustomCspNet(**backbone_cfg) 26 | elif name == 'RepVGG': 27 | return RepVGG(**backbone_cfg) 28 | else: 29 | raise NotImplementedError 30 | 31 | -------------------------------------------------------------------------------- /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.neck = 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 | -------------------------------------------------------------------------------- /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 | 21 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /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.0.1/nanodet_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 | ![](Android_demo.jpg) 30 | 31 | 32 | # Reference 33 | 34 | * [ncnn](https://github.com/tencent/ncnn) 35 | * [YOLOv5_NCNN](https://github.com/WZTENG/YOLOv5_NCNN) 36 | 37 | -------------------------------------------------------------------------------- /nanodet/trainer/dist_trainer.py: -------------------------------------------------------------------------------- 1 | import torch.distributed as dist 2 | from .trainer import Trainer 3 | from ..util import DDP 4 | 5 | 6 | def average_gradients(model): 7 | """ Gradient averaging. """ 8 | size = float(dist.get_world_size()) 9 | for param in model.parameters(): 10 | if param.grad is not None: 11 | dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM) 12 | param.grad.data /= size 13 | 14 | 15 | 16 | class DistTrainer(Trainer): 17 | """ 18 | Distributed trainer for multi-gpu training. (not finish yet) 19 | """ 20 | def run_step(self, model, batch, mode='train'): 21 | output, loss, loss_stats = model.module.forward_train(batch) 22 | loss = loss.mean() 23 | if mode == 'train': 24 | self.optimizer.zero_grad() 25 | loss.backward() 26 | average_gradients(model) 27 | self.optimizer.step() 28 | return output, loss, loss_stats 29 | 30 | def set_device(self, batch_per_gpu, rank, device): 31 | """ 32 | Set model device for Distributed-Data-Parallel 33 | :param batch_per_gpu: batch size of each gpu 34 | :param rank: distributed training process rank 35 | :param device: cuda 36 | """ 37 | self.rank = rank 38 | self.model = DDP(batch_per_gpu, module=self.model.cuda(), device_ids=[rank], output_device=rank) 39 | 40 | 41 | -------------------------------------------------------------------------------- /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 os 16 | import argparse 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 | parser.add_argument('--file_path', 27 | type=str, 28 | help='Path to .pth checkpoint.') 29 | parser.add_argument('--out_path', 30 | type=str, 31 | help='Path to .ckpt checkpoint.') 32 | return parser.parse_args() 33 | 34 | 35 | if __name__ == '__main__': 36 | args = parse_args() 37 | file_path = args.file_path 38 | out_path = args.out_path 39 | old_check_point = torch.load(file_path) 40 | new_check_point = convert_old_model(old_check_point) 41 | torch.save(new_check_point, out_path) 42 | print("Checkpoint saved to:", out_path) 43 | -------------------------------------------------------------------------------- /nanodet/data/transform/mosaic.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def merge_box(box1, box2): 5 | xl = max(box1[0],box2[0]) 6 | xr = min(box1[2],box2[2]) 7 | yl = max(box1[1],box2[1]) 8 | yr = min(box1[3],box2[3]) 9 | if xl >= xr or yl >= yr: 10 | return None 11 | else: 12 | area_before = (box1[2] - box1[0]) * (box1[3] - box1[1]) 13 | area_after = (xr - xl) * (yr - yl) 14 | 15 | #eliminate incomplete boxes 16 | if area_after / area_before < 0.2 or (xr - xl) / (box1[2] - box1[0]) < 0.4 or (yr - yl) / (box1[3] - box1[1]) < 0.4: 17 | return None 18 | return [xl,yl,xr,yr] 19 | 20 | 21 | def merge_bboxes(bboxes, labels, cutx, cuty, im_w, im_h): 22 | merge_bbox = [] 23 | merge_label = [] 24 | for i in range(len(bboxes)): 25 | for j in range(len(bboxes[i])): 26 | # print(labels[i][j], bboxes[i][j]) 27 | if i == 0: 28 | box_now = merge_box(bboxes[i][j], [0, 0, cutx, cuty]) 29 | if i == 1: 30 | box_now = merge_box(bboxes[i][j], [0, cuty, cutx, im_h]) 31 | if i == 2: 32 | box_now = merge_box(bboxes[i][j], [cutx, cuty, im_w, im_h]) 33 | if i == 3: 34 | box_now = merge_box(bboxes[i][j], [cutx, 0, im_w, cuty]) 35 | 36 | if box_now == None: 37 | continue 38 | 39 | merge_bbox.append(box_now) 40 | merge_label.append(labels[i][j]) 41 | # print(merge_bbox) 42 | if len(merge_bbox) == 0: 43 | return [], [] 44 | else: 45 | return merge_label, merge_bbox 46 | -------------------------------------------------------------------------------- /nanodet/model/arch/one_stage_detector.py: -------------------------------------------------------------------------------- 1 | import time 2 | import torch 3 | import torch.nn as nn 4 | from ..backbone import build_backbone 5 | from ..fpn import build_fpn 6 | from ..head import build_head 7 | 8 | 9 | class OneStageDetector(nn.Module): 10 | def __init__(self, 11 | backbone_cfg, 12 | fpn_cfg=None, 13 | head_cfg=None,): 14 | super(OneStageDetector, self).__init__() 15 | self.backbone = build_backbone(backbone_cfg) 16 | if fpn_cfg is not None: 17 | self.fpn = build_fpn(fpn_cfg) 18 | if head_cfg is not None: 19 | self.head = build_head(head_cfg) 20 | 21 | def forward(self, x): 22 | x = self.backbone(x) 23 | if hasattr(self, 'fpn'): 24 | x = self.fpn(x) 25 | if hasattr(self, 'head'): 26 | x = self.head(x) 27 | return x 28 | 29 | def inference(self, meta): 30 | with torch.no_grad(): 31 | torch.cuda.synchronize() 32 | time1 = time.time() 33 | preds = self(meta['img']) 34 | torch.cuda.synchronize() 35 | time2 = time.time() 36 | print('forward time: {:.3f}s'.format((time2 - time1)), end=' | ') 37 | results = self.head.post_process(preds, meta) 38 | torch.cuda.synchronize() 39 | print('decode time: {:.3f}s'.format((time.time() - time2)), end=' | ') 40 | return results 41 | 42 | def forward_train(self, gt_meta): 43 | preds = self(gt_meta['img']) 44 | loss, loss_states = self.head.loss(preds, gt_meta) 45 | 46 | return preds, loss, loss_states 47 | -------------------------------------------------------------------------------- /nanodet/model/module/init_weights.py: -------------------------------------------------------------------------------- 1 | """ 2 | from MMDetection 3 | """ 4 | import torch.nn as nn 5 | 6 | 7 | def kaiming_init(module, 8 | a=0, 9 | mode='fan_out', 10 | nonlinearity='relu', 11 | bias=0, 12 | distribution='normal'): 13 | assert distribution in ['uniform', 'normal'] 14 | if distribution == 'uniform': 15 | nn.init.kaiming_uniform_( 16 | module.weight, a=a, mode=mode, nonlinearity=nonlinearity) 17 | else: 18 | nn.init.kaiming_normal_( 19 | module.weight, a=a, mode=mode, nonlinearity=nonlinearity) 20 | if hasattr(module, 'bias') and module.bias is not None: 21 | nn.init.constant_(module.bias, bias) 22 | 23 | 24 | def xavier_init(module, gain=1, bias=0, distribution='normal'): 25 | assert distribution in ['uniform', 'normal'] 26 | if distribution == 'uniform': 27 | nn.init.xavier_uniform_(module.weight, gain=gain) 28 | else: 29 | nn.init.xavier_normal_(module.weight, gain=gain) 30 | if hasattr(module, 'bias') and module.bias is not None: 31 | nn.init.constant_(module.bias, bias) 32 | 33 | 34 | def normal_init(module, mean=0, std=1, bias=0): 35 | nn.init.normal_(module.weight, mean, std) 36 | if hasattr(module, 'bias') and module.bias is not None: 37 | nn.init.constant_(module.bias, bias) 38 | 39 | 40 | def constant_init(module, val, bias=0): 41 | if hasattr(module, 'weight') and module.weight is not None: 42 | nn.init.constant_(module.weight, val) 43 | if hasattr(module, 'bias') and module.bias is not None: 44 | nn.init.constant_(module.bias, bias) -------------------------------------------------------------------------------- /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_ -------------------------------------------------------------------------------- /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) -------------------------------------------------------------------------------- /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 | import torch 17 | from functools import partial 18 | 19 | 20 | def multi_apply(func, *args, **kwargs): 21 | pfunc = partial(func, **kwargs) if kwargs else func 22 | map_results = map(pfunc, *args) 23 | return tuple(map(list, zip(*map_results))) 24 | 25 | 26 | def images_to_levels(target, num_level_anchors): 27 | """Convert targets by image to targets by feature level. 28 | 29 | [target_img0, target_img1] -> [target_level0, target_level1, ...] 30 | """ 31 | target = torch.stack(target, 0) 32 | level_targets = [] 33 | start = 0 34 | for n in num_level_anchors: 35 | end = start + n 36 | level_targets.append(target[:, start:end].squeeze(0)) 37 | start = end 38 | return level_targets 39 | 40 | 41 | def unmap(data, count, inds, fill=0): 42 | """ Unmap a subset of item (data) back to the original set of items (of 43 | size count) """ 44 | if data.dim() == 1: 45 | ret = data.new_full((count, ), fill) 46 | ret[inds.type(torch.bool)] = data 47 | else: 48 | new_size = (count, ) + data.size()[1:] 49 | ret = data.new_full(new_size, fill) 50 | ret[inds.type(torch.bool), :] = data 51 | return ret -------------------------------------------------------------------------------- /nanodet/util/scatter_gather.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Variable 3 | from torch.nn.parallel._functions import Scatter 4 | 5 | 6 | def list_scatter(input, target_gpus, chunk_sizes): 7 | ret = [] 8 | for idx, size in enumerate(chunk_sizes): 9 | ret.append(input[:size]) 10 | del input[:size] 11 | return tuple(ret) 12 | 13 | def scatter(inputs, target_gpus, dim=0, chunk_sizes=None): 14 | """ 15 | Slices variables into approximately equal chunks and 16 | distributes them across given GPUs. Duplicates 17 | references to objects that are not variables. Does not 18 | support Tensors. 19 | """ 20 | def scatter_map(obj): 21 | if isinstance(obj, Variable): 22 | return Scatter.apply(target_gpus, chunk_sizes, dim, obj) 23 | assert not torch.is_tensor(obj), "Tensors not supported in scatter." 24 | if isinstance(obj, list): 25 | return list_scatter(obj, target_gpus, chunk_sizes) 26 | if isinstance(obj, tuple): 27 | return list(zip(*map(scatter_map, obj))) 28 | if isinstance(obj, dict): 29 | return list(map(type(obj), zip(*map(scatter_map, obj.items())))) 30 | return [obj for targets in target_gpus] 31 | 32 | return scatter_map(inputs) 33 | 34 | 35 | def scatter_kwargs(inputs, kwargs, target_gpus, dim=0, chunk_sizes=None): 36 | r"""Scatter with support for kwargs dictionary""" 37 | inputs = scatter(inputs, target_gpus, dim, chunk_sizes) if inputs else [] 38 | kwargs = scatter(kwargs, target_gpus, dim, chunk_sizes) if kwargs else [] 39 | if len(inputs) < len(kwargs): 40 | inputs.extend([() for _ in range(len(kwargs) - len(inputs))]) 41 | elif len(kwargs) < len(inputs): 42 | kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))]) 43 | inputs = tuple(inputs) 44 | kwargs = tuple(kwargs) 45 | return inputs, kwargs -------------------------------------------------------------------------------- /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': 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 | -------------------------------------------------------------------------------- /.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 | workspace/ 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | .vscode 107 | .idea 108 | .DS_Store 109 | 110 | # custom 111 | *.pkl 112 | *.pkl.json 113 | *.log.json 114 | work_dirs/ 115 | 116 | # Pytorch 117 | *.pth 118 | *.py~ 119 | *.sh~ 120 | -------------------------------------------------------------------------------- /demo_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/cpp/YoloV4.h: -------------------------------------------------------------------------------- 1 | #ifndef YOLOV4_H 2 | #define YOLOV4_H 3 | 4 | #include "net.h" 5 | #include "YoloV5.h" 6 | 7 | 8 | class YoloV4 { 9 | public: 10 | YoloV4(AAssetManager *mgr, const char *param, const char *bin, bool useGPU); 11 | 12 | ~YoloV4(); 13 | 14 | std::vector detect(JNIEnv *env, jobject image, float threshold, float nms_threshold); 15 | std::vector labels{"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", 16 | "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", 17 | "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 18 | "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", 19 | "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 20 | "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", 21 | "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", 22 | "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", 23 | "hair drier", "toothbrush"}; 24 | private: 25 | static std::vector 26 | decode_infer(ncnn::Mat &data, const yolocv::YoloSize &frame_size, int net_size, int num_classes, float threshold); 27 | 28 | // static void nms(std::vector& result,float nms_threshold); 29 | ncnn::Net *Net; 30 | int input_size = 640 / 2; 31 | int num_class = 80; 32 | public: 33 | static YoloV4 *detector; 34 | static bool hasGPU; 35 | }; 36 | 37 | 38 | #endif //YOLOV4_H 39 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/res/drawable-v24/ic_launcher_foreground.xml: -------------------------------------------------------------------------------- 1 | 7 | 12 | 13 | 19 | 22 | 25 | 26 | 27 | 28 | 34 | 35 | -------------------------------------------------------------------------------- /tools/inference.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import time 4 | import torch 5 | 6 | from nanodet.model.arch import build_model 7 | from nanodet.util import load_model_weight 8 | from nanodet.data.transform import Pipeline 9 | 10 | 11 | class Predictor(object): 12 | def __init__(self, cfg, model_path, logger, device='cuda:0'): 13 | self.cfg = cfg 14 | self.device = device 15 | model = build_model(cfg.model) 16 | ckpt = torch.load(model_path, map_location=lambda storage, loc: storage) 17 | load_model_weight(model, ckpt, logger) 18 | if cfg.model.arch.backbone.name == 'RepVGG': 19 | deploy_config = cfg.model 20 | deploy_config.arch.backbone.update({'deploy': True}) 21 | deploy_model = build_model(deploy_config) 22 | from nanodet.model.backbone.repvgg import repvgg_det_model_convert 23 | model = repvgg_det_model_convert(model, deploy_model) 24 | self.model = model.to(device).eval() 25 | self.pipeline = Pipeline(cfg.data.val.pipeline, cfg.data.val.keep_ratio) 26 | 27 | def inference(self, img): 28 | img_info = {} 29 | if isinstance(img, str): 30 | img_info['file_name'] = os.path.basename(img) 31 | img = cv2.imread(img) 32 | else: 33 | img_info['file_name'] = None 34 | 35 | height, width = img.shape[:2] 36 | img_info['height'] = height 37 | img_info['width'] = width 38 | meta = dict(img_info=img_info, 39 | raw_img=img, 40 | img=img) 41 | meta = self.pipeline(meta, self.cfg.data.val.input_size) 42 | meta['img'] = torch.from_numpy(meta['img'].transpose(2, 0, 1)).unsqueeze(0).to(self.device) 43 | with torch.no_grad(): 44 | results = self.model.inference(meta) 45 | return meta, results 46 | 47 | def visualize(self, dets, meta, class_names, score_thres, wait=0): 48 | time1 = time.time() 49 | self.model.head.show_result(meta['raw_img'], dets, class_names, score_thres=score_thres, show=True) 50 | print('viz time: {:.3f}s'.format(time.time()-time1)) 51 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/java/com/rangi/nanodet/NcnnApp.java: -------------------------------------------------------------------------------- 1 | package com.rangi.nanodet; 2 | 3 | import android.app.Application; 4 | import android.content.Context; 5 | import android.util.Log; 6 | 7 | import androidx.multidex.MultiDex; 8 | 9 | import com.zxy.recovery.callback.RecoveryCallback; 10 | import com.zxy.recovery.core.Recovery; 11 | 12 | 13 | public class NcnnApp extends Application { 14 | 15 | @Override 16 | public void onCreate() { 17 | super.onCreate(); 18 | 19 | //崩溃界面 20 | initRecovery(); 21 | } 22 | 23 | @Override 24 | protected void attachBaseContext(Context base) { 25 | super.attachBaseContext(base); 26 | MultiDex.install(base); 27 | } 28 | 29 | private void initRecovery() { 30 | Recovery.getInstance() 31 | .debug(BuildConfig.DEBUG) 32 | .recoverInBackground(true) 33 | .recoverStack(true) 34 | .mainPage(MainActivity.class) 35 | .recoverEnabled(true) 36 | .callback(new MyCrashCallback()) 37 | .silent(false, Recovery.SilentMode.RECOVER_ACTIVITY_STACK) 38 | // .skip(TestActivity.class) 39 | .init(this); 40 | AppCrashHandler.register(); 41 | } 42 | 43 | static final class MyCrashCallback implements RecoveryCallback { 44 | @Override 45 | public void stackTrace(String exceptionMessage) { 46 | Log.e("wzt", "exceptionMessage:" + exceptionMessage); 47 | } 48 | 49 | @Override 50 | public void cause(String cause) { 51 | Log.e("wzt", "cause:" + cause); 52 | } 53 | 54 | @Override 55 | public void exception(String exceptionType, String throwClassName, String throwMethodName, int throwLineNumber) { 56 | Log.e("wzt", "exceptionClassName:" + exceptionType); 57 | Log.e("wzt", "throwClassName:" + throwClassName); 58 | Log.e("wzt", "throwMethodName:" + throwMethodName); 59 | Log.e("wzt", "throwLineNumber:" + throwLineNumber); 60 | } 61 | 62 | @Override 63 | public void throwable(Throwable throwable) { 64 | 65 | } 66 | } 67 | 68 | } 69 | -------------------------------------------------------------------------------- /nanodet/evaluator/coco_detection.py: -------------------------------------------------------------------------------- 1 | import pycocotools.coco as coco 2 | from pycocotools.cocoeval import COCOeval 3 | import json 4 | import os 5 | import copy 6 | 7 | 8 | def xyxy2xywh(bbox): 9 | """ 10 | change bbox to coco format 11 | :param bbox: [x1, y1, x2, y2] 12 | :return: [x, y, w, h] 13 | """ 14 | return [ 15 | bbox[0], 16 | bbox[1], 17 | bbox[2] - bbox[0], 18 | bbox[3] - bbox[1], 19 | ] 20 | 21 | 22 | class CocoDetectionEvaluator: 23 | def __init__(self, dataset): 24 | assert hasattr(dataset, 'coco_api') 25 | self.coco_api = dataset.coco_api 26 | self.cat_ids = dataset.cat_ids 27 | self.metric_names = ['mAP', 'AP_50', 'AP_75', 'AP_small', 'AP_m', 'AP_l'] 28 | 29 | def results2json(self, results): 30 | """ 31 | results: {image_id: {label: [bboxes...] } } 32 | :return coco json format: {image_id: 33 | category_id: 34 | bbox: 35 | score: } 36 | """ 37 | json_results = [] 38 | for image_id, dets in results.items(): 39 | for label, bboxes in dets.items(): 40 | category_id = self.cat_ids[label] 41 | for bbox in bboxes: 42 | score = float(bbox[4]) 43 | detection = dict( 44 | image_id=int(image_id), 45 | category_id=int(category_id), 46 | bbox=xyxy2xywh(bbox), 47 | score=score) 48 | json_results.append(detection) 49 | return json_results 50 | 51 | def evaluate(self, results, save_dir, rank=-1): 52 | results_json = self.results2json(results) 53 | json_path = os.path.join(save_dir, 'results{}.json'.format(rank)) 54 | json.dump(results_json, open(json_path, 'w')) 55 | coco_dets = self.coco_api.loadRes(json_path) 56 | coco_eval = COCOeval(copy.deepcopy(self.coco_api), copy.deepcopy(coco_dets), "bbox") 57 | coco_eval.evaluate() 58 | coco_eval.accumulate() 59 | coco_eval.summarize() 60 | aps = coco_eval.stats[:6] 61 | eval_results = {} 62 | for k, v in zip(self.metric_names, aps): 63 | eval_results[k] = v 64 | return eval_results 65 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/build.gradle: -------------------------------------------------------------------------------- 1 | apply plugin: 'com.android.application' 2 | 3 | android { 4 | compileSdkVersion 29 5 | buildToolsVersion "29.0.3" 6 | defaultConfig { 7 | applicationId "com.rangi.nanodet" 8 | minSdkVersion 26 9 | targetSdkVersion 29 10 | versionCode 1 11 | versionName "1.0" 12 | testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner" 13 | externalNativeBuild { 14 | cmake { 15 | cppFlags "" 16 | arguments '-DANDROID_PLATFORM=android-24', '-DANDROID_STL=c++_static', '-DANDROID_STL=c++_shared' 17 | } 18 | } 19 | 20 | ndk { 21 | moduleName "NcnnJniLog" 22 | ldLibs "log", "z", "m" 23 | abiFilters "armeabi-v7a", "arm64-v8a" 24 | } 25 | 26 | multiDexEnabled true 27 | } 28 | buildTypes { 29 | release { 30 | minifyEnabled false 31 | proguardFiles getDefaultProguardFile('proguard-android-optimize.txt'), 'proguard-rules.pro' 32 | } 33 | } 34 | externalNativeBuild { 35 | cmake { 36 | path "src/main/cpp/CMakeLists.txt" 37 | version "3.10.2" 38 | } 39 | } 40 | sourceSets { 41 | main { 42 | jniLibs.srcDirs = ['libs'] 43 | } 44 | } 45 | 46 | repositories { 47 | flatDir { 48 | dirs 'libs' 49 | } 50 | } 51 | } 52 | 53 | dependencies { 54 | implementation fileTree(dir: 'libs', include: ['*.jar']) 55 | implementation 'androidx.appcompat:appcompat:1.1.0' 56 | implementation 'androidx.constraintlayout:constraintlayout:1.1.3' 57 | testImplementation 'junit:junit:4.12' 58 | androidTestImplementation 'androidx.test.ext:junit:1.1.1' 59 | androidTestImplementation 'androidx.test.espresso:espresso-core:3.2.0' 60 | 61 | // Use the most recent version of CameraX, currently that is alpha04 62 | def camerax_version = "1.0.0-alpha05" 63 | //noinspection GradleDependency 64 | implementation "androidx.camera:camera-core:${camerax_version}" 65 | //noinspection GradleDependency 66 | implementation "androidx.camera:camera-camera2:${camerax_version}" 67 | 68 | implementation 'com.android.support:multidex:1.0.3' 69 | // crash 70 | implementation 'com.zxy.android:recovery:1.0.0' 71 | // photoview 72 | implementation 'com.github.chrisbanes:PhotoView:2.3.0' 73 | // implementation 'com.bm.photoview:library:1.4.1' 74 | 75 | } 76 | -------------------------------------------------------------------------------- /demo_android_ncnn/gradlew.bat: -------------------------------------------------------------------------------- 1 | @if "%DEBUG%" == "" @echo off 2 | @rem ########################################################################## 3 | @rem 4 | @rem Gradle startup script for Windows 5 | @rem 6 | @rem ########################################################################## 7 | 8 | @rem Set local scope for the variables with windows NT shell 9 | if "%OS%"=="Windows_NT" setlocal 10 | 11 | set DIRNAME=%~dp0 12 | if "%DIRNAME%" == "" set DIRNAME=. 13 | set APP_BASE_NAME=%~n0 14 | set APP_HOME=%DIRNAME% 15 | 16 | @rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. 17 | set DEFAULT_JVM_OPTS= 18 | 19 | @rem Find java.exe 20 | if defined JAVA_HOME goto findJavaFromJavaHome 21 | 22 | set JAVA_EXE=java.exe 23 | %JAVA_EXE% -version >NUL 2>&1 24 | if "%ERRORLEVEL%" == "0" goto init 25 | 26 | echo. 27 | echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 28 | echo. 29 | echo Please set the JAVA_HOME variable in your environment to match the 30 | echo location of your Java installation. 31 | 32 | goto fail 33 | 34 | :findJavaFromJavaHome 35 | set JAVA_HOME=%JAVA_HOME:"=% 36 | set JAVA_EXE=%JAVA_HOME%/bin/java.exe 37 | 38 | if exist "%JAVA_EXE%" goto init 39 | 40 | echo. 41 | echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 42 | echo. 43 | echo Please set the JAVA_HOME variable in your environment to match the 44 | echo location of your Java installation. 45 | 46 | goto fail 47 | 48 | :init 49 | @rem Get command-line arguments, handling Windows variants 50 | 51 | if not "%OS%" == "Windows_NT" goto win9xME_args 52 | 53 | :win9xME_args 54 | @rem Slurp the command line arguments. 55 | set CMD_LINE_ARGS= 56 | set _SKIP=2 57 | 58 | :win9xME_args_slurp 59 | if "x%~1" == "x" goto execute 60 | 61 | set CMD_LINE_ARGS=%* 62 | 63 | :execute 64 | @rem Setup the command line 65 | 66 | set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar 67 | 68 | @rem Execute Gradle 69 | "%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" org.gradle.wrapper.GradleWrapperMain %CMD_LINE_ARGS% 70 | 71 | :end 72 | @rem End local scope for the variables with windows NT shell 73 | if "%ERRORLEVEL%"=="0" goto mainEnd 74 | 75 | :fail 76 | rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of 77 | rem the _cmd.exe /c_ return code! 78 | if not "" == "%GRADLE_EXIT_CONSOLE%" exit 1 79 | exit /b 1 80 | 81 | :mainEnd 82 | if "%OS%"=="Windows_NT" endlocal 83 | 84 | :omega 85 | -------------------------------------------------------------------------------- /nanodet/data/collate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | import re 5 | from torch._six import container_abcs, string_classes, int_classes 6 | 7 | 8 | np_str_obj_array_pattern = re.compile(r'[SaUO]') 9 | 10 | 11 | default_collate_err_msg_format = ( 12 | "default_collate: batch must contain tensors, numpy arrays, numbers, " 13 | "dicts or lists; found {}") 14 | 15 | 16 | def collate_function(batch): 17 | r"""Puts each data field into a tensor with outer dimension batch size""" 18 | 19 | elem = batch[0] 20 | elem_type = type(elem) 21 | if isinstance(elem, torch.Tensor): 22 | out = None 23 | # TODO: support pytorch < 1.3 24 | # if torch.utils.data.get_worker_info() is not None: 25 | # # If we're in a background process, concatenate directly into a 26 | # # shared memory tensor to avoid an extra copy 27 | # numel = sum([x.numel() for x in batch]) 28 | # storage = elem.storage()._new_shared(numel) 29 | # out = elem.new(storage) 30 | return torch.stack(batch, 0, out=out) 31 | elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \ 32 | and elem_type.__name__ != 'string_': 33 | elem = batch[0] 34 | if elem_type.__name__ == 'ndarray': 35 | # array of string classes and object 36 | if np_str_obj_array_pattern.search(elem.dtype.str) is not None: 37 | raise TypeError(default_collate_err_msg_format.format(elem.dtype)) 38 | 39 | # return collate_function([torch.as_tensor(b) for b in batch]) 40 | return batch 41 | elif elem.shape == (): # scalars 42 | # return torch.as_tensor(batch) 43 | return batch 44 | elif isinstance(elem, float): 45 | return torch.tensor(batch, dtype=torch.float64) 46 | elif isinstance(elem, int_classes): 47 | return torch.tensor(batch) 48 | elif isinstance(elem, string_classes): 49 | return batch 50 | elif isinstance(elem, container_abcs.Mapping): 51 | return {key: collate_function([d[key] for d in batch]) for key in elem} 52 | elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple 53 | return elem_type(*(collate_function(samples) for samples in zip(*batch))) 54 | elif isinstance(elem, container_abcs.Sequence): 55 | transposed = zip(*batch) 56 | return [collate_function(samples) for samples in transposed] 57 | 58 | raise TypeError(default_collate_err_msg_format.format(elem_type)) 59 | -------------------------------------------------------------------------------- /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 | 73 | -------------------------------------------------------------------------------- /demo_android_ncnn/app/src/main/cpp/NanoDet.h: -------------------------------------------------------------------------------- 1 | // 2 | // Create by RangiLyu 3 | // 2020 / 10 / 2 4 | // 5 | 6 | #ifndef NANODET_H 7 | #define NANODET_H 8 | 9 | #include "net.h" 10 | #include "YoloV5.h" 11 | 12 | typedef struct HeadInfo 13 | { 14 | std::string cls_layer; 15 | std::string dis_layer; 16 | int stride; 17 | } HeadInfo; 18 | 19 | 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_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 | -------------------------------------------------------------------------------- /nanodet/util/check_point.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import pytorch_lightning as pl 3 | from collections import OrderedDict 4 | from .rank_filter import rank_filter 5 | 6 | 7 | def load_model_weight(model, checkpoint, logger): 8 | state_dict = checkpoint['state_dict'] 9 | # strip prefix of state_dict 10 | if list(state_dict.keys())[0].startswith('module.'): 11 | state_dict = {k[7:]: v for k, v in checkpoint['state_dict'].items()} 12 | if list(state_dict.keys())[0].startswith('model.'): 13 | state_dict = {k[6:]: v for k, v in checkpoint['state_dict'].items()} 14 | 15 | model_state_dict = model.module.state_dict() if hasattr(model, 'module') else model.state_dict() 16 | 17 | # check loaded parameters and created model parameters 18 | for k in state_dict: 19 | if k in model_state_dict: 20 | if state_dict[k].shape != model_state_dict[k].shape: 21 | logger.log('Skip loading parameter {}, required shape{}, loaded shape{}.'.format( 22 | k, model_state_dict[k].shape, state_dict[k].shape)) 23 | state_dict[k] = model_state_dict[k] 24 | else: 25 | logger.log('Drop parameter {}.'.format(k)) 26 | for k in model_state_dict: 27 | if not (k in state_dict): 28 | logger.log('No param {}.'.format(k)) 29 | state_dict[k] = model_state_dict[k] 30 | model.load_state_dict(state_dict, strict=False) 31 | 32 | 33 | @rank_filter 34 | def save_model(model, path, epoch, iter, optimizer=None): 35 | model_state_dict = model.module.state_dict() if hasattr(model, 'module') else model.state_dict() 36 | data = {'epoch': epoch, 37 | 'state_dict': model_state_dict, 38 | 'iter': iter} 39 | if optimizer is not None: 40 | data['optimizer'] = optimizer.state_dict() 41 | 42 | torch.save(data, path) 43 | 44 | 45 | def convert_old_model(old_model_dict): 46 | if 'pytorch-lightning_version' in old_model_dict: 47 | raise ValueError('This model is not old format. No need to convert!') 48 | version = pl.__version__ 49 | epoch = old_model_dict['epoch'] 50 | global_step = old_model_dict['iter'] 51 | state_dict = old_model_dict['state_dict'] 52 | new_state_dict = OrderedDict() 53 | for name, value in state_dict.items(): 54 | new_state_dict['model.' + name] = value 55 | 56 | new_checkpoint = {'epoch': epoch, 57 | 'global_step': global_step, 58 | 'pytorch-lightning_version': version, 59 | 'state_dict': new_state_dict, 60 | 'lr_schedulers': []} 61 | 62 | if 'optimizer' in old_model_dict: 63 | optimizer_states = [old_model_dict['optimizer']] 64 | new_checkpoint['optimizer_states'] = optimizer_states 65 | 66 | return new_checkpoint 67 | -------------------------------------------------------------------------------- /tools/show_dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import argparse 4 | import numpy as np 5 | import cv2 6 | import warnings 7 | import pytorch_lightning as pl 8 | from pytorch_lightning.callbacks import ProgressBar 9 | 10 | from nanodet.util import mkdir, Logger, cfg, load_config, convert_old_model 11 | from nanodet.data.collate import collate_function 12 | from nanodet.data.dataset import build_dataset 13 | from nanodet.trainer.task import TrainingTask 14 | from nanodet.evaluator import build_evaluator 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('config', help='train config file path') 19 | parser.add_argument('--local_rank', default=-1, type=int, 20 | help='node rank for distributed training') 21 | parser.add_argument('--seed', type=int, default=None, 22 | help='random seed') 23 | args = parser.parse_args() 24 | return args 25 | 26 | 27 | def main(args): 28 | load_config(cfg, args.config) 29 | if cfg.model.arch.head.num_classes != len(cfg.class_names): 30 | raise ValueError('cfg.model.arch.head.num_classes must equal len(cfg.class_names),but got {} and {}'.format(cfg.model.arch.head.num_classes,len(cfg.class_names))) 31 | local_rank = int(args.local_rank) 32 | torch.backends.cudnn.enabled = True 33 | torch.backends.cudnn.benchmark = True 34 | mkdir(local_rank, cfg.save_dir) 35 | logger = Logger(local_rank, cfg.save_dir) 36 | 37 | if args.seed is not None: 38 | logger.log('Set random seed to {}'.format(args.seed)) 39 | pl.seed_everything(args.seed) 40 | 41 | logger.log('Setting up data...') 42 | # print(cfg.data.train) 43 | train_dataset = build_dataset(cfg.data.train, 'train') 44 | val_dataset = build_dataset(cfg.data.val, 'test') 45 | 46 | evaluator = build_evaluator(cfg, val_dataset) 47 | 48 | show_cnt = 0 49 | for meta in train_dataset: 50 | img = meta['img'].numpy().transpose(1,2,0).copy() 51 | #cv2.imshow('gray',np.ones(img.shape,dtype=np.float32) * 0.5) 52 | #print(meta['gt_bboxes']) 53 | #print(img.shape) 54 | for idx in range(len(meta['gt_bboxes'])): 55 | rect = meta['gt_bboxes'][idx] 56 | #print(rect[0]) 57 | label = meta['gt_labels'][idx] 58 | #print(label) 59 | cv2.rectangle(img, (int(rect[0]),int(rect[1])),(int(rect[2]),int(rect[3])),(0, 255, 0), 1) 60 | cv2.putText(img,str(cfg.class_names[label]),(int(rect[0]),int(rect[1] + 1.5)),cv2.FONT_HERSHEY_SIMPLEX,0.7,(0,255,0),1) 61 | cv2.imshow('dataset_img',img) 62 | cv2.waitKey(0) 63 | 64 | show_cnt += 1 65 | if show_cnt > 50: 66 | break 67 | 68 | 69 | if __name__ == '__main__': 70 | args = parse_args() 71 | main(args) -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /config/nanodet-m-640-mod.yml: -------------------------------------------------------------------------------- 1 | #nanodet-m-416 2 | # COCO mAP(0.5:0.95) = 0.235 3 | # AP_50 = 0.384 4 | # AP_75 = 0.242 5 | # AP_small = 0.082 6 | # AP_m = 0.240 7 | # AP_l = 0.375 8 | save_dir: workspace/rmcv_m 9 | model: 10 | arch: 11 | name: OneStageDetector 12 | backbone: 13 | name: ShuffleNetV2 14 | model_size: 1.0x 15 | out_stages: [2,3,4] 16 | activation: LeakyReLU 17 | fpn: 18 | name: PAN 19 | in_channels: [116, 232, 464] 20 | out_channels: 96 21 | start_level: 0 22 | num_outs: 3 23 | head: 24 | name: NanoDetHead 25 | num_classes: 18 26 | input_channel: 96 27 | feat_channels: 96 28 | stacked_convs: 2 29 | share_cls_reg: True 30 | octave_base_scale: 5 31 | scales_per_octave: 1 32 | strides: [8, 16, 32] 33 | reg_max: 7 34 | norm_cfg: 35 | type: BN 36 | loss: 37 | loss_qfl: 38 | name: QualityFocalLoss 39 | use_sigmoid: True 40 | beta: 2.0 41 | loss_weight: 1.0 42 | loss_dfl: 43 | name: DistributionFocalLoss 44 | loss_weight: 0.25 45 | loss_bbox: 46 | name: GIoULoss 47 | loss_weight: 2.0 48 | data: 49 | train: 50 | name: coco 51 | img_path: ../rmcv_yolo/images 52 | ann_path: ../rmcv_yolo/annotations/train.json 53 | input_size: [640,640] #[w,h] 54 | keep_ratio: True 55 | mixup: True 56 | load_mosaic: False 57 | pipeline: 58 | perspective: 0.0002 59 | scale: [0.5, 1.4] 60 | stretch: [[1, 1], [1, 1]] 61 | rotation: 5 62 | shear: 3 63 | translate: 0.1 64 | flip: 0.5 65 | hsv_h: 0.0138 66 | hsv_s: 0.464 67 | hsv_v: 0.464 68 | val: 69 | name: coco 70 | img_path: ../rmcv_yolo/images 71 | ann_path: ../rmcv_yolo/annotations/val.json 72 | input_size: [640,640] #[w,h] 73 | keep_ratio: True 74 | pipeline: 75 | normalize: [[103.53, 116.28, 123.675], [57.375, 57.12, 58.395]] 76 | device: 77 | gpu_ids: [0] 78 | workers_per_gpu: 8 79 | batchsize_per_gpu: 32 80 | schedule: 81 | resume: 82 | # load_model: nanodet_m.ckpt 83 | optimizer: 84 | name: SGD 85 | lr: 0.01 86 | momentum: 0.9 87 | weight_decay: 0.0001 88 | warmup: 89 | name: linear 90 | steps: 300 91 | ratio: 0.01 92 | total_epochs: 100 93 | lr_schedule: 94 | name: MultiStepLR 95 | milestones: [30,60,90] 96 | gamma: 0.1 97 | val_intervals: 2 98 | evaluator: 99 | name: CocoDetectionEvaluator 100 | save_key: mAP 101 | 102 | log: 103 | interval: 10 104 | 105 | class_names: ['R1', 'B1', 'R2', 'B2', 'R3', 'B3', 'R4', 'B4','R5', 'B5', 'R7', 'B7', 'R10', 'B10', 'R11', 'B11', 'RE', 'BE'] 106 | -------------------------------------------------------------------------------- /nanodet/util/logger.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | import torch 4 | import numpy as np 5 | from termcolor import colored 6 | from .rank_filter import rank_filter 7 | from .path import mkdir 8 | 9 | 10 | class Logger: 11 | def __init__(self, local_rank, save_dir='./', use_tensorboard=True): 12 | mkdir(local_rank, save_dir) 13 | self.rank = local_rank 14 | fmt = colored('[%(name)s]', 'magenta', attrs=['bold']) + colored('[%(asctime)s]', 'blue') + \ 15 | colored('%(levelname)s:', 'green') + colored('%(message)s', 'white') 16 | logging.basicConfig(level=logging.INFO, 17 | filename=os.path.join(save_dir, 'logs.txt'), 18 | filemode='w') 19 | self.log_dir = os.path.join(save_dir, 'logs') 20 | console = logging.StreamHandler() 21 | console.setLevel(logging.INFO) 22 | formatter = logging.Formatter(fmt, datefmt="%m-%d %H:%M:%S") 23 | console.setFormatter(formatter) 24 | logging.getLogger().addHandler(console) 25 | if use_tensorboard: 26 | try: 27 | from torch.utils.tensorboard import SummaryWriter 28 | except ImportError: 29 | raise ImportError( 30 | 'Please run "pip install future tensorboard" to install ' 31 | 'the dependencies to use torch.utils.tensorboard ' 32 | '(applicable to PyTorch 1.1 or higher)') 33 | if self.rank < 1: 34 | logging.info('Using Tensorboard, logs will be saved in {}'.format(self.log_dir)) 35 | self.writer = SummaryWriter(log_dir=self.log_dir) 36 | 37 | def log(self, string): 38 | if self.rank < 1: 39 | logging.info(string) 40 | 41 | def scalar_summary(self, tag, phase, value, step): 42 | if self.rank < 1: 43 | self.writer.add_scalars(tag, {phase: value}, step) 44 | 45 | 46 | class MovingAverage(object): 47 | def __init__(self, val, window_size=50): 48 | self.window_size = window_size 49 | self.reset() 50 | self.push(val) 51 | 52 | def reset(self): 53 | self.queue = [] 54 | 55 | def push(self, val): 56 | self.queue.append(val) 57 | if len(self.queue) > self.window_size: 58 | self.queue.pop(0) 59 | 60 | def avg(self): 61 | return np.mean(self.queue) 62 | 63 | 64 | class AverageMeter(object): 65 | """Computes and stores the average and current value""" 66 | 67 | def __init__(self, val): 68 | self.reset() 69 | self.update(val) 70 | 71 | def reset(self): 72 | self.val = 0 73 | self.avg = 0 74 | self.sum = 0 75 | self.count = 0 76 | 77 | def update(self, val, n=1): 78 | self.val = val 79 | self.sum += val * n 80 | self.count += n 81 | if self.count > 0: 82 | self.avg = self.sum / self.count 83 | -------------------------------------------------------------------------------- /nanodet/data/transform/color.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import random 4 | 5 | 6 | # def random_brightness(img, delta): 7 | # img += random.uniform(-delta, delta) 8 | # return img 9 | 10 | #Gamma trans 11 | def gamma_trans(img, gamma): # gamma函数处理 12 | gamma_table = [np.power(x / 255.0, gamma) * 255.0 for x in range(256)] # 建立映射表 13 | gamma_table = np.round(np.array(gamma_table)).astype(np.uint8) # 颜色值为整数 14 | return cv2.LUT(img, gamma_table) # 图片颜色查表。另外可以根据光强（颜色）均匀化原则设计自适应算法。 15 | 16 | def random_gamma(img, gamma_low, gamme_up): 17 | gamma = random.uniform(gamma_low, gamme_up) 18 | return gamma_trans(img,gamma) 19 | 20 | 21 | def augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5): 22 | r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1 # random gains 23 | hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV)) 24 | dtype = img.dtype # uint8 25 | 26 | x = np.arange(0, 256, dtype=np.int16) 27 | lut_hue = ((x * r[0]) % 180).astype(dtype) 28 | lut_sat = np.clip(x * r[1], 0, 255).astype(dtype) 29 | lut_val = np.clip(x * r[2], 0, 255).astype(dtype) 30 | 31 | img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype) 32 | cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) # no return needed 33 | 34 | def normalize(meta, mean, std): 35 | img = meta['img'].astype(np.float32) 36 | mean = np.array(mean, dtype=np.float64).reshape(1, -1) 37 | stdinv = 1 / np.array(std, dtype=np.float64).reshape(1, -1) 38 | cv2.subtract(img, mean, img) 39 | cv2.multiply(img, stdinv, img) 40 | meta['img'] = img 41 | return meta 42 | 43 | 44 | def _normalize(img, mean, std): 45 | img = img.astype(np.float32) 46 | mean = np.array(mean, dtype=np.float32).reshape(1, 1, 3) 47 | std = np.array(std, dtype=np.float32).reshape(1, 1, 3) 48 | img = (img - mean) / std 49 | return img 50 | 51 | 52 | def color_aug_and_norm(meta, kwargs): 53 | img = meta['img'] #.astype(np.float32) / 255 54 | if 'gamma' in kwargs and random.randint(0, 1): 55 | img = random_gamma(img, *kwargs['gamma']) 56 | #print(kwargs) 57 | if 'hsv_h' in kwargs and 'hsv_s' in kwargs and 'hsv_v' in kwargs: 58 | hgain=kwargs['hsv_h'] 59 | sgain=kwargs['hsv_s'] 60 | vgain=kwargs['hsv_v'] 61 | augment_hsv(img,hgain,sgain,vgain) 62 | # if 'brightness' in kwargs and random.randint(0, 1): 63 | # hsv_img = random_brightness(hsv_img, kwargs['brightness']) 64 | 65 | # if 'saturation' in kwargs and random.randint(0, 1): 66 | # hsv_img = random_saturation(hsv_img, *kwargs['saturation']) 67 | 68 | # if 'contrast' in kwargs and random.randint(0, 1): 69 | # img = random_contrast(img, *kwargs['contrast']) 70 | # cv2.imshow('trans', img) 71 | # cv2.waitKey(0) 72 | # img = _normalize(img, *kwargs['normalize']) 73 | meta['img'] = img.astype(np.float32) / 255.0 74 | return meta 75 | 76 | 77 | -------------------------------------------------------------------------------- /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.0.1/nanodet_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 | ![bench_mark](benchmark.jpg) 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 | -------------------------------------------------------------------------------- /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: xml_dataset 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: xml_dataset 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 | -------------------------------------------------------------------------------- /tools/export_torchscript.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import torch 4 | from nanodet.model.arch import build_model 5 | from nanodet.util import Logger, cfg, load_config, load_model_weight 6 | 7 | 8 | def main(config, model_path: str, output_path: str, input_shape=(320, 320)): 9 | logger = Logger(local_rank=-1, save_dir=config.save_dir, use_tensorboard=False) 10 | 11 | # Create model and load weights 12 | model = build_model(config.model) 13 | checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) 14 | load_model_weight(model, checkpoint, logger) 15 | 16 | # Convert backbone weights for RepVGG models 17 | if config.model.arch.backbone.name == 'RepVGG': 18 | deploy_config = config.model 19 | deploy_config.arch.backbone.update({'deploy': True}) 20 | deploy_model = build_model(deploy_config) 21 | from nanodet.model.backbone.repvgg import repvgg_det_model_convert 22 | model = repvgg_det_model_convert(model, deploy_model) 23 | 24 | # TorchScript: tracing the model with dummy inputs 25 | with torch.no_grad(): 26 | dummy_input = torch.zeros(1, 3, input_shape[0], input_shape[1]) # Batch size = 1 27 | model.eval().cpu() 28 | model_traced = torch.jit.trace(model, example_inputs=dummy_input).eval() 29 | model_traced.save(output_path) 30 | print('Finished export to TorchScript') 31 | 32 | 33 | def parse_args(): 34 | parser = argparse.ArgumentParser( 35 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, 36 | description='Convert .pth model weights to TorchScript.') 37 | parser.add_argument('--cfg_path', 38 | type=str, 39 | help='Path to .yml config file.') 40 | parser.add_argument('--model_path', 41 | type=str, 42 | default=None, 43 | help='Path to .ckpt model.') 44 | parser.add_argument('--out_path', 45 | type=str, 46 | default='nanodet.torchscript.pth', 47 | help='TorchScript model output path.') 48 | parser.add_argument('--input_shape', 49 | type=str, 50 | default=None, 51 | help='Model input shape.') 52 | return parser.parse_args() 53 | 54 | 55 | if __name__ == '__main__': 56 | args = parse_args() 57 | cfg_path = args.cfg_path 58 | model_path = args.model_path 59 | out_path = args.out_path 60 | input_shape = args.input_shape 61 | load_config(cfg, cfg_path) 62 | if input_shape is None: 63 | input_shape = cfg.data.train.input_size 64 | else: 65 | input_shape = tuple(map(int, input_shape.split(','))) 66 | assert len(input_shape) == 2 67 | if model_path is None: 68 | model_path = os.path.join(cfg.save_dir, "model_best/model_best.ckpt") 69 | main(cfg, model_path, out_path, input_shape) 70 | print("Model saved to:", out_path) 71 | -------------------------------------------------------------------------------- /tools/deprecated/test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import json 4 | import datetime 5 | import argparse 6 | import warnings 7 | 8 | from nanodet.util import mkdir, Logger, cfg, load_config 9 | from nanodet.trainer import build_trainer 10 | from nanodet.data.collate import collate_function 11 | from nanodet.data.dataset import build_dataset 12 | from nanodet.model.arch import build_model 13 | from nanodet.evaluator import build_evaluator 14 | 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser() 18 | parser.add_argument('--task', type=str, default='val', help='task to run, test or val') 19 | parser.add_argument('--config', type=str, help='model config file(.yml) path') 20 | parser.add_argument('--model', type=str, help='model weight file(.pth) path') 21 | parser.add_argument('--save_result', action='store_true', default=True, help='save val results to txt') 22 | args = parser.parse_args() 23 | return args 24 | 25 | 26 | def main(args): 27 | warnings.warn('Warning! Old testing code is deprecated and will be deleted ' 28 | 'in next version. Please use tools/test.py') 29 | load_config(cfg, args.config) 30 | local_rank = -1 31 | torch.backends.cudnn.enabled = True 32 | torch.backends.cudnn.benchmark = True 33 | cfg.defrost() 34 | timestr = datetime.datetime.now().__format__('%Y%m%d%H%M%S') 35 | cfg.save_dir = os.path.join(cfg.save_dir, timestr) 36 | cfg.freeze() 37 | mkdir(local_rank, cfg.save_dir) 38 | logger = Logger(local_rank, cfg.save_dir) 39 | 40 | logger.log('Creating model...') 41 | model = build_model(cfg.model) 42 | 43 | logger.log('Setting up data...') 44 | val_dataset = build_dataset(cfg.data.val, args.task) 45 | val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1, 46 | pin_memory=True, collate_fn=collate_function, drop_last=True) 47 | trainer = build_trainer(local_rank, cfg, model, logger) 48 | cfg.schedule.update({'load_model': args.model}) 49 | trainer.load_model(cfg) 50 | evaluator = build_evaluator(cfg, val_dataset) 51 | logger.log('Starting testing...') 52 | with torch.no_grad(): 53 | results, val_loss_dict = trainer.run_epoch(0, val_dataloader, mode=args.task) 54 | if args.task == 'test': 55 | res_json = evaluator.results2json(results) 56 | json_path = os.path.join(cfg.save_dir, 'results{}.json'.format(timestr)) 57 | json.dump(res_json, open(json_path, 'w')) 58 | elif args.task == 'val': 59 | eval_results = evaluator.evaluate(results, cfg.save_dir, rank=local_rank) 60 | if args.save_result: 61 | txt_path = os.path.join(cfg.save_dir, "eval_results{}.txt".format(timestr)) 62 | with open(txt_path, "a") as f: 63 | for k, v in eval_results.items(): 64 | f.write("{}: {}\n".format(k, v)) 65 | 66 | 67 | if __name__ == '__main__': 68 | args = parse_args() 69 | main(args) 70 | -------------------------------------------------------------------------------- /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 | -------------------------------------------------------------------------------- /tools/export_onnx.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import torch 4 | from nanodet.model.arch import build_model 5 | from nanodet.util import Logger, cfg, load_config, load_model_weight 6 | 7 | 8 | def generate_ouput_names(head_cfg): 9 | cls_names, dis_names = [], [] 10 | for stride in head_cfg.strides: 11 | cls_names.append('cls_pred_stride_{}'.format(stride)) 12 | dis_names.append('dis_pred_stride_{}'.format(stride)) 13 | return cls_names + dis_names 14 | 15 | 16 | def main(config, model_path, output_path, input_shape=(320, 320)): 17 | logger = Logger(-1, config.save_dir, False) 18 | model = build_model(config.model) 19 | checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) 20 | load_model_weight(model, checkpoint, logger) 21 | if config.model.arch.backbone.name == 'RepVGG': 22 | deploy_config = config.model 23 | deploy_config.arch.backbone.update({'deploy': True}) 24 | deploy_model = build_model(deploy_config) 25 | from nanodet.model.backbone.repvgg import repvgg_det_model_convert 26 | model = repvgg_det_model_convert(model, deploy_model) 27 | dummy_input = torch.autograd.Variable(torch.randn(1, 3, input_shape[0], input_shape[1])) 28 | output_names = None 29 | # if config.model.arch.head.name == 'NanoDetHead': 30 | # output_names = generate_ouput_names(config.model.arch.head) 31 | torch.onnx.export(model, 32 | dummy_input, 33 | output_path, 34 | verbose=True, 35 | keep_initializers_as_inputs=True, 36 | opset_version=11, 37 | output_names=output_names) 38 | logger.log('finished exporting onnx ') 39 | 40 | 41 | def parse_args(): 42 | parser = argparse.ArgumentParser( 43 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, 44 | description='Convert .pth model to onnx.') 45 | parser.add_argument('--cfg_path', 46 | type=str, 47 | help='Path to .yml config file.') 48 | parser.add_argument('--model_path', 49 | type=str, 50 | default=None, 51 | help='Path to .ckpt model.') 52 | parser.add_argument('--out_path', 53 | type=str, 54 | default='nanodet.onnx', 55 | help='Onnx model output path.') 56 | parser.add_argument('--input_shape', 57 | type=str, 58 | default=None, 59 | help='Model intput shape.') 60 | return parser.parse_args() 61 | 62 | 63 | if __name__ == '__main__': 64 | args = parse_args() 65 | cfg_path = args.cfg_path 66 | model_path = args.model_path 67 | out_path = args.out_path 68 | input_shape = args.input_shape 69 | load_config(cfg, cfg_path) 70 | if input_shape is None: 71 | input_shape = cfg.data.train.input_size 72 | else: 73 | input_shape = tuple(map(int, input_shape.split(','))) 74 | assert len(input_shape) == 2 75 | if model_path is None: 76 | model_path = os.path.join(cfg.save_dir, "model_best/model_best.ckpt") 77 | main(cfg, model_path, out_path, input_shape) 78 | print("Model saved to:", out_path) 79 | -------------------------------------------------------------------------------- /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 os 16 | import torch 17 | import json 18 | import datetime 19 | import argparse 20 | import warnings 21 | import pytorch_lightning as pl 22 | 23 | from nanodet.util import mkdir, Logger, cfg, load_config, convert_old_model 24 | from nanodet.data.collate import collate_function 25 | from nanodet.data.dataset import build_dataset 26 | from nanodet.trainer.task import TrainingTask 27 | from nanodet.evaluator import build_evaluator 28 | 29 | 30 | def parse_args(): 31 | parser = argparse.ArgumentParser() 32 | parser.add_argument('--task', type=str, default='val', help='task to run, test or val') 33 | parser.add_argument('--config', type=str, help='model config file(.yml) path') 34 | parser.add_argument('--model', type=str, help='ckeckpoint file(.ckpt) path') 35 | args = parser.parse_args() 36 | return args 37 | 38 | 39 | def main(args): 40 | load_config(cfg, args.config) 41 | local_rank = -1 42 | torch.backends.cudnn.enabled = True 43 | torch.backends.cudnn.benchmark = True 44 | cfg.defrost() 45 | timestr = datetime.datetime.now().__format__('%Y%m%d%H%M%S') 46 | cfg.save_dir = os.path.join(cfg.save_dir, timestr) 47 | mkdir(local_rank, cfg.save_dir) 48 | logger = Logger(local_rank, cfg.save_dir) 49 | 50 | assert args.task in ['val', 'test'] 51 | cfg.update({'test_mode': args.task}) 52 | 53 | logger.log('Setting up data...') 54 | val_dataset = build_dataset(cfg.data.val, args.task) 55 | val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, 56 | num_workers=cfg.device.workers_per_gpu, 57 | pin_memory=True, collate_fn=collate_function, drop_last=True) 58 | evaluator = build_evaluator(cfg, val_dataset) 59 | 60 | logger.log('Creating model...') 61 | task = TrainingTask(cfg, evaluator) 62 | 63 | ckpt = torch.load(args.model) 64 | if 'pytorch-lightning_version' not in ckpt: 65 | warnings.warn('Warning! Old .pth checkpoint is deprecated. ' 66 | 'Convert the checkpoint with tools/convert_old_checkpoint.py ') 67 | ckpt = convert_old_model(ckpt) 68 | task.load_state_dict(ckpt['state_dict']) 69 | 70 | trainer = pl.Trainer(default_root_dir=cfg.save_dir, 71 | gpus=cfg.device.gpu_ids, 72 | accelerator='ddp', 73 | log_every_n_steps=cfg.log.interval, 74 | num_sanity_val_steps=0, 75 | ) 76 | logger.log('Starting testing...') 77 | trainer.test(task, val_dataloader) 78 | 79 | 80 | if __name__ == '__main__': 81 | args = parse_args() 82 | main(args) 83 | -------------------------------------------------------------------------------- /nanodet/model/fpn/pan.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch.nn.functional as F 3 | from ..module.conv import ConvModule 4 | from .fpn import FPN 5 | 6 | 7 | class PAN(FPN): 8 | """Path Aggregation Network for Instance Segmentation. 9 | 10 | This is an implementation of the `PAN in Path Aggregation Network 11 | `_. 12 | 13 | Args: 14 | in_channels (List[int]): Number of input channels per scale. 15 | out_channels (int): Number of output channels (used at each scale) 16 | num_outs (int): Number of output scales. 17 | start_level (int): Index of the start input backbone level used to 18 | build the feature pyramid. Default: 0. 19 | end_level (int): Index of the end input backbone level (exclusive) to 20 | build the feature pyramid. Default: -1, which means the last level. 21 | add_extra_convs (bool): Whether to add conv layers on top of the 22 | original feature maps. Default: False. 23 | extra_convs_on_inputs (bool): Whether to apply extra conv on 24 | the original feature from the backbone. Default: False. 25 | relu_before_extra_convs (bool): Whether to apply relu before the extra 26 | conv. Default: False. 27 | no_norm_on_lateral (bool): Whether to apply norm on lateral. 28 | Default: False. 29 | conv_cfg (dict): Config dict for convolution layer. Default: None. 30 | norm_cfg (dict): Config dict for normalization layer. Default: None. 31 | act_cfg (str): Config dict for activation layer in ConvModule. 32 | Default: None. 33 | """ 34 | 35 | def __init__(self, 36 | in_channels, 37 | out_channels, 38 | num_outs, 39 | start_level=0, 40 | end_level=-1, 41 | conv_cfg=None, 42 | norm_cfg=None, 43 | activation=None): 44 | super(PAN, 45 | self).__init__(in_channels, out_channels, num_outs, start_level, 46 | end_level, conv_cfg, norm_cfg, activation) 47 | self.init_weights() 48 | 49 | def forward(self, inputs): 50 | """Forward function.""" 51 | assert len(inputs) == len(self.in_channels) 52 | 53 | # build laterals 54 | laterals = [ 55 | lateral_conv(inputs[i + self.start_level]) 56 | for i, lateral_conv in enumerate(self.lateral_convs) 57 | ] 58 | 59 | # build top-down path 60 | used_backbone_levels = len(laterals) 61 | for i in range(used_backbone_levels - 1, 0, -1): 62 | laterals[i - 1] += F.interpolate( 63 | laterals[i], scale_factor=2, mode='bilinear') 64 | 65 | # build outputs 66 | # part 1: from original levels 67 | inter_outs = [ 68 | laterals[i] for i in range(used_backbone_levels) 69 | ] 70 | 71 | # part 2: add bottom-up path 72 | for i in range(0, used_backbone_levels - 1): 73 | inter_outs[i + 1] += F.interpolate(inter_outs[i], scale_factor=0.5, mode='bilinear') 74 | 75 | outs = [] 76 | outs.append(inter_outs[0]) 77 | outs.extend([ 78 | inter_outs[i] for i in range(1, used_backbone_levels) 79 | ]) 80 | return tuple(outs) 81 | -------------------------------------------------------------------------------- /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 |