├── Dataset ├── REAME.md └── download_dataset.py ├── Deepstream-app ├── config_infer_primary_yoloV5.txt ├── deepstream_app_config (CSI).txt ├── deepstream_app_config.txt ├── deepstream_app_config_USB.txt ├── labels.txt └── nvdsinfer_custom_impl_Yolo │ ├── Makefile │ ├── calibrator.cpp │ ├── calibrator.h │ ├── layers │ ├── activation_layer.cpp │ ├── activation_layer.h │ ├── activation_layer.o │ ├── channels_layer.cpp │ ├── channels_layer.h │ ├── channels_layer.o │ ├── convolutional_layer.cpp │ ├── convolutional_layer.h │ ├── convolutional_layer.o │ ├── dropout_layer.cpp │ ├── dropout_layer.h │ ├── dropout_layer.o │ ├── implicit_layer.cpp │ ├── implicit_layer.h │ ├── implicit_layer.o │ ├── maxpool_layer.cpp │ ├── maxpool_layer.h │ ├── maxpool_layer.o │ ├── route_layer.cpp │ ├── route_layer.h │ ├── route_layer.o │ ├── shortcut_layer.cpp │ ├── shortcut_layer.h │ ├── shortcut_layer.o │ ├── upsample_layer.cpp │ ├── upsample_layer.h │ └── upsample_layer.o │ ├── libnvdsinfer_custom_impl_Yolo.so │ ├── nvdsinfer_yolo_engine.cpp │ ├── nvdsinfer_yolo_engine.o │ ├── nvdsparsebbox_Yolo.cpp │ ├── nvdsparsebbox_Yolo.o │ ├── utils.cpp │ ├── utils.h │ ├── utils.o │ ├── yolo.cpp │ ├── yolo.h │ ├── yolo.o │ ├── yoloForward.cu │ ├── yoloForward.o │ ├── yoloForward_nc.cu │ ├── yoloForward_nc.o │ ├── yoloForward_r.cu │ ├── yoloForward_r.o │ ├── yoloForward_v2.cu │ ├── yoloForward_v2.o │ ├── yoloPlugins.cpp │ ├── yoloPlugins.h │ └── yoloPlugins.o ├── LICENSE ├── README.md ├── Weights ├── best.cfg ├── best.pt └── best.wts ├── YoloV5 ├── .ipynb_checkpoints │ └── Custom_YoloV5_Deepstream_6_0-checkpoint.ipynb ├── Custom_YoloV5_Deepstream_6_0.ipynb └── gen_wts_yoloV5.py └── assets ├── map.png ├── map.svg ├── optical.jpg ├── pcbinsp.jpg ├── pcbscreenshot.png ├── precision.png ├── precision.svg ├── recall.png └── recall.svg /Dataset/REAME.md: -------------------------------------------------------------------------------- 1 | ## Downloading Dataset 2 | 3 | The dataset for this project is hosted in Roboflow. To download it, you will first need to install Roboflow using Pip, then run the download_dataset.py script. 4 | 5 | `Pip install Roboflow` 6 | 7 | `Python3 download_dataset.py` 8 | -------------------------------------------------------------------------------- /Dataset/download_dataset.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """Download_dataset.ipynb 3 | 4 | Automatically generated by Colaboratory. 5 | 6 | Original file is located at 7 | https://colab.research.google.com/drive/1KKVLgZdBo-w8cLdJpmPCRZXjyHJ67LE2 8 | """ 9 | 10 | from roboflow import Roboflow 11 | rf = Roboflow(api_key="UexS2wDYgkC64r7Zdy4X") 12 | project = rf.workspace("new-workspace-cfyg7").project("pcb-defect-yolov5-deepstream-6.0") 13 | dataset = project.version(2).download("yolov5") -------------------------------------------------------------------------------- /Deepstream-app/config_infer_primary_yoloV5.txt: -------------------------------------------------------------------------------- 1 | [property] 2 | gpu-id=0 3 | net-scale-factor=0.0039215697906911373 4 | model-color-format=0 5 | custom-network-config=best.cfg 6 | model-file=best.wts 7 | model-engine-file=model_b1_gpu0_fp32.engine 8 | #int8-calib-file=calib.table 9 | labelfile-path=labels.txt 10 | batch-size=1 11 | network-mode=0 12 | num-detected-classes=6 13 | interval=0 14 | gie-unique-id=1 15 | process-mode=1 16 | network-type=0 17 | cluster-mode=4 18 | maintain-aspect-ratio=1 19 | parse-bbox-func-name=NvDsInferParseYolo 20 | custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so 21 | engine-create-func-name=NvDsInferYoloCudaEngineGet 22 | 23 | [class-attrs-all] 24 | pre-cluster-threshold=0.5 25 | -------------------------------------------------------------------------------- /Deepstream-app/deepstream_app_config (CSI).txt: -------------------------------------------------------------------------------- 1 | [application] 2 | enable-perf-measurement=1 3 | perf-measurement-interval-sec=5 4 | 5 | [tiled-display] 6 | enable=1 7 | rows=1 8 | columns=1 9 | width=1280 10 | height=720 11 | gpu-id=0 12 | nvbuf-memory-type=0 13 | 14 | [source0] 15 | enable=1 16 | #Type - 1=CameraV4L2 2=URI 3=MultiURI 4=RTSP 5=CSI 17 | type=5 18 | camera-width=1280 19 | camera-height=720 20 | camera-fps-n=30 21 | camera-fps-d=1 22 | 23 | 24 | 25 | [sink0] 26 | enable=1 27 | #Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 5=Overlay 28 | type=5 29 | sync=0 30 | display-id=0 31 | offset-x=0 32 | offset-y=0 33 | width=0 34 | height=0 35 | overlay-id=1 36 | source-id=0 37 | 38 | [sink1] 39 | enable=0 40 | type=3 41 | #1=mp4 2=mkv 42 | container=1 43 | #1=h264 2=h265 3=mpeg4 44 | codec=1 45 | #encoder type 0=Hardware 1=Software 46 | enc-type=0 47 | sync=0 48 | bitrate=2000000 49 | #H264 Profile - 0=Baseline 2=Main 4=High 50 | #H265 Profile - 0=Main 1=Main10 51 | profile=0 52 | output-file=out.mp4 53 | source-id=0 54 | 55 | [sink2] 56 | enable=0 57 | #Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 5=Overlay 58 | type=4 59 | #1=h264 2=h265 60 | codec=1 61 | #encoder type 0=Hardware 1=Software 62 | enc-type=0 63 | sync=0 64 | bitrate=4000000 65 | #H264 Profile - 0=Baseline 2=Main 4=High 66 | #H265 Profile - 0=Main 1=Main10 67 | profile=0 68 | # set below properties in case of RTSPStreaming 69 | rtsp-port=8554 70 | udp-port=5400 71 | 72 | 73 | [osd] 74 | enable=1 75 | gpu-id=0 76 | border-width=1 77 | text-size=15 78 | text-color=1;1;1;1; 79 | text-bg-color=0.3;0.3;0.3;1 80 | font=Serif 81 | show-clock=0 82 | clock-x-offset=800 83 | clock-y-offset=820 84 | clock-text-size=12 85 | clock-color=1;0;0;0 86 | nvbuf-memory-type=0 87 | 88 | [streammux] 89 | ##Boolean property to inform muxer that sources are live 90 | live-source=1 91 | batch-size=1 92 | ##time out in usec, to wait after the first buffer is available 93 | ##to push the batch even if the complete batch is not formed 94 | batched-push-timeout=40000 95 | ## Set muxer output width and height 96 | width=1280 97 | height=720 98 | 99 | 100 | [primary-gie] 101 | enable=1 102 | gpu-id=0 103 | gie-unique-id=1 104 | nvbuf-memory-type=0 105 | config-file=config_infer_primary.txt 106 | 107 | [tracker] 108 | enable = 1 109 | tracker-width=640 110 | tracker-height=384 111 | ll-lib-file=/opt/nvidia/deepstream/deepstream-6.0/lib/libnvds_nvmultiobjecttracker.so 112 | # ll-config-file required to set different tracker types 113 | # ll-config-file=config_tracker_IOU.yml 114 | #ll-config-file=config_tracker_NvDCF_perf.yml 115 | # ll-config-file=config_tracker_NvDCF_accuracy.yml 116 | ll-config-file=config_tracker_DeepSORT.yml 117 | gpu-id=0 118 | enable-batch-process=1 119 | enable-past-frame=1 120 | display-tracking-id=1 121 | 122 | [tests] 123 | file-loop=0 124 | -------------------------------------------------------------------------------- /Deepstream-app/deepstream_app_config.txt: -------------------------------------------------------------------------------- 1 | [application] 2 | enable-perf-measurement=1 3 | perf-measurement-interval-sec=5 4 | 5 | [tiled-display] 6 | enable=1 7 | rows=1 8 | columns=1 9 | width=1280 10 | height=720 11 | gpu-id=0 12 | nvbuf-memory-type=0 13 | 14 | [source0] 15 | enable=1 16 | type=3 17 | uri=file:///opt/nvidia/deepstream/deepstream/samples/streams/pcb2.mp4 18 | num-sources=1 19 | gpu-id=0 20 | cudadec-memtype=0 21 | 22 | 23 | [sink0] 24 | enable=1 25 | type=2 26 | sync=0 27 | gpu-id=0 28 | nvbuf-memory-type=0 29 | 30 | [sink1] 31 | enable=1 32 | type=3 33 | container=1 34 | codec=1 35 | enc-type=0 36 | sync=0 37 | bitrate=2000000 38 | profile=0 39 | output-file=out.mp4 40 | source-id=1 41 | 42 | [osd] 43 | enable=1 44 | gpu-id=0 45 | border-width=1 46 | text-size=15 47 | text-color=1;1;1;1; 48 | text-bg-color=0.3;0.3;0.3;1 49 | font=Serif 50 | show-clock=0 51 | clock-x-offset=800 52 | clock-y-offset=820 53 | clock-text-size=12 54 | clock-color=1;0;0;0 55 | nvbuf-memory-type=0 56 | 57 | [streammux] 58 | gpu-id=0 59 | live-source=0 60 | batch-size=1 61 | batched-push-timeout=40000 62 | width=1280 63 | height=720 64 | enable-padding=0 65 | nvbuf-memory-type=0 66 | 67 | [primary-gie] 68 | enable=1 69 | gpu-id=0 70 | gie-unique-id=1 71 | nvbuf-memory-type=0 72 | config-file=config_infer_primary_yoloV5.txt 73 | 74 | [tests] 75 | file-loop=0 76 | -------------------------------------------------------------------------------- /Deepstream-app/deepstream_app_config_USB.txt: -------------------------------------------------------------------------------- 1 | [application] 2 | enable-perf-measurement=1 3 | perf-measurement-interval-sec=5 4 | 5 | [tiled-display] 6 | enable=1 7 | rows=4 8 | columns=4 9 | width=1280 10 | height=720 11 | gpu-id=0 12 | nvbuf-memory-type=0 13 | 14 | [source0] 15 | enable=1 16 | #Type - 1=CameraV4L2 2=URI 3=MultiURI 17 | type=1 18 | camera-width=320 19 | camera-height=240 20 | camera-fps-n=15 21 | camera-fps-d=1 22 | camera-v4l2-dev-node=0 23 | num-sources=16 24 | 25 | 26 | 27 | [sink0] 28 | enable=1 29 | #Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 5=Overlay 30 | type=5 31 | sync=0 32 | display-id=0 33 | offset-x=0 34 | offset-y=0 35 | width=0 36 | height=0 37 | overlay-id=1 38 | source-id=0 39 | 40 | [sink1] 41 | enable=0 42 | type=3 43 | #1=mp4 2=mkv 44 | container=1 45 | #1=h264 2=h265 3=mpeg4 46 | codec=1 47 | #encoder type 0=Hardware 1=Software 48 | enc-type=0 49 | sync=0 50 | bitrate=2000000 51 | #H264 Profile - 0=Baseline 2=Main 4=High 52 | #H265 Profile - 0=Main 1=Main10 53 | profile=0 54 | output-file=out.mp4 55 | source-id=0 56 | 57 | [sink2] 58 | enable=0 59 | #Type - 1=FakeSink 2=EglSink 3=File 4=RTSPStreaming 5=Overlay 60 | type=4 61 | #1=h264 2=h265 62 | codec=1 63 | #encoder type 0=Hardware 1=Software 64 | enc-type=0 65 | sync=0 66 | bitrate=4000000 67 | #H264 Profile - 0=Baseline 2=Main 4=High 68 | #H265 Profile - 0=Main 1=Main10 69 | profile=0 70 | # set below properties in case of RTSPStreaming 71 | rtsp-port=8554 72 | udp-port=5400 73 | 74 | 75 | [osd] 76 | enable=1 77 | gpu-id=0 78 | border-width=1 79 | text-size=15 80 | text-color=1;1;1;1; 81 | text-bg-color=0.3;0.3;0.3;1 82 | font=Serif 83 | show-clock=0 84 | clock-x-offset=800 85 | clock-y-offset=820 86 | clock-text-size=12 87 | clock-color=1;0;0;0 88 | nvbuf-memory-type=0 89 | 90 | [streammux] 91 | ##Boolean property to inform muxer that sources are live 92 | live-source=1 93 | batch-size=16 94 | ##time out in usec, to wait after the first buffer is available 95 | ##to push the batch even if the complete batch is not formed 96 | batched-push-timeout=40000 97 | ## Set muxer output width and height 98 | width=320 99 | height=240 100 | 101 | 102 | [primary-gie] 103 | enable=1 104 | gpu-id=0 105 | gie-unique-id=1 106 | nvbuf-memory-type=0 107 | config-file=config_infer_primary.txt 108 | 109 | [tracker] 110 | enable = 1 111 | tracker-width=640 112 | tracker-height=384 113 | ll-lib-file=/opt/nvidia/deepstream/deepstream-6.0/lib/libnvds_nvmultiobjecttracker.so 114 | # ll-config-file required to set different tracker types 115 | # ll-config-file=config_tracker_IOU.yml 116 | #ll-config-file=config_tracker_NvDCF_perf.yml 117 | # ll-config-file=config_tracker_NvDCF_accuracy.yml 118 | ll-config-file=config_tracker_DeepSORT.yml 119 | gpu-id=0 120 | enable-batch-process=1 121 | enable-past-frame=1 122 | display-tracking-id=1 123 | 124 | [tests] 125 | file-loop=0 126 | -------------------------------------------------------------------------------- /Deepstream-app/labels.txt: -------------------------------------------------------------------------------- 1 | missing_hole 2 | mouse_bite 3 | open_circuit 4 | short 5 | spur 6 | spurious_copper 7 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/Makefile: -------------------------------------------------------------------------------- 1 | ################################################################################ 2 | # Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | # 4 | # Permission is hereby granted, free of charge, to any person obtaining a 5 | # copy of this software and associated documentation files (the "Software"), 6 | # to deal in the Software without restriction, including without limitation 7 | # the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | # and/or sell copies of the Software, and to permit persons to whom the 9 | # Software is furnished to do so, subject to the following conditions: 10 | # 11 | # The above copyright notice and this permission notice shall be included in 12 | # all copies or substantial portions of the Software. 13 | # 14 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | # DEALINGS IN THE SOFTWARE. 21 | # 22 | # Edited by Marcos Luciano 23 | # https://www.github.com/marcoslucianops 24 | ################################################################################ 25 | 26 | CUDA_VER?= 27 | ifeq ($(CUDA_VER),) 28 | $(error "CUDA_VER is not set") 29 | endif 30 | 31 | OPENCV?= 32 | ifeq ($(OPENCV),) 33 | OPENCV=0 34 | endif 35 | 36 | CC:= g++ 37 | NVCC:=/usr/local/cuda-$(CUDA_VER)/bin/nvcc 38 | 39 | CFLAGS:= -Wall -std=c++11 -shared -fPIC -Wno-error=deprecated-declarations 40 | CFLAGS+= -I/opt/nvidia/deepstream/deepstream/sources/includes -I/usr/local/cuda-$(CUDA_VER)/include 41 | 42 | ifeq ($(OPENCV), 1) 43 | COMMON= -DOPENCV 44 | CFLAGS+= $(shell pkg-config --cflags opencv4 2> /dev/null || pkg-config --cflags opencv) 45 | LIBS+= $(shell pkg-config --libs opencv4 2> /dev/null || pkg-config --libs opencv) 46 | endif 47 | 48 | LIBS+= -lnvinfer_plugin -lnvinfer -lnvparsers -L/usr/local/cuda-$(CUDA_VER)/lib64 -lcudart -lcublas -lstdc++fs 49 | LFLAGS:= -shared -Wl,--start-group $(LIBS) -Wl,--end-group 50 | 51 | INCS:= $(wildcard *.h) 52 | SRCFILES:= nvdsinfer_yolo_engine.cpp \ 53 | nvdsparsebbox_Yolo.cpp \ 54 | yoloPlugins.cpp \ 55 | layers/convolutional_layer.cpp \ 56 | layers/implicit_layer.cpp \ 57 | layers/channels_layer.cpp \ 58 | layers/dropout_layer.cpp \ 59 | layers/shortcut_layer.cpp \ 60 | layers/route_layer.cpp \ 61 | layers/upsample_layer.cpp \ 62 | layers/maxpool_layer.cpp \ 63 | layers/activation_layer.cpp \ 64 | utils.cpp \ 65 | yolo.cpp \ 66 | yoloForward.cu \ 67 | yoloForward_v2.cu \ 68 | yoloForward_nc.cu \ 69 | yoloForward_r.cu 70 | 71 | ifeq ($(OPENCV), 1) 72 | SRCFILES+= calibrator.cpp 73 | endif 74 | 75 | TARGET_LIB:= libnvdsinfer_custom_impl_Yolo.so 76 | 77 | TARGET_OBJS:= $(SRCFILES:.cpp=.o) 78 | TARGET_OBJS:= $(TARGET_OBJS:.cu=.o) 79 | 80 | all: $(TARGET_LIB) 81 | 82 | %.o: %.cpp $(INCS) Makefile 83 | $(CC) -c $(COMMON) -o $@ $(CFLAGS) $< 84 | 85 | %.o: %.cu $(INCS) Makefile 86 | $(NVCC) -c -o $@ --compiler-options '-fPIC' $< 87 | 88 | $(TARGET_LIB) : $(TARGET_OBJS) 89 | $(CC) -o $@ $(TARGET_OBJS) $(LFLAGS) 90 | 91 | clean: 92 | rm -rf $(TARGET_LIB) 93 | rm -rf $(TARGET_OBJS) 94 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/calibrator.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "calibrator.h" 7 | #include 8 | #include 9 | 10 | namespace nvinfer1 11 | { 12 | Int8EntropyCalibrator2::Int8EntropyCalibrator2(const int &batchsize, const int &channels, const int &height, const int &width, const int &letterbox, const std::string &imgPath, 13 | const std::string &calibTablePath):batchSize(batchsize), inputC(channels), inputH(height), inputW(width), letterBox(letterbox), calibTablePath(calibTablePath), imageIndex(0) 14 | { 15 | inputCount = batchsize * channels * height * width; 16 | std::fstream f(imgPath); 17 | if (f.is_open()) 18 | { 19 | std::string temp; 20 | while (std::getline(f, temp)) imgPaths.push_back(temp); 21 | } 22 | batchData = new float[inputCount]; 23 | CUDA_CHECK(cudaMalloc(&deviceInput, inputCount * sizeof(float))); 24 | } 25 | 26 | Int8EntropyCalibrator2::~Int8EntropyCalibrator2() 27 | { 28 | CUDA_CHECK(cudaFree(deviceInput)); 29 | if (batchData) 30 | delete[] batchData; 31 | } 32 | 33 | int Int8EntropyCalibrator2::getBatchSize() const noexcept 34 | { 35 | return batchSize; 36 | } 37 | 38 | bool Int8EntropyCalibrator2::getBatch(void **bindings, const char **names, int nbBindings) noexcept 39 | { 40 | if (imageIndex + batchSize > uint(imgPaths.size())) 41 | return false; 42 | 43 | float* ptr = batchData; 44 | for (size_t j = imageIndex; j < imageIndex + batchSize; ++j) 45 | { 46 | cv::Mat img = cv::imread(imgPaths[j], cv::IMREAD_COLOR); 47 | std::vectorinputData = prepareImage(img, inputC, inputH, inputW, letterBox); 48 | 49 | int len = (int)(inputData.size()); 50 | memcpy(ptr, inputData.data(), len * sizeof(float)); 51 | 52 | ptr += inputData.size(); 53 | std::cout << "Load image: " << imgPaths[j] << std::endl; 54 | std::cout << "Progress: " << (j + 1)*100. / imgPaths.size() << "%" << std::endl; 55 | } 56 | imageIndex += batchSize; 57 | CUDA_CHECK(cudaMemcpy(deviceInput, batchData, inputCount * sizeof(float), cudaMemcpyHostToDevice)); 58 | bindings[0] = deviceInput; 59 | return true; 60 | } 61 | 62 | const void* Int8EntropyCalibrator2::readCalibrationCache(std::size_t &length) noexcept 63 | { 64 | calibrationCache.clear(); 65 | std::ifstream input(calibTablePath, std::ios::binary); 66 | input >> std::noskipws; 67 | if (readCache && input.good()) 68 | { 69 | std::copy(std::istream_iterator(input), std::istream_iterator(), 70 | std::back_inserter(calibrationCache)); 71 | } 72 | length = calibrationCache.size(); 73 | return length ? calibrationCache.data() : nullptr; 74 | } 75 | 76 | void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, std::size_t length) noexcept 77 | { 78 | std::ofstream output(calibTablePath, std::ios::binary); 79 | output.write(reinterpret_cast(cache), length); 80 | } 81 | } 82 | 83 | std::vector prepareImage(cv::Mat& img, int input_c, int input_h, int input_w, int letter_box) 84 | { 85 | cv::Mat out; 86 | int image_w = img.cols; 87 | int image_h = img.rows; 88 | if (image_w != input_w || image_h != input_h) 89 | { 90 | if (letter_box == 1) 91 | { 92 | float ratio_w = (float)image_w / (float)input_w; 93 | float ratio_h = (float)image_h / (float)input_h; 94 | if (ratio_w > ratio_h) 95 | { 96 | int new_width = input_w * ratio_h; 97 | int x = (image_w - new_width) / 2; 98 | cv::Rect roi(abs(x), 0, new_width, image_h); 99 | out = img(roi); 100 | } 101 | else if (ratio_w < ratio_h) 102 | { 103 | int new_height = input_h * ratio_w; 104 | int y = (image_h - new_height) / 2; 105 | cv::Rect roi(0, abs(y), image_w, new_height); 106 | out = img(roi); 107 | } 108 | else { 109 | out = img; 110 | } 111 | cv::resize(out, out, cv::Size(input_w, input_h), 0, 0, cv::INTER_CUBIC); 112 | } 113 | else 114 | { 115 | cv::resize(img, out, cv::Size(input_w, input_h), 0, 0, cv::INTER_CUBIC); 116 | } 117 | cv::cvtColor(out, out, cv::COLOR_BGR2RGB); 118 | } 119 | else 120 | { 121 | cv::cvtColor(img, out, cv::COLOR_BGR2RGB); 122 | } 123 | if (input_c == 3) 124 | { 125 | out.convertTo(out, CV_32FC3, 1.0 / 255.0); 126 | } 127 | else 128 | { 129 | out.convertTo(out, CV_32FC1, 1.0 / 255.0); 130 | } 131 | std::vector input_channels(input_c); 132 | cv::split(out, input_channels); 133 | std::vector result(input_h * input_w * input_c); 134 | auto data = result.data(); 135 | int channelLength = input_h * input_w; 136 | for (int i = 0; i < input_c; ++i) 137 | { 138 | memcpy(data, input_channels[i].data, channelLength * sizeof(float)); 139 | data += channelLength; 140 | } 141 | return result; 142 | } 143 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/calibrator.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef CALIBRATOR_H 7 | #define CALIBRATOR_H 8 | 9 | #include "opencv2/opencv.hpp" 10 | #include "cuda_runtime.h" 11 | #include "NvInfer.h" 12 | #include 13 | #include 14 | 15 | #ifndef CUDA_CHECK 16 | #define CUDA_CHECK(callstr) \ 17 | { \ 18 | cudaError_t error_code = callstr; \ 19 | if (error_code != cudaSuccess) { \ 20 | std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ 21 | assert(0); \ 22 | } \ 23 | } 24 | #endif 25 | 26 | namespace nvinfer1 { 27 | class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 { 28 | public: 29 | Int8EntropyCalibrator2(const int &batchsize, 30 | const int &channels, 31 | const int &height, 32 | const int &width, 33 | const int &letterbox, 34 | const std::string &imgPath, 35 | const std::string &calibTablePath); 36 | 37 | virtual ~Int8EntropyCalibrator2(); 38 | int getBatchSize() const noexcept override; 39 | bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override; 40 | const void* readCalibrationCache(std::size_t& length) noexcept override; 41 | void writeCalibrationCache(const void* cache, size_t length) noexcept override; 42 | 43 | private: 44 | int batchSize; 45 | int inputC; 46 | int inputH; 47 | int inputW; 48 | int letterBox; 49 | std::string calibTablePath; 50 | size_t imageIndex; 51 | size_t inputCount; 52 | std::vector imgPaths; 53 | float *batchData{ nullptr }; 54 | void *deviceInput{ nullptr }; 55 | bool readCache; 56 | std::vector calibrationCache; 57 | }; 58 | } 59 | 60 | std::vector prepareImage(cv::Mat& img, int input_c, int input_h, int input_w, int letter_box); 61 | 62 | #endif //CALIBRATOR_H -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "activation_layer.h" 7 | 8 | nvinfer1::ILayer* activationLayer( 9 | int layerIdx, 10 | std::string activation, 11 | nvinfer1::ILayer* output, 12 | nvinfer1::ITensor* input, 13 | nvinfer1::INetworkDefinition* network) 14 | { 15 | if (activation == "linear") { 16 | // Pass 17 | } 18 | else if (activation == "relu") 19 | { 20 | nvinfer1::IActivationLayer* relu = network->addActivation( 21 | *input, nvinfer1::ActivationType::kRELU); 22 | assert(relu != nullptr); 23 | std::string reluLayerName = "relu_" + std::to_string(layerIdx); 24 | relu->setName(reluLayerName.c_str()); 25 | output = relu; 26 | } 27 | else if (activation == "sigmoid" || activation == "logistic") 28 | { 29 | nvinfer1::IActivationLayer* sigmoid = network->addActivation( 30 | *input, nvinfer1::ActivationType::kSIGMOID); 31 | assert(sigmoid != nullptr); 32 | std::string sigmoidLayerName = "sigmoid_" + std::to_string(layerIdx); 33 | sigmoid->setName(sigmoidLayerName.c_str()); 34 | output = sigmoid; 35 | } 36 | else if (activation == "tanh") 37 | { 38 | nvinfer1::IActivationLayer* tanh = network->addActivation( 39 | *input, nvinfer1::ActivationType::kTANH); 40 | assert(tanh != nullptr); 41 | std::string tanhLayerName = "tanh_" + std::to_string(layerIdx); 42 | tanh->setName(tanhLayerName.c_str()); 43 | output = tanh; 44 | } 45 | else if (activation == "leaky") 46 | { 47 | nvinfer1::IActivationLayer* leaky = network->addActivation( 48 | *input, nvinfer1::ActivationType::kLEAKY_RELU); 49 | leaky->setAlpha(0.1); 50 | assert(leaky != nullptr); 51 | std::string leakyLayerName = "leaky_" + std::to_string(layerIdx); 52 | leaky->setName(leakyLayerName.c_str()); 53 | output = leaky; 54 | } 55 | else if (activation == "softplus") 56 | { 57 | nvinfer1::IActivationLayer* softplus = network->addActivation( 58 | *input, nvinfer1::ActivationType::kSOFTPLUS); 59 | assert(softplus != nullptr); 60 | std::string softplusLayerName = "softplus_" + std::to_string(layerIdx); 61 | softplus->setName(softplusLayerName.c_str()); 62 | output = softplus; 63 | } 64 | else if (activation == "mish") 65 | { 66 | nvinfer1::IActivationLayer* softplus = network->addActivation( 67 | *input, nvinfer1::ActivationType::kSOFTPLUS); 68 | assert(softplus != nullptr); 69 | std::string softplusLayerName = "softplus_" + std::to_string(layerIdx); 70 | softplus->setName(softplusLayerName.c_str()); 71 | nvinfer1::IActivationLayer* tanh = network->addActivation( 72 | *softplus->getOutput(0), nvinfer1::ActivationType::kTANH); 73 | assert(tanh != nullptr); 74 | std::string tanhLayerName = "tanh_" + std::to_string(layerIdx); 75 | tanh->setName(tanhLayerName.c_str()); 76 | nvinfer1::IElementWiseLayer* mish = network->addElementWise( 77 | *tanh->getOutput(0), *input, 78 | nvinfer1::ElementWiseOperation::kPROD); 79 | assert(mish != nullptr); 80 | std::string mishLayerName = "mish_" + std::to_string(layerIdx); 81 | mish->setName(mishLayerName.c_str()); 82 | output = mish; 83 | } 84 | else if (activation == "silu") 85 | { 86 | nvinfer1::IActivationLayer* sigmoid = network->addActivation( 87 | *input, nvinfer1::ActivationType::kSIGMOID); 88 | assert(sigmoid != nullptr); 89 | std::string sigmoidLayerName = "sigmoid_" + std::to_string(layerIdx); 90 | sigmoid->setName(sigmoidLayerName.c_str()); 91 | nvinfer1::IElementWiseLayer* silu = network->addElementWise( 92 | *sigmoid->getOutput(0), *input, 93 | nvinfer1::ElementWiseOperation::kPROD); 94 | assert(silu != nullptr); 95 | std::string siluLayerName = "silu_" + std::to_string(layerIdx); 96 | silu->setName(siluLayerName.c_str()); 97 | output = silu; 98 | } 99 | else { 100 | std::cerr << "Activation not supported: " << activation << std::endl; 101 | std::abort(); 102 | } 103 | return output; 104 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/activation_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __ACTIVATION_LAYER_H__ 7 | #define __ACTIVATION_LAYER_H__ 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include "NvInfer.h" 14 | 15 | #include "activation_layer.h" 16 | 17 | nvinfer1::ILayer* activationLayer( 18 | int layerIdx, 19 | std::string activation, 20 | nvinfer1::ILayer* output, 21 | nvinfer1::ITensor* input, 22 | nvinfer1::INetworkDefinition* network); 23 | 24 | #endif 25 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/activation_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/activation_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/channels_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "channels_layer.h" 7 | 8 | nvinfer1::ILayer* channelsLayer( 9 | std::string type, 10 | nvinfer1::ITensor* input, 11 | nvinfer1::ITensor* implicitTensor, 12 | nvinfer1::INetworkDefinition* network) 13 | { 14 | nvinfer1::ILayer* output; 15 | 16 | if (type == "shift") { 17 | nvinfer1::IElementWiseLayer* ew = network->addElementWise( 18 | *input, *implicitTensor, 19 | nvinfer1::ElementWiseOperation::kSUM); 20 | assert(ew != nullptr); 21 | output = ew; 22 | } 23 | else if (type == "control") { 24 | nvinfer1::IElementWiseLayer* ew = network->addElementWise( 25 | *input, *implicitTensor, 26 | nvinfer1::ElementWiseOperation::kPROD); 27 | assert(ew != nullptr); 28 | output = ew; 29 | } 30 | 31 | return output; 32 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/channels_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __CHANNELS_LAYER_H__ 7 | #define __CHANNELS_LAYER_H__ 8 | 9 | #include 10 | #include 11 | 12 | #include "NvInfer.h" 13 | 14 | nvinfer1::ILayer* channelsLayer( 15 | std::string type, 16 | nvinfer1::ITensor* input, 17 | nvinfer1::ITensor* implicitTensor, 18 | nvinfer1::INetworkDefinition* network); 19 | 20 | #endif 21 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/channels_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/channels_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include 7 | #include "convolutional_layer.h" 8 | 9 | nvinfer1::ILayer* convolutionalLayer( 10 | int layerIdx, 11 | std::map& block, 12 | std::vector& weights, 13 | std::vector& trtWeights, 14 | int& weightPtr, 15 | std::string weightsType, 16 | int& inputChannels, 17 | nvinfer1::ITensor* input, 18 | nvinfer1::INetworkDefinition* network) 19 | { 20 | assert(block.at("type") == "convolutional"); 21 | assert(block.find("filters") != block.end()); 22 | assert(block.find("pad") != block.end()); 23 | assert(block.find("size") != block.end()); 24 | assert(block.find("stride") != block.end()); 25 | 26 | int filters = std::stoi(block.at("filters")); 27 | int padding = std::stoi(block.at("pad")); 28 | int kernelSize = std::stoi(block.at("size")); 29 | int stride = std::stoi(block.at("stride")); 30 | std::string activation = block.at("activation"); 31 | int bias = filters; 32 | 33 | bool batchNormalize = false; 34 | if (block.find("batch_normalize") != block.end()) 35 | { 36 | bias = 0; 37 | batchNormalize = (block.at("batch_normalize") == "1"); 38 | } 39 | 40 | int groups = 1; 41 | if (block.find("groups") != block.end()) 42 | { 43 | groups = std::stoi(block.at("groups")); 44 | } 45 | 46 | int pad; 47 | if (padding) 48 | pad = (kernelSize - 1) / 2; 49 | else 50 | pad = 0; 51 | 52 | int size = filters * inputChannels * kernelSize * kernelSize / groups; 53 | std::vector bnBiases; 54 | std::vector bnWeights; 55 | std::vector bnRunningMean; 56 | std::vector bnRunningVar; 57 | nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, size}; 58 | nvinfer1::Weights convBias{nvinfer1::DataType::kFLOAT, nullptr, bias}; 59 | 60 | if (weightsType == "weights") { 61 | if (batchNormalize == false) 62 | { 63 | float* val = new float[filters]; 64 | for (int i = 0; i < filters; ++i) 65 | { 66 | val[i] = weights[weightPtr]; 67 | weightPtr++; 68 | } 69 | convBias.values = val; 70 | trtWeights.push_back(convBias); 71 | val = new float[size]; 72 | for (int i = 0; i < size; ++i) 73 | { 74 | val[i] = weights[weightPtr]; 75 | weightPtr++; 76 | } 77 | convWt.values = val; 78 | trtWeights.push_back(convWt); 79 | } 80 | else 81 | { 82 | for (int i = 0; i < filters; ++i) 83 | { 84 | bnBiases.push_back(weights[weightPtr]); 85 | weightPtr++; 86 | } 87 | for (int i = 0; i < filters; ++i) 88 | { 89 | bnWeights.push_back(weights[weightPtr]); 90 | weightPtr++; 91 | } 92 | for (int i = 0; i < filters; ++i) 93 | { 94 | bnRunningMean.push_back(weights[weightPtr]); 95 | weightPtr++; 96 | } 97 | for (int i = 0; i < filters; ++i) 98 | { 99 | bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5)); 100 | weightPtr++; 101 | } 102 | float* val = new float[size]; 103 | for (int i = 0; i < size; ++i) 104 | { 105 | val[i] = weights[weightPtr]; 106 | weightPtr++; 107 | } 108 | convWt.values = val; 109 | trtWeights.push_back(convWt); 110 | trtWeights.push_back(convBias); 111 | } 112 | } 113 | else { 114 | if (batchNormalize == false) 115 | { 116 | float* val = new float[size]; 117 | for (int i = 0; i < size; ++i) 118 | { 119 | val[i] = weights[weightPtr]; 120 | weightPtr++; 121 | } 122 | convWt.values = val; 123 | trtWeights.push_back(convWt); 124 | val = new float[filters]; 125 | for (int i = 0; i < filters; ++i) 126 | { 127 | val[i] = weights[weightPtr]; 128 | weightPtr++; 129 | } 130 | convBias.values = val; 131 | trtWeights.push_back(convBias); 132 | } 133 | else 134 | { 135 | float* val = new float[size]; 136 | for (int i = 0; i < size; ++i) 137 | { 138 | val[i] = weights[weightPtr]; 139 | weightPtr++; 140 | } 141 | convWt.values = val; 142 | for (int i = 0; i < filters; ++i) 143 | { 144 | bnWeights.push_back(weights[weightPtr]); 145 | weightPtr++; 146 | } 147 | for (int i = 0; i < filters; ++i) 148 | { 149 | bnBiases.push_back(weights[weightPtr]); 150 | weightPtr++; 151 | } 152 | for (int i = 0; i < filters; ++i) 153 | { 154 | bnRunningMean.push_back(weights[weightPtr]); 155 | weightPtr++; 156 | } 157 | for (int i = 0; i < filters; ++i) 158 | { 159 | bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5)); 160 | weightPtr++; 161 | } 162 | trtWeights.push_back(convWt); 163 | trtWeights.push_back(convBias); 164 | } 165 | } 166 | 167 | nvinfer1::IConvolutionLayer* conv = network->addConvolutionNd( 168 | *input, filters, nvinfer1::DimsHW{kernelSize, kernelSize}, convWt, convBias); 169 | assert(conv != nullptr); 170 | std::string convLayerName = "conv_" + std::to_string(layerIdx); 171 | conv->setName(convLayerName.c_str()); 172 | conv->setStrideNd(nvinfer1::DimsHW{stride, stride}); 173 | conv->setPaddingNd(nvinfer1::DimsHW{pad, pad}); 174 | 175 | if (block.find("groups") != block.end()) 176 | { 177 | conv->setNbGroups(groups); 178 | } 179 | 180 | nvinfer1::ILayer* output = conv; 181 | 182 | if (batchNormalize == true) 183 | { 184 | size = filters; 185 | nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size}; 186 | nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size}; 187 | nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size}; 188 | float* shiftWt = new float[size]; 189 | for (int i = 0; i < size; ++i) 190 | { 191 | shiftWt[i] 192 | = bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i)); 193 | } 194 | shift.values = shiftWt; 195 | float* scaleWt = new float[size]; 196 | for (int i = 0; i < size; ++i) 197 | { 198 | scaleWt[i] = bnWeights.at(i) / bnRunningVar[i]; 199 | } 200 | scale.values = scaleWt; 201 | float* powerWt = new float[size]; 202 | for (int i = 0; i < size; ++i) 203 | { 204 | powerWt[i] = 1.0; 205 | } 206 | power.values = powerWt; 207 | trtWeights.push_back(shift); 208 | trtWeights.push_back(scale); 209 | trtWeights.push_back(power); 210 | 211 | nvinfer1::IScaleLayer* bn = network->addScale( 212 | *output->getOutput(0), nvinfer1::ScaleMode::kCHANNEL, shift, scale, power); 213 | assert(bn != nullptr); 214 | std::string bnLayerName = "batch_norm_" + std::to_string(layerIdx); 215 | bn->setName(bnLayerName.c_str()); 216 | output = bn; 217 | } 218 | 219 | output = activationLayer(layerIdx, activation, output, output->getOutput(0), network); 220 | assert(output != nullptr); 221 | 222 | return output; 223 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __CONVOLUTIONAL_LAYER_H__ 7 | #define __CONVOLUTIONAL_LAYER_H__ 8 | 9 | #include 10 | #include 11 | 12 | #include "NvInfer.h" 13 | 14 | #include "activation_layer.h" 15 | 16 | nvinfer1::ILayer* convolutionalLayer( 17 | int layerIdx, 18 | std::map& block, 19 | std::vector& weights, 20 | std::vector& trtWeights, 21 | int& weightPtr, 22 | std::string weightsType, 23 | int& inputChannels, 24 | nvinfer1::ITensor* input, 25 | nvinfer1::INetworkDefinition* network); 26 | 27 | #endif 28 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "dropout_layer.h" 7 | 8 | nvinfer1::ILayer* dropoutLayer( 9 | float probability, 10 | nvinfer1::ITensor* input, 11 | nvinfer1::INetworkDefinition* network) 12 | { 13 | nvinfer1::ILayer* output; 14 | return output; 15 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __DROPOUT_LAYER_H__ 7 | #define __DROPOUT_LAYER_H__ 8 | 9 | #include "NvInfer.h" 10 | 11 | nvinfer1::ILayer* dropoutLayer( 12 | float probability, 13 | nvinfer1::ITensor* input, 14 | nvinfer1::INetworkDefinition* network); 15 | 16 | #endif 17 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/dropout_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/implicit_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include 7 | #include "implicit_layer.h" 8 | 9 | nvinfer1::ILayer* implicitLayer( 10 | int channels, 11 | std::vector& weights, 12 | std::vector& trtWeights, 13 | int& weightPtr, 14 | nvinfer1::INetworkDefinition* network) 15 | { 16 | nvinfer1::Weights convWt{nvinfer1::DataType::kFLOAT, nullptr, channels}; 17 | 18 | float* val = new float[channels]; 19 | for (int i = 0; i < channels; ++i) 20 | { 21 | val[i] = weights[weightPtr]; 22 | weightPtr++; 23 | } 24 | convWt.values = val; 25 | trtWeights.push_back(convWt); 26 | 27 | nvinfer1::IConstantLayer* implicit = network->addConstant(nvinfer1::Dims3{static_cast(channels), 1, 1}, convWt); 28 | assert(implicit != nullptr); 29 | 30 | return implicit; 31 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/implicit_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __IMPLICIT_LAYER_H__ 7 | #define __IMPLICIT_LAYER_H__ 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include "NvInfer.h" 14 | 15 | nvinfer1::ILayer* implicitLayer( 16 | int channels, 17 | std::vector& weights, 18 | std::vector& trtWeights, 19 | int& weightPtr, 20 | nvinfer1::INetworkDefinition* network); 21 | 22 | #endif 23 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/implicit_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/implicit_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "maxpool_layer.h" 7 | 8 | nvinfer1::ILayer* maxpoolLayer( 9 | int layerIdx, 10 | std::map& block, 11 | nvinfer1::ITensor* input, 12 | nvinfer1::INetworkDefinition* network) 13 | { 14 | assert(block.at("type") == "maxpool"); 15 | assert(block.find("size") != block.end()); 16 | assert(block.find("stride") != block.end()); 17 | 18 | int size = std::stoi(block.at("size")); 19 | int stride = std::stoi(block.at("stride")); 20 | 21 | nvinfer1::IPoolingLayer* pool 22 | = network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{size, size}); 23 | assert(pool); 24 | std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx); 25 | pool->setStrideNd(nvinfer1::DimsHW{stride, stride}); 26 | pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER); 27 | pool->setName(maxpoolLayerName.c_str()); 28 | 29 | return pool; 30 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __MAXPOOL_LAYER_H__ 7 | #define __MAXPOOL_LAYER_H__ 8 | 9 | #include 10 | #include 11 | 12 | #include "NvInfer.h" 13 | 14 | nvinfer1::ILayer* maxpoolLayer( 15 | int layerIdx, 16 | std::map& block, 17 | nvinfer1::ITensor* input, 18 | nvinfer1::INetworkDefinition* network); 19 | 20 | #endif 21 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "route_layer.h" 7 | 8 | nvinfer1::ILayer* routeLayer( 9 | int layerIdx, 10 | std::map& block, 11 | std::vector tensorOutputs, 12 | nvinfer1::INetworkDefinition* network) 13 | { 14 | std::string strLayers = block.at("layers"); 15 | std::vector idxLayers; 16 | size_t lastPos = 0, pos = 0; 17 | while ((pos = strLayers.find(',', lastPos)) != std::string::npos) { 18 | int vL = std::stoi(trim(strLayers.substr(lastPos, pos - lastPos))); 19 | idxLayers.push_back (vL); 20 | lastPos = pos + 1; 21 | } 22 | if (lastPos < strLayers.length()) { 23 | std::string lastV = trim(strLayers.substr(lastPos)); 24 | if (!lastV.empty()) { 25 | idxLayers.push_back (std::stoi(lastV)); 26 | } 27 | } 28 | assert (!idxLayers.empty()); 29 | std::vector concatInputs; 30 | for (int idxLayer : idxLayers) { 31 | if (idxLayer < 0) { 32 | idxLayer = tensorOutputs.size() + idxLayer; 33 | } 34 | assert (idxLayer >= 0 && idxLayer < (int)tensorOutputs.size()); 35 | concatInputs.push_back (tensorOutputs[idxLayer]); 36 | } 37 | 38 | nvinfer1::IConcatenationLayer* concat = 39 | network->addConcatenation(concatInputs.data(), concatInputs.size()); 40 | assert(concat != nullptr); 41 | std::string concatLayerName = "route_" + std::to_string(layerIdx - 1); 42 | concat->setName(concatLayerName.c_str()); 43 | concat->setAxis(0); 44 | 45 | nvinfer1::ILayer* output = concat; 46 | 47 | if (block.find("groups") != block.end()) { 48 | nvinfer1::Dims prevTensorDims = output->getOutput(0)->getDimensions(); 49 | int groups = stoi(block.at("groups")); 50 | int group_id = stoi(block.at("group_id")); 51 | int startSlice = (prevTensorDims.d[0] / groups) * group_id; 52 | int channelSlice = (prevTensorDims.d[0] / groups); 53 | nvinfer1::ISliceLayer* sl = network->addSlice( 54 | *output->getOutput(0), 55 | nvinfer1::Dims3{startSlice, 0, 0}, 56 | nvinfer1::Dims3{channelSlice, prevTensorDims.d[1], prevTensorDims.d[2]}, 57 | nvinfer1::Dims3{1, 1, 1}); 58 | assert(sl != nullptr); 59 | output = sl; 60 | } 61 | 62 | return output; 63 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/route_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __ROUTE_LAYER_H__ 7 | #define __ROUTE_LAYER_H__ 8 | 9 | #include "NvInfer.h" 10 | #include "../utils.h" 11 | 12 | nvinfer1::ILayer* routeLayer( 13 | int layerIdx, 14 | std::map& block, 15 | std::vector tensorOutputs, 16 | nvinfer1::INetworkDefinition* network); 17 | 18 | #endif -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/route_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/route_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "shortcut_layer.h" 7 | 8 | nvinfer1::ILayer* shortcutLayer( 9 | int layerIdx, 10 | std::string activation, 11 | std::string inputVol, 12 | std::string shortcutVol, 13 | nvinfer1::ITensor* input, 14 | nvinfer1::ITensor* shortcutTensor, 15 | nvinfer1::INetworkDefinition* network) 16 | { 17 | nvinfer1::ILayer* output; 18 | nvinfer1::ITensor* outputTensor; 19 | 20 | if (inputVol != shortcutVol) 21 | { 22 | nvinfer1::ISliceLayer* sl = network->addSlice( 23 | *shortcutTensor, 24 | nvinfer1::Dims3{0, 0, 0}, 25 | input->getDimensions(), 26 | nvinfer1::Dims3{1, 1, 1}); 27 | assert(sl != nullptr); 28 | outputTensor = sl->getOutput(0); 29 | assert(outputTensor != nullptr); 30 | } else 31 | { 32 | outputTensor = shortcutTensor; 33 | assert(outputTensor != nullptr); 34 | } 35 | 36 | nvinfer1::IElementWiseLayer* ew = network->addElementWise( 37 | *input, *outputTensor, 38 | nvinfer1::ElementWiseOperation::kSUM); 39 | assert(ew != nullptr); 40 | 41 | output = activationLayer(layerIdx, activation, ew, ew->getOutput(0), network); 42 | assert(output != nullptr); 43 | 44 | return output; 45 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __SHORTCUT_LAYER_H__ 7 | #define __SHORTCUT_LAYER_H__ 8 | 9 | #include "NvInfer.h" 10 | 11 | #include "activation_layer.h" 12 | 13 | nvinfer1::ILayer* shortcutLayer( 14 | int layerIdx, 15 | std::string activation, 16 | std::string inputVol, 17 | std::string shortcutVol, 18 | nvinfer1::ITensor* input, 19 | nvinfer1::ITensor* shortcutTensor, 20 | nvinfer1::INetworkDefinition* network); 21 | 22 | #endif 23 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include "upsample_layer.h" 7 | 8 | nvinfer1::ILayer* upsampleLayer( 9 | int layerIdx, 10 | std::map& block, 11 | nvinfer1::ITensor* input, 12 | nvinfer1::INetworkDefinition* network) 13 | { 14 | assert(block.at("type") == "upsample"); 15 | int stride = std::stoi(block.at("stride")); 16 | 17 | nvinfer1::IResizeLayer* resize_layer = network->addResize(*input); 18 | resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST); 19 | float scale[3] = {1, static_cast(stride), static_cast(stride)}; 20 | resize_layer->setScales(scale, 3); 21 | std::string layer_name = "upsample_" + std::to_string(layerIdx); 22 | resize_layer->setName(layer_name.c_str()); 23 | return resize_layer; 24 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #ifndef __UPSAMPLE_LAYER_H__ 7 | #define __UPSAMPLE_LAYER_H__ 8 | 9 | #include 10 | #include 11 | #include 12 | 13 | #include "NvInfer.h" 14 | 15 | nvinfer1::ILayer* upsampleLayer( 16 | int layerIdx, 17 | std::map& block, 18 | nvinfer1::ITensor* input, 19 | nvinfer1::INetworkDefinition* network); 20 | 21 | #endif 22 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/layers/upsample_layer.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #include "nvdsinfer_custom_impl.h" 27 | #include "nvdsinfer_context.h" 28 | #include "yoloPlugins.h" 29 | #include "yolo.h" 30 | 31 | #include 32 | 33 | #define USE_CUDA_ENGINE_GET_API 1 34 | 35 | static bool getYoloNetworkInfo (NetworkInfo &networkInfo, const NvDsInferContextInitParams* initParams) 36 | { 37 | std::string yoloCfg = initParams->customNetworkConfigFilePath; 38 | std::string yoloType; 39 | 40 | std::transform (yoloCfg.begin(), yoloCfg.end(), yoloCfg.begin(), [] (uint8_t c) { 41 | return std::tolower (c);}); 42 | 43 | yoloType = yoloCfg.substr(0, yoloCfg.find(".cfg")); 44 | 45 | networkInfo.networkType = yoloType; 46 | networkInfo.configFilePath = initParams->customNetworkConfigFilePath; 47 | networkInfo.wtsFilePath = initParams->modelFilePath; 48 | networkInfo.int8CalibPath = initParams->int8CalibrationFilePath; 49 | networkInfo.deviceType = (initParams->useDLA ? "kDLA" : "kGPU"); 50 | networkInfo.inputBlobName = "data"; 51 | 52 | if(initParams->networkMode == 0) { 53 | networkInfo.networkMode = "FP32"; 54 | } 55 | else if(initParams->networkMode == 1) { 56 | networkInfo.networkMode = "INT8"; 57 | } 58 | else if(initParams->networkMode == 2) { 59 | networkInfo.networkMode = "FP16"; 60 | } 61 | 62 | if (networkInfo.configFilePath.empty() || 63 | networkInfo.wtsFilePath.empty()) { 64 | std::cerr << "YOLO config file or weights file is not specified" 65 | << std::endl; 66 | return false; 67 | } 68 | 69 | if (!fileExists(networkInfo.configFilePath) || 70 | !fileExists(networkInfo.wtsFilePath)) { 71 | std::cerr << "YOLO config file or weights file is not exist" 72 | << std::endl; 73 | return false; 74 | } 75 | 76 | return true; 77 | } 78 | 79 | #if !USE_CUDA_ENGINE_GET_API 80 | IModelParser* NvDsInferCreateModelParser( 81 | const NvDsInferContextInitParams* initParams) { 82 | NetworkInfo networkInfo; 83 | if (!getYoloNetworkInfo(networkInfo, initParams)) { 84 | return nullptr; 85 | } 86 | 87 | return new Yolo(networkInfo); 88 | } 89 | #else 90 | extern "C" 91 | bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, 92 | nvinfer1::IBuilderConfig * const builderConfig, 93 | const NvDsInferContextInitParams * const initParams, 94 | nvinfer1::DataType dataType, 95 | nvinfer1::ICudaEngine *& cudaEngine); 96 | 97 | extern "C" 98 | bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder, 99 | nvinfer1::IBuilderConfig * const builderConfig, 100 | const NvDsInferContextInitParams * const initParams, 101 | nvinfer1::DataType dataType, 102 | nvinfer1::ICudaEngine *& cudaEngine) 103 | { 104 | NetworkInfo networkInfo; 105 | if (!getYoloNetworkInfo(networkInfo, initParams)) { 106 | return false; 107 | } 108 | 109 | Yolo yolo(networkInfo); 110 | cudaEngine = yolo.createEngine (builder); 111 | if (cudaEngine == nullptr) 112 | { 113 | std::cerr << "Failed to build CUDA engine on " 114 | << networkInfo.configFilePath << std::endl; 115 | return false; 116 | } 117 | 118 | return true; 119 | } 120 | #endif 121 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #include 27 | #include 28 | #include 29 | #include "nvdsinfer_custom_impl.h" 30 | #include "utils.h" 31 | 32 | #include "yoloPlugins.h" 33 | 34 | extern "C" bool NvDsInferParseYolo( 35 | std::vector const& outputLayersInfo, 36 | NvDsInferNetworkInfo const& networkInfo, 37 | NvDsInferParseDetectionParams const& detectionParams, 38 | std::vector& objectList); 39 | 40 | static std::vector 41 | nonMaximumSuppression(const float nmsThresh, std::vector binfo) 42 | { 43 | auto overlap1D = [](float x1min, float x1max, float x2min, float x2max) -> float { 44 | if (x1min > x2min) 45 | { 46 | std::swap(x1min, x2min); 47 | std::swap(x1max, x2max); 48 | } 49 | return x1max < x2min ? 0 : std::min(x1max, x2max) - x2min; 50 | }; 51 | auto computeIoU 52 | = [&overlap1D](NvDsInferParseObjectInfo& bbox1, NvDsInferParseObjectInfo& bbox2) -> float { 53 | float overlapX 54 | = overlap1D(bbox1.left, bbox1.left + bbox1.width, bbox2.left, bbox2.left + bbox2.width); 55 | float overlapY 56 | = overlap1D(bbox1.top, bbox1.top + bbox1.height, bbox2.top, bbox2.top + bbox2.height); 57 | float area1 = (bbox1.width) * (bbox1.height); 58 | float area2 = (bbox2.width) * (bbox2.height); 59 | float overlap2D = overlapX * overlapY; 60 | float u = area1 + area2 - overlap2D; 61 | return u == 0 ? 0 : overlap2D / u; 62 | }; 63 | 64 | std::stable_sort(binfo.begin(), binfo.end(), 65 | [](const NvDsInferParseObjectInfo& b1, const NvDsInferParseObjectInfo& b2) { 66 | return b1.detectionConfidence > b2.detectionConfidence; 67 | }); 68 | std::vector out; 69 | for (auto i : binfo) 70 | { 71 | bool keep = true; 72 | for (auto j : out) 73 | { 74 | if (keep) 75 | { 76 | float overlap = computeIoU(i, j); 77 | keep = overlap <= nmsThresh; 78 | } 79 | else 80 | break; 81 | } 82 | if (keep) out.push_back(i); 83 | } 84 | return out; 85 | } 86 | 87 | static std::vector 88 | nmsAllClasses(const float nmsThresh, 89 | std::vector& binfo, 90 | const uint numClasses) 91 | { 92 | std::vector result; 93 | std::vector> splitBoxes(numClasses); 94 | for (auto& box : binfo) 95 | { 96 | splitBoxes.at(box.classId).push_back(box); 97 | } 98 | 99 | for (auto& boxes : splitBoxes) 100 | { 101 | boxes = nonMaximumSuppression(nmsThresh, boxes); 102 | result.insert(result.end(), boxes.begin(), boxes.end()); 103 | } 104 | return result; 105 | } 106 | 107 | static NvDsInferParseObjectInfo convertBBox(const float& bx, const float& by, const float& bw, 108 | const float& bh, const int& stride, const uint& netW, 109 | const uint& netH) 110 | { 111 | NvDsInferParseObjectInfo b; 112 | float xCenter = bx * stride; 113 | float yCenter = by * stride; 114 | float x0 = xCenter - bw / 2; 115 | float y0 = yCenter - bh / 2; 116 | float x1 = x0 + bw; 117 | float y1 = y0 + bh; 118 | 119 | x0 = clamp(x0, 0, netW); 120 | y0 = clamp(y0, 0, netH); 121 | x1 = clamp(x1, 0, netW); 122 | y1 = clamp(y1, 0, netH); 123 | 124 | b.left = x0; 125 | b.width = clamp(x1 - x0, 0, netW); 126 | b.top = y0; 127 | b.height = clamp(y1 - y0, 0, netH); 128 | 129 | return b; 130 | } 131 | 132 | static void addBBoxProposal(const float bx, const float by, const float bw, const float bh, 133 | const uint stride, const uint& netW, const uint& netH, const int maxIndex, 134 | const float maxProb, std::vector& binfo) 135 | { 136 | NvDsInferParseObjectInfo bbi = convertBBox(bx, by, bw, bh, stride, netW, netH); 137 | if (bbi.width < 1 || bbi.height < 1) return; 138 | 139 | bbi.detectionConfidence = maxProb; 140 | bbi.classId = maxIndex; 141 | binfo.push_back(bbi); 142 | } 143 | 144 | static std::vector 145 | decodeYoloTensor( 146 | const float* detections, const std::vector &mask, const std::vector &anchors, 147 | const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, 148 | const uint numOutputClasses, const uint& netW, 149 | const uint& netH, 150 | const float confThresh) 151 | { 152 | std::vector binfo; 153 | for (uint y = 0; y < gridSizeH; ++y) { 154 | for (uint x = 0; x < gridSizeW; ++x) { 155 | for (uint b = 0; b < numBBoxes; ++b) 156 | { 157 | const float pw = anchors[mask[b] * 2]; 158 | const float ph = anchors[mask[b] * 2 + 1]; 159 | 160 | const int numGridCells = gridSizeH * gridSizeW; 161 | const int bbindex = y * gridSizeW + x; 162 | const float bx 163 | = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; 164 | const float by 165 | = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; 166 | const float bw 167 | = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]; 168 | const float bh 169 | = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]; 170 | 171 | const float objectness 172 | = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; 173 | 174 | float maxProb = 0.0f; 175 | int maxIndex = -1; 176 | 177 | for (uint i = 0; i < numOutputClasses; ++i) 178 | { 179 | float prob 180 | = (detections[bbindex 181 | + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); 182 | 183 | if (prob > maxProb) 184 | { 185 | maxProb = prob; 186 | maxIndex = i; 187 | } 188 | } 189 | maxProb = objectness * maxProb; 190 | 191 | if (maxProb > confThresh) 192 | { 193 | addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); 194 | } 195 | } 196 | } 197 | } 198 | return binfo; 199 | } 200 | 201 | static std::vector 202 | decodeYoloV2Tensor( 203 | const float* detections, const std::vector &anchors, 204 | const uint gridSizeW, const uint gridSizeH, const uint stride, const uint numBBoxes, 205 | const uint numOutputClasses, const uint& netW, 206 | const uint& netH) 207 | { 208 | std::vector binfo; 209 | for (uint y = 0; y < gridSizeH; ++y) { 210 | for (uint x = 0; x < gridSizeW; ++x) { 211 | for (uint b = 0; b < numBBoxes; ++b) 212 | { 213 | const float pw = anchors[b * 2]; 214 | const float ph = anchors[b * 2 + 1]; 215 | 216 | const int numGridCells = gridSizeH * gridSizeW; 217 | const int bbindex = y * gridSizeW + x; 218 | const float bx 219 | = x + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 0)]; 220 | const float by 221 | = y + detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 1)]; 222 | const float bw 223 | = pw * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 2)]; 224 | const float bh 225 | = ph * detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 3)]; 226 | 227 | const float objectness 228 | = detections[bbindex + numGridCells * (b * (5 + numOutputClasses) + 4)]; 229 | 230 | float maxProb = 0.0f; 231 | int maxIndex = -1; 232 | 233 | for (uint i = 0; i < numOutputClasses; ++i) 234 | { 235 | float prob 236 | = (detections[bbindex 237 | + numGridCells * (b * (5 + numOutputClasses) + (5 + i))]); 238 | 239 | if (prob > maxProb) 240 | { 241 | maxProb = prob; 242 | maxIndex = i; 243 | } 244 | } 245 | maxProb = objectness * maxProb; 246 | 247 | addBBoxProposal(bx, by, bw, bh, stride, netW, netH, maxIndex, maxProb, binfo); 248 | } 249 | } 250 | } 251 | return binfo; 252 | } 253 | 254 | static inline std::vector 255 | SortLayers(const std::vector & outputLayersInfo) 256 | { 257 | std::vector outLayers; 258 | for (auto const &layer : outputLayersInfo) { 259 | outLayers.push_back (&layer); 260 | } 261 | std::sort(outLayers.begin(), outLayers.end(), 262 | [](const NvDsInferLayerInfo* a, const NvDsInferLayerInfo* b) { 263 | return a->inferDims.d[1] < b->inferDims.d[1]; 264 | }); 265 | return outLayers; 266 | } 267 | 268 | static bool NvDsInferParseYolo( 269 | std::vector const& outputLayersInfo, 270 | NvDsInferNetworkInfo const& networkInfo, 271 | NvDsInferParseDetectionParams const& detectionParams, 272 | std::vector& objectList, 273 | const std::vector &anchors, 274 | const std::vector> &masks, 275 | const uint &num_classes, 276 | const float &beta_nms) 277 | { 278 | const float kCONF_THRESH = detectionParams.perClassThreshold[0]; 279 | 280 | const std::vector sortedLayers = 281 | SortLayers (outputLayersInfo); 282 | 283 | if (sortedLayers.size() != masks.size()) { 284 | std::cerr << "ERROR: YOLO output layer.size: " << sortedLayers.size() 285 | << " does not match mask.size: " << masks.size() << std::endl; 286 | return false; 287 | } 288 | 289 | if (num_classes != detectionParams.numClassesConfigured) 290 | { 291 | std::cerr << "WARNING: Num classes mismatch. Configured: " 292 | << detectionParams.numClassesConfigured 293 | << ", detected by network: " << num_classes << std::endl; 294 | } 295 | 296 | std::vector objects; 297 | 298 | for (uint idx = 0; idx < masks.size(); ++idx) { 299 | const NvDsInferLayerInfo &layer = *sortedLayers[idx]; // 255 x Grid x Grid 300 | 301 | assert(layer.inferDims.numDims == 3); 302 | const uint gridSizeH = layer.inferDims.d[1]; 303 | const uint gridSizeW = layer.inferDims.d[2]; 304 | const uint stride = DIVUP(networkInfo.width, gridSizeW); 305 | 306 | std::vector outObjs = 307 | decodeYoloTensor((const float*)(layer.buffer), masks[idx], anchors, gridSizeW, gridSizeH, stride, masks[idx].size(), 308 | num_classes, networkInfo.width, networkInfo.height, kCONF_THRESH); 309 | objects.insert(objects.end(), outObjs.begin(), outObjs.end()); 310 | } 311 | 312 | 313 | objectList.clear(); 314 | objectList = nmsAllClasses(beta_nms, objects, num_classes); 315 | 316 | return true; 317 | } 318 | 319 | static bool NvDsInferParseYoloV2( 320 | std::vector const& outputLayersInfo, 321 | NvDsInferNetworkInfo const& networkInfo, 322 | NvDsInferParseDetectionParams const& detectionParams, 323 | std::vector& objectList, 324 | std::vector &anchors, 325 | const uint &num_classes) 326 | { 327 | if (outputLayersInfo.empty()) { 328 | std::cerr << "Could not find output layer in bbox parsing" << std::endl;; 329 | return false; 330 | } 331 | const uint kNUM_BBOXES = anchors.size() / 2; 332 | 333 | const NvDsInferLayerInfo &layer = outputLayersInfo[0]; 334 | 335 | if (num_classes != detectionParams.numClassesConfigured) 336 | { 337 | std::cerr << "WARNING: Num classes mismatch. Configured: " 338 | << detectionParams.numClassesConfigured 339 | << ", detected by network: " << num_classes << std::endl; 340 | } 341 | 342 | assert(layer.inferDims.numDims == 3); 343 | const uint gridSizeH = layer.inferDims.d[1]; 344 | const uint gridSizeW = layer.inferDims.d[2]; 345 | const uint stride = DIVUP(networkInfo.width, gridSizeW); 346 | for (auto& anchor : anchors) { 347 | anchor *= stride; 348 | } 349 | std::vector objects = 350 | decodeYoloV2Tensor((const float*)(layer.buffer), anchors, gridSizeW, gridSizeH, stride, kNUM_BBOXES, 351 | num_classes, networkInfo.width, networkInfo.height); 352 | 353 | objectList = objects; 354 | 355 | return true; 356 | } 357 | 358 | extern "C" bool NvDsInferParseYolo( 359 | std::vector const& outputLayersInfo, 360 | NvDsInferNetworkInfo const& networkInfo, 361 | NvDsInferParseDetectionParams const& detectionParams, 362 | std::vector& objectList) 363 | { 364 | 365 | int num_classes = kNUM_CLASSES; 366 | float beta_nms = kBETA_NMS; 367 | std::vector anchors = kANCHORS; 368 | std::vector> mask = kMASK; 369 | 370 | if (mask.size() > 0) { 371 | return NvDsInferParseYolo (outputLayersInfo, networkInfo, detectionParams, objectList, anchors, mask, num_classes, beta_nms); 372 | } 373 | else { 374 | return NvDsInferParseYoloV2 (outputLayersInfo, networkInfo, detectionParams, objectList, anchors, num_classes); 375 | } 376 | } 377 | 378 | CHECK_CUSTOM_PARSE_FUNC_PROTOTYPE(NvDsInferParseYolo); -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/utils.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #include "utils.h" 27 | 28 | #include 29 | #include 30 | #include 31 | #include 32 | 33 | static void leftTrim(std::string& s) 34 | { 35 | s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](int ch) { return !isspace(ch); })); 36 | } 37 | 38 | static void rightTrim(std::string& s) 39 | { 40 | s.erase(std::find_if(s.rbegin(), s.rend(), [](int ch) { return !isspace(ch); }).base(), s.end()); 41 | } 42 | 43 | std::string trim(std::string s) 44 | { 45 | leftTrim(s); 46 | rightTrim(s); 47 | return s; 48 | } 49 | 50 | float clamp(const float val, const float minVal, const float maxVal) 51 | { 52 | assert(minVal <= maxVal); 53 | return std::min(maxVal, std::max(minVal, val)); 54 | } 55 | 56 | bool fileExists(const std::string fileName, bool verbose) 57 | { 58 | if (!std::experimental::filesystem::exists(std::experimental::filesystem::path(fileName))) 59 | { 60 | if (verbose) std::cout << "File does not exist: " << fileName << std::endl; 61 | return false; 62 | } 63 | return true; 64 | } 65 | 66 | std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType) 67 | { 68 | assert(fileExists(weightsFilePath)); 69 | std::cout << "\nLoading pre-trained weights" << std::endl; 70 | 71 | std::vector weights; 72 | 73 | if (weightsFilePath.find(".weights") != std::string::npos) { 74 | std::ifstream file(weightsFilePath, std::ios_base::binary); 75 | assert(file.good()); 76 | std::string line; 77 | 78 | if (networkType.find("yolov2") != std::string::npos && networkType.find("yolov2-tiny") == std::string::npos) 79 | { 80 | // Remove 4 int32 bytes of data from the stream belonging to the header 81 | file.ignore(4 * 4); 82 | } 83 | else 84 | { 85 | // Remove 5 int32 bytes of data from the stream belonging to the header 86 | file.ignore(4 * 5); 87 | } 88 | 89 | char floatWeight[4]; 90 | while (!file.eof()) 91 | { 92 | file.read(floatWeight, 4); 93 | assert(file.gcount() == 4); 94 | weights.push_back(*reinterpret_cast(floatWeight)); 95 | if (file.peek() == std::istream::traits_type::eof()) break; 96 | } 97 | } 98 | 99 | else if (weightsFilePath.find(".wts") != std::string::npos) { 100 | std::ifstream file(weightsFilePath); 101 | assert(file.good()); 102 | int32_t count; 103 | file >> count; 104 | assert(count > 0 && "Invalid .wts file."); 105 | 106 | uint32_t floatWeight; 107 | std::string name; 108 | uint32_t size; 109 | 110 | while (count--) { 111 | file >> name >> std::dec >> size; 112 | for (uint32_t x = 0, y = size; x < y; ++x) 113 | { 114 | file >> std::hex >> floatWeight; 115 | weights.push_back(*reinterpret_cast(&floatWeight)); 116 | }; 117 | } 118 | } 119 | 120 | else { 121 | std::cerr << "File " << weightsFilePath << " is not supported" << std::endl; 122 | std::abort(); 123 | } 124 | 125 | std::cout << "Loading weights of " << networkType << " complete" 126 | << std::endl; 127 | std::cout << "Total weights read: " << weights.size() << std::endl; 128 | return weights; 129 | } 130 | 131 | std::string dimsToString(const nvinfer1::Dims d) 132 | { 133 | std::stringstream s; 134 | assert(d.nbDims >= 1); 135 | for (int i = 0; i < d.nbDims - 1; ++i) 136 | { 137 | s << std::setw(4) << d.d[i] << " x"; 138 | } 139 | s << std::setw(4) << d.d[d.nbDims - 1]; 140 | 141 | return s.str(); 142 | } 143 | 144 | int getNumChannels(nvinfer1::ITensor* t) 145 | { 146 | nvinfer1::Dims d = t->getDimensions(); 147 | assert(d.nbDims == 3); 148 | 149 | return d.d[0]; 150 | } 151 | 152 | uint64_t get3DTensorVolume(nvinfer1::Dims inputDims) 153 | { 154 | assert(inputDims.nbDims == 3); 155 | return inputDims.d[0] * inputDims.d[1] * inputDims.d[2]; 156 | } 157 | 158 | void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, 159 | std::string layerOutput, std::string weightPtr) 160 | { 161 | std::cout << std::setw(6) << std::left << layerIndex << std::setw(24) << std::left << layerName; 162 | std::cout << std::setw(20) << std::left << layerInput << std::setw(20) << std::left 163 | << layerOutput; 164 | std::cout << std::setw(7) << std::left << weightPtr << std::endl; 165 | } -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/utils.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | 27 | #ifndef __UTILS_H__ 28 | #define __UTILS_H__ 29 | 30 | #include 31 | #include 32 | #include 33 | #include 34 | #include 35 | 36 | #include "NvInfer.h" 37 | 38 | #define UNUSED(expr) (void)(expr) 39 | #define DIVUP(n, d) ((n) + (d)-1) / (d) 40 | 41 | std::string trim(std::string s); 42 | float clamp(const float val, const float minVal, const float maxVal); 43 | bool fileExists(const std::string fileName, bool verbose = true); 44 | std::vector loadWeights(const std::string weightsFilePath, const std::string& networkType); 45 | std::string dimsToString(const nvinfer1::Dims d); 46 | void displayDimType(const nvinfer1::Dims d); 47 | int getNumChannels(nvinfer1::ITensor* t); 48 | uint64_t get3DTensorVolume(nvinfer1::Dims inputDims); 49 | 50 | void printLayerInfo(std::string layerIndex, std::string layerName, std::string layerInput, 51 | std::string layerOutput, std::string weightPtr); 52 | 53 | #endif 54 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/utils.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/utils.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yolo.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #include "yolo.h" 27 | #include "yoloPlugins.h" 28 | #include 29 | 30 | #ifdef OPENCV 31 | #include "calibrator.h" 32 | #endif 33 | 34 | void orderParams(std::vector> *maskVector) { 35 | std::vector> maskinput = *maskVector; 36 | std::vector maskPartial; 37 | for (uint i = 0; i < maskinput.size(); i++) { 38 | for (uint j = i + 1; j < maskinput.size(); j++) { 39 | if (maskinput[i][0] <= maskinput[j][0]) { 40 | maskPartial = maskinput[i]; 41 | maskinput[i] = maskinput[j]; 42 | maskinput[j] = maskPartial; 43 | } 44 | } 45 | } 46 | *maskVector = maskinput; 47 | } 48 | 49 | Yolo::Yolo(const NetworkInfo& networkInfo) 50 | : m_NetworkType(networkInfo.networkType), // YOLO type 51 | m_ConfigFilePath(networkInfo.configFilePath), // YOLO cfg 52 | m_WtsFilePath(networkInfo.wtsFilePath), // YOLO weights 53 | m_Int8CalibPath(networkInfo.int8CalibPath), // INT8 calibration path 54 | m_NetworkMode(networkInfo.networkMode), // FP32, INT8, FP16 55 | m_DeviceType(networkInfo.deviceType), // kDLA, kGPU 56 | m_InputBlobName(networkInfo.inputBlobName), // data 57 | m_InputH(0), 58 | m_InputW(0), 59 | m_InputC(0), 60 | m_InputSize(0) 61 | {} 62 | 63 | Yolo::~Yolo() 64 | { 65 | destroyNetworkUtils(); 66 | } 67 | 68 | nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder) 69 | { 70 | assert (builder); 71 | 72 | m_ConfigBlocks = parseConfigFile(m_ConfigFilePath); 73 | parseConfigBlocks(); 74 | orderParams(&m_OutputMasks); 75 | 76 | nvinfer1::INetworkDefinition *network = builder->createNetworkV2(0); 77 | if (parseModel(*network) != NVDSINFER_SUCCESS) { 78 | delete network; 79 | return nullptr; 80 | } 81 | 82 | std::cout << "Building the TensorRT Engine" << std::endl; 83 | 84 | if (m_LetterBox == 1) { 85 | std::cout << "\nNOTE: letter_box is set in cfg file, make sure to set maintain-aspect-ratio=1 in config_infer file to get better accuracy\n" << std::endl; 86 | } 87 | 88 | nvinfer1::IBuilderConfig *config = builder->createBuilderConfig(); 89 | 90 | if (m_NetworkMode == "INT8" && !fileExists(m_Int8CalibPath)) { 91 | assert(builder->platformHasFastInt8()); 92 | #ifdef OPENCV 93 | std::string calib_image_list; 94 | int calib_batch_size; 95 | if (getenv("INT8_CALIB_IMG_PATH")) { 96 | calib_image_list = getenv("INT8_CALIB_IMG_PATH"); 97 | } 98 | else { 99 | std::cerr << "INT8_CALIB_IMG_PATH not set" << std::endl; 100 | std::abort(); 101 | } 102 | if (getenv("INT8_CALIB_BATCH_SIZE")) { 103 | calib_batch_size = std::stoi(getenv("INT8_CALIB_BATCH_SIZE")); 104 | } 105 | else { 106 | std::cerr << "INT8_CALIB_BATCH_SIZE not set" << std::endl; 107 | std::abort(); 108 | } 109 | nvinfer1::Int8EntropyCalibrator2 *calibrator = new nvinfer1::Int8EntropyCalibrator2(calib_batch_size, m_InputC, m_InputH, m_InputW, m_LetterBox, calib_image_list, m_Int8CalibPath); 110 | config->setFlag(nvinfer1::BuilderFlag::kINT8); 111 | config->setInt8Calibrator(calibrator); 112 | #else 113 | std::cerr << "OpenCV is required to run INT8 calibrator" << std::endl; 114 | std::abort(); 115 | #endif 116 | } 117 | 118 | nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config); 119 | if (engine) { 120 | std::cout << "Building complete\n" << std::endl; 121 | } else { 122 | std::cerr << "Building engine failed\n" << std::endl; 123 | } 124 | 125 | delete network; 126 | delete config; 127 | return engine; 128 | } 129 | 130 | NvDsInferStatus Yolo::parseModel(nvinfer1::INetworkDefinition& network) { 131 | destroyNetworkUtils(); 132 | 133 | std::vector weights = loadWeights(m_WtsFilePath, m_NetworkType); 134 | std::cout << "Building YOLO network\n" << std::endl; 135 | NvDsInferStatus status = buildYoloNetwork(weights, network); 136 | 137 | if (status == NVDSINFER_SUCCESS) { 138 | std::cout << "Building YOLO network complete" << std::endl; 139 | } else { 140 | std::cerr << "Building YOLO network failed" << std::endl; 141 | } 142 | 143 | return status; 144 | } 145 | 146 | NvDsInferStatus Yolo::buildYoloNetwork( 147 | std::vector& weights, nvinfer1::INetworkDefinition& network) { 148 | int weightPtr = 0; 149 | int channels = m_InputC; 150 | 151 | std::string weightsType; 152 | 153 | if (m_WtsFilePath.find(".weights") != std::string::npos) { 154 | weightsType = "weights"; 155 | } 156 | else { 157 | weightsType = "wts"; 158 | } 159 | 160 | nvinfer1::ITensor* data = 161 | network.addInput(m_InputBlobName.c_str(), nvinfer1::DataType::kFLOAT, 162 | nvinfer1::Dims3{static_cast(m_InputC), 163 | static_cast(m_InputH), static_cast(m_InputW)}); 164 | assert(data != nullptr && data->getDimensions().nbDims > 0); 165 | 166 | nvinfer1::ITensor* previous = data; 167 | std::vector tensorOutputs; 168 | uint outputTensorCount = 0; 169 | 170 | for (uint i = 0; i < m_ConfigBlocks.size(); ++i) { 171 | assert(getNumChannels(previous) == channels); 172 | std::string layerIndex = "(" + std::to_string(tensorOutputs.size()) + ")"; 173 | 174 | if (m_ConfigBlocks.at(i).at("type") == "net") { 175 | printLayerInfo("", "layer", " input", " output", "weightPtr"); 176 | } 177 | 178 | else if (m_ConfigBlocks.at(i).at("type") == "convolutional") { 179 | std::string inputVol = dimsToString(previous->getDimensions()); 180 | nvinfer1::ILayer* out = convolutionalLayer(i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, weightsType, channels, previous, &network); 181 | previous = out->getOutput(0); 182 | assert(previous != nullptr); 183 | channels = getNumChannels(previous); 184 | std::string outputVol = dimsToString(previous->getDimensions()); 185 | tensorOutputs.push_back(previous); 186 | std::string layerType = "conv_" + m_ConfigBlocks.at(i).at("activation"); 187 | printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr)); 188 | } 189 | 190 | else if (m_ConfigBlocks.at(i).at("type") == "implicit_add" || m_ConfigBlocks.at(i).at("type") == "implicit_mul") { 191 | std::string type; 192 | if (m_ConfigBlocks.at(i).at("type") == "implicit_add") { 193 | type = "add"; 194 | } 195 | else if (m_ConfigBlocks.at(i).at("type") == "implicit_mul") { 196 | type = "mul"; 197 | } 198 | assert(m_ConfigBlocks.at(i).find("filters") != m_ConfigBlocks.at(i).end()); 199 | int filters = std::stoi(m_ConfigBlocks.at(i).at("filters")); 200 | nvinfer1::ILayer* out = implicitLayer(filters, weights, m_TrtWeights, weightPtr, &network); 201 | previous = out->getOutput(0); 202 | assert(previous != nullptr); 203 | channels = getNumChannels(previous); 204 | std::string outputVol = dimsToString(previous->getDimensions()); 205 | tensorOutputs.push_back(previous); 206 | std::string layerType = "implicit_" + type; 207 | printLayerInfo(layerIndex, layerType, " -", outputVol, std::to_string(weightPtr)); 208 | } 209 | 210 | else if (m_ConfigBlocks.at(i).at("type") == "shift_channels" || m_ConfigBlocks.at(i).at("type") == "control_channels") { 211 | std::string type; 212 | if (m_ConfigBlocks.at(i).at("type") == "shift_channels") { 213 | type = "shift"; 214 | } 215 | else if (m_ConfigBlocks.at(i).at("type") == "control_channels") { 216 | type = "control"; 217 | } 218 | assert(m_ConfigBlocks.at(i).find("from") != m_ConfigBlocks.at(i).end()); 219 | int from = stoi(m_ConfigBlocks.at(i).at("from")); 220 | if (from > 0) { 221 | from = from - i + 1; 222 | } 223 | assert((i - 2 >= 0) && (i - 2 < tensorOutputs.size())); 224 | assert((i + from - 1 >= 0) && (i + from - 1 < tensorOutputs.size())); 225 | assert(i + from - 1 < i - 2); 226 | nvinfer1::ILayer* out = channelsLayer(type, previous, tensorOutputs[i + from - 1], &network); 227 | previous = out->getOutput(0); 228 | assert(previous != nullptr); 229 | std::string outputVol = dimsToString(previous->getDimensions()); 230 | tensorOutputs.push_back(previous); 231 | std::string layerType = type + "_channels" + ": " + std::to_string(i + from - 1); 232 | printLayerInfo(layerIndex, layerType, " -", outputVol, " -"); 233 | } 234 | 235 | else if (m_ConfigBlocks.at(i).at("type") == "dropout") { // Skip dropout layer 236 | assert(m_ConfigBlocks.at(i).find("probability") != m_ConfigBlocks.at(i).end()); 237 | //float probability = std::stof(m_ConfigBlocks.at(i).at("probability")); 238 | //nvinfer1::ILayer* out = dropoutLayer(probability, previous, &network); 239 | //previous = out->getOutput(0); 240 | assert(previous != nullptr); 241 | tensorOutputs.push_back(previous); 242 | printLayerInfo(layerIndex, "dropout", " -", " -", " -"); 243 | } 244 | 245 | else if (m_ConfigBlocks.at(i).at("type") == "shortcut") { 246 | assert(m_ConfigBlocks.at(i).find("activation") != m_ConfigBlocks.at(i).end()); 247 | assert(m_ConfigBlocks.at(i).find("from") != m_ConfigBlocks.at(i).end()); 248 | std::string activation = m_ConfigBlocks.at(i).at("activation"); 249 | int from = stoi(m_ConfigBlocks.at(i).at("from")); 250 | if (from > 0) { 251 | from = from - i + 1; 252 | } 253 | assert((i - 2 >= 0) && (i - 2 < tensorOutputs.size())); 254 | assert((i + from - 1 >= 0) && (i + from - 1 < tensorOutputs.size())); 255 | assert(i + from - 1 < i - 2); 256 | std::string inputVol = dimsToString(previous->getDimensions()); 257 | std::string shortcutVol = dimsToString(tensorOutputs[i + from - 1]->getDimensions()); 258 | nvinfer1::ILayer* out = shortcutLayer(i, activation, inputVol, shortcutVol, previous, tensorOutputs[i + from - 1], &network); 259 | previous = out->getOutput(0); 260 | assert(previous != nullptr); 261 | std::string outputVol = dimsToString(previous->getDimensions()); 262 | tensorOutputs.push_back(previous); 263 | std::string layerType = "shortcut_" + m_ConfigBlocks.at(i).at("activation") + ": " + std::to_string(i + from - 1); 264 | printLayerInfo(layerIndex, layerType, " -", outputVol, " -"); 265 | if (inputVol != shortcutVol) { 266 | std::cout << inputVol << " +" << shortcutVol << std::endl; 267 | } 268 | } 269 | 270 | else if (m_ConfigBlocks.at(i).at("type") == "route") { 271 | assert(m_ConfigBlocks.at(i).find("layers") != m_ConfigBlocks.at(i).end()); 272 | nvinfer1::ILayer* out = routeLayer(i, m_ConfigBlocks.at(i), tensorOutputs, &network); 273 | previous = out->getOutput(0); 274 | assert(previous != nullptr); 275 | channels = getNumChannels(previous); 276 | std::string outputVol = dimsToString(previous->getDimensions()); 277 | tensorOutputs.push_back(previous); 278 | printLayerInfo(layerIndex, "route", " -", outputVol, std::to_string(weightPtr)); 279 | } 280 | 281 | else if (m_ConfigBlocks.at(i).at("type") == "upsample") { 282 | std::string inputVol = dimsToString(previous->getDimensions()); 283 | nvinfer1::ILayer* out = upsampleLayer(i - 1, m_ConfigBlocks[i], previous, &network); 284 | previous = out->getOutput(0); 285 | assert(previous != nullptr); 286 | std::string outputVol = dimsToString(previous->getDimensions()); 287 | tensorOutputs.push_back(previous); 288 | printLayerInfo(layerIndex, "upsample", inputVol, outputVol, " -"); 289 | } 290 | 291 | else if (m_ConfigBlocks.at(i).at("type") == "maxpool") { 292 | std::string inputVol = dimsToString(previous->getDimensions()); 293 | nvinfer1::ILayer* out = maxpoolLayer(i, m_ConfigBlocks.at(i), previous, &network); 294 | previous = out->getOutput(0); 295 | assert(previous != nullptr); 296 | std::string outputVol = dimsToString(previous->getDimensions()); 297 | tensorOutputs.push_back(previous); 298 | printLayerInfo(layerIndex, "maxpool", inputVol, outputVol, std::to_string(weightPtr)); 299 | } 300 | 301 | else if (m_ConfigBlocks.at(i).at("type") == "yolo") { 302 | uint model_type; 303 | if (m_NetworkType.find("yolor") != std::string::npos) { 304 | model_type = 2; 305 | } 306 | else { 307 | model_type = 1; 308 | } 309 | nvinfer1::Dims prevTensorDims = previous->getDimensions(); 310 | TensorInfo& curYoloTensor = m_OutputTensors.at(outputTensorCount); 311 | curYoloTensor.gridSizeY = prevTensorDims.d[1]; 312 | curYoloTensor.gridSizeX = prevTensorDims.d[2]; 313 | curYoloTensor.stride = m_InputH / curYoloTensor.gridSizeY; 314 | m_OutputTensors.at(outputTensorCount).volume = curYoloTensor.gridSizeY 315 | * curYoloTensor.gridSizeX 316 | * (curYoloTensor.numBBoxes * (5 + curYoloTensor.numClasses)); 317 | std::string layerName = "yolo_" + std::to_string(i); 318 | curYoloTensor.blobName = layerName; 319 | int new_coords = 0; 320 | float scale_x_y = 1; 321 | float beta_nms = 0.45; 322 | if (m_ConfigBlocks.at(i).find("new_coords") != m_ConfigBlocks.at(i).end()) { 323 | new_coords = std::stoi(m_ConfigBlocks.at(i).at("new_coords")); 324 | } 325 | if (m_ConfigBlocks.at(i).find("scale_x_y") != m_ConfigBlocks.at(i).end()) { 326 | scale_x_y = std::stof(m_ConfigBlocks.at(i).at("scale_x_y")); 327 | } 328 | if (m_ConfigBlocks.at(i).find("beta_nms") != m_ConfigBlocks.at(i).end()) { 329 | beta_nms = std::stof(m_ConfigBlocks.at(i).at("beta_nms")); 330 | } 331 | nvinfer1::IPluginV2* yoloPlugin 332 | = new YoloLayer(curYoloTensor.numBBoxes, 333 | curYoloTensor.numClasses, 334 | curYoloTensor.gridSizeX, 335 | curYoloTensor.gridSizeY, 336 | model_type, new_coords, scale_x_y, beta_nms, 337 | curYoloTensor.anchors, 338 | m_OutputMasks); 339 | assert(yoloPlugin != nullptr); 340 | nvinfer1::IPluginV2Layer* yolo = 341 | network.addPluginV2(&previous, 1, *yoloPlugin); 342 | assert(yolo != nullptr); 343 | yolo->setName(layerName.c_str()); 344 | std::string inputVol = dimsToString(previous->getDimensions()); 345 | previous = yolo->getOutput(0); 346 | assert(previous != nullptr); 347 | previous->setName(layerName.c_str()); 348 | std::string outputVol = dimsToString(previous->getDimensions()); 349 | network.markOutput(*previous); 350 | channels = getNumChannels(previous); 351 | tensorOutputs.push_back(yolo->getOutput(0)); 352 | printLayerInfo(layerIndex, "yolo", inputVol, outputVol, std::to_string(weightPtr)); 353 | ++outputTensorCount; 354 | } 355 | 356 | //YOLOv2 support 357 | else if (m_ConfigBlocks.at(i).at("type") == "region") { 358 | nvinfer1::Dims prevTensorDims = previous->getDimensions(); 359 | TensorInfo& curRegionTensor = m_OutputTensors.at(outputTensorCount); 360 | curRegionTensor.gridSizeY = prevTensorDims.d[1]; 361 | curRegionTensor.gridSizeX = prevTensorDims.d[2]; 362 | curRegionTensor.stride = m_InputH / curRegionTensor.gridSizeY; 363 | m_OutputTensors.at(outputTensorCount).volume = curRegionTensor.gridSizeY 364 | * curRegionTensor.gridSizeX 365 | * (curRegionTensor.numBBoxes * (5 + curRegionTensor.numClasses)); 366 | std::string layerName = "region_" + std::to_string(i); 367 | curRegionTensor.blobName = layerName; 368 | std::vector> mask; 369 | nvinfer1::IPluginV2* regionPlugin 370 | = new YoloLayer(curRegionTensor.numBBoxes, 371 | curRegionTensor.numClasses, 372 | curRegionTensor.gridSizeX, 373 | curRegionTensor.gridSizeY, 374 | 0, 0, 1.0, 0, 375 | curRegionTensor.anchors, 376 | mask); 377 | assert(regionPlugin != nullptr); 378 | nvinfer1::IPluginV2Layer* region = 379 | network.addPluginV2(&previous, 1, *regionPlugin); 380 | assert(region != nullptr); 381 | region->setName(layerName.c_str()); 382 | std::string inputVol = dimsToString(previous->getDimensions()); 383 | previous = region->getOutput(0); 384 | assert(previous != nullptr); 385 | previous->setName(layerName.c_str()); 386 | std::string outputVol = dimsToString(previous->getDimensions()); 387 | network.markOutput(*previous); 388 | channels = getNumChannels(previous); 389 | tensorOutputs.push_back(region->getOutput(0)); 390 | printLayerInfo(layerIndex, "region", inputVol, outputVol, std::to_string(weightPtr)); 391 | ++outputTensorCount; 392 | } 393 | else if (m_ConfigBlocks.at(i).at("type") == "reorg") { 394 | std::string inputVol = dimsToString(previous->getDimensions()); 395 | nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2); 396 | assert(reorgPlugin != nullptr); 397 | nvinfer1::IPluginV2Layer* reorg = 398 | network.addPluginV2(&previous, 1, *reorgPlugin); 399 | assert(reorg != nullptr); 400 | std::string layerName = "reorg_" + std::to_string(i); 401 | reorg->setName(layerName.c_str()); 402 | previous = reorg->getOutput(0); 403 | assert(previous != nullptr); 404 | std::string outputVol = dimsToString(previous->getDimensions()); 405 | channels = getNumChannels(previous); 406 | tensorOutputs.push_back(reorg->getOutput(0)); 407 | printLayerInfo(layerIndex, "reorg", inputVol, outputVol, std::to_string(weightPtr)); 408 | } 409 | 410 | else 411 | { 412 | std::cout << "Unsupported layer type --> \"" 413 | << m_ConfigBlocks.at(i).at("type") << "\"" << std::endl; 414 | assert(0); 415 | } 416 | } 417 | 418 | if ((int)weights.size() != weightPtr) 419 | { 420 | std::cout << "Number of unused weights left: " << weights.size() - weightPtr << std::endl; 421 | assert(0); 422 | } 423 | 424 | std::cout << "Output YOLO blob names: " << std::endl; 425 | for (auto& tensor : m_OutputTensors) { 426 | std::cout << tensor.blobName << std::endl; 427 | } 428 | 429 | int nbLayers = network.getNbLayers(); 430 | std::cout << "Total number of YOLO layers: " << nbLayers << std::endl; 431 | 432 | return NVDSINFER_SUCCESS; 433 | } 434 | 435 | std::vector> 436 | Yolo::parseConfigFile (const std::string cfgFilePath) 437 | { 438 | assert(fileExists(cfgFilePath)); 439 | std::ifstream file(cfgFilePath); 440 | assert(file.good()); 441 | std::string line; 442 | std::vector> blocks; 443 | std::map block; 444 | 445 | while (getline(file, line)) 446 | { 447 | if (line.size() == 0) continue; 448 | if (line.front() == '#') continue; 449 | line = trim(line); 450 | if (line.front() == '[') 451 | { 452 | if (block.size() > 0) 453 | { 454 | blocks.push_back(block); 455 | block.clear(); 456 | } 457 | std::string key = "type"; 458 | std::string value = trim(line.substr(1, line.size() - 2)); 459 | block.insert(std::pair(key, value)); 460 | } 461 | else 462 | { 463 | int cpos = line.find('='); 464 | std::string key = trim(line.substr(0, cpos)); 465 | std::string value = trim(line.substr(cpos + 1)); 466 | block.insert(std::pair(key, value)); 467 | } 468 | } 469 | blocks.push_back(block); 470 | return blocks; 471 | } 472 | 473 | void Yolo::parseConfigBlocks() 474 | { 475 | for (auto block : m_ConfigBlocks) { 476 | if (block.at("type") == "net") 477 | { 478 | assert((block.find("height") != block.end()) 479 | && "Missing 'height' param in network cfg"); 480 | assert((block.find("width") != block.end()) && "Missing 'width' param in network cfg"); 481 | assert((block.find("channels") != block.end()) 482 | && "Missing 'channels' param in network cfg"); 483 | 484 | m_InputH = std::stoul(block.at("height")); 485 | m_InputW = std::stoul(block.at("width")); 486 | m_InputC = std::stoul(block.at("channels")); 487 | m_InputSize = m_InputC * m_InputH * m_InputW; 488 | 489 | if (block.find("letter_box") != block.end()) { 490 | m_LetterBox = std::stoul(block.at("letter_box")); 491 | } 492 | else { 493 | m_LetterBox = 0; 494 | } 495 | } 496 | else if ((block.at("type") == "region") || (block.at("type") == "yolo") || (block.at("type") == "detect")) 497 | { 498 | assert((block.find("num") != block.end()) 499 | && std::string("Missing 'num' param in " + block.at("type") + " layer").c_str()); 500 | assert((block.find("classes") != block.end()) 501 | && std::string("Missing 'classes' param in " + block.at("type") + " layer") 502 | .c_str()); 503 | assert((block.find("anchors") != block.end()) 504 | && std::string("Missing 'anchors' param in " + block.at("type") + " layer") 505 | .c_str()); 506 | 507 | TensorInfo outputTensor; 508 | std::string anchorString = block.at("anchors"); 509 | while (!anchorString.empty()) 510 | { 511 | int npos = anchorString.find_first_of(','); 512 | if (npos != -1) 513 | { 514 | float anchor = std::stof(trim(anchorString.substr(0, npos))); 515 | outputTensor.anchors.push_back(anchor); 516 | anchorString.erase(0, npos + 1); 517 | } 518 | else 519 | { 520 | float anchor = std::stof(trim(anchorString)); 521 | outputTensor.anchors.push_back(anchor); 522 | break; 523 | } 524 | } 525 | 526 | if (block.find("mask") != block.end()) { 527 | std::string maskString = block.at("mask"); 528 | std::vector pMASKS; 529 | while (!maskString.empty()) 530 | { 531 | int npos = maskString.find_first_of(','); 532 | if (npos != -1) 533 | { 534 | int mask = std::stoul(trim(maskString.substr(0, npos))); 535 | pMASKS.push_back(mask); 536 | outputTensor.masks.push_back(mask); 537 | maskString.erase(0, npos + 1); 538 | } 539 | else 540 | { 541 | int mask = std::stoul(trim(maskString)); 542 | pMASKS.push_back(mask); 543 | outputTensor.masks.push_back(mask); 544 | break; 545 | } 546 | } 547 | m_OutputMasks.push_back(pMASKS); 548 | } 549 | 550 | outputTensor.numBBoxes = outputTensor.masks.size() > 0 551 | ? outputTensor.masks.size() 552 | : std::stoul(trim(block.at("num"))); 553 | outputTensor.numClasses = std::stoul(block.at("classes")); 554 | m_OutputTensors.push_back(outputTensor); 555 | } 556 | } 557 | } 558 | 559 | void Yolo::destroyNetworkUtils() { 560 | for (uint i = 0; i < m_TrtWeights.size(); ++i) { 561 | if (m_TrtWeights[i].count > 0) 562 | free(const_cast(m_TrtWeights[i].values)); 563 | } 564 | m_TrtWeights.clear(); 565 | } 566 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yolo.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #ifndef _YOLO_H_ 27 | #define _YOLO_H_ 28 | 29 | #include "layers/convolutional_layer.h" 30 | #include "layers/implicit_layer.h" 31 | #include "layers/channels_layer.h" 32 | #include "layers/dropout_layer.h" 33 | #include "layers/shortcut_layer.h" 34 | #include "layers/route_layer.h" 35 | #include "layers/upsample_layer.h" 36 | #include "layers/maxpool_layer.h" 37 | 38 | #include "nvdsinfer_custom_impl.h" 39 | 40 | struct NetworkInfo 41 | { 42 | std::string networkType; 43 | std::string configFilePath; 44 | std::string wtsFilePath; 45 | std::string int8CalibPath; 46 | std::string networkMode; 47 | std::string deviceType; 48 | std::string inputBlobName; 49 | }; 50 | 51 | struct TensorInfo 52 | { 53 | std::string blobName; 54 | uint stride{0}; 55 | uint gridSizeY{0}; 56 | uint gridSizeX{0}; 57 | uint numClasses{0}; 58 | uint numBBoxes{0}; 59 | uint64_t volume{0}; 60 | std::vector masks; 61 | std::vector anchors; 62 | int bindingIndex{-1}; 63 | float* hostBuffer{nullptr}; 64 | }; 65 | 66 | class Yolo : public IModelParser { 67 | public: 68 | Yolo(const NetworkInfo& networkInfo); 69 | ~Yolo() override; 70 | bool hasFullDimsSupported() const override { return false; } 71 | const char* getModelName() const override { 72 | return m_ConfigFilePath.empty() ? m_NetworkType.c_str() 73 | : m_ConfigFilePath.c_str(); 74 | } 75 | NvDsInferStatus parseModel(nvinfer1::INetworkDefinition& network) override; 76 | 77 | nvinfer1::ICudaEngine *createEngine (nvinfer1::IBuilder* builder); 78 | 79 | protected: 80 | const std::string m_NetworkType; 81 | const std::string m_ConfigFilePath; 82 | const std::string m_WtsFilePath; 83 | const std::string m_Int8CalibPath; 84 | const std::string m_NetworkMode; 85 | const std::string m_DeviceType; 86 | const std::string m_InputBlobName; 87 | std::vector m_OutputTensors; 88 | std::vector> m_OutputMasks; 89 | std::vector> m_ConfigBlocks; 90 | uint m_InputH; 91 | uint m_InputW; 92 | uint m_InputC; 93 | uint64_t m_InputSize; 94 | uint m_LetterBox; 95 | 96 | std::vector m_TrtWeights; 97 | 98 | private: 99 | NvDsInferStatus buildYoloNetwork( 100 | std::vector& weights, nvinfer1::INetworkDefinition& network); 101 | std::vector> parseConfigFile( 102 | const std::string cfgFilePath); 103 | void parseConfigBlocks(); 104 | void destroyNetworkUtils(); 105 | }; 106 | 107 | #endif // _YOLO_H_ -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yolo.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/yolo.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018-2019 NVIDIA Corporation. All rights reserved. 3 | * 4 | * NVIDIA Corporation and its licensors retain all intellectual property 5 | * and proprietary rights in and to this software, related documentation 6 | * and any modifications thereto. Any use, reproduction, disclosure or 7 | * distribution of this software and related documentation without an express 8 | * license agreement from NVIDIA Corporation is strictly prohibited. 9 | * 10 | * Edited by Marcos Luciano 11 | * https://www.github.com/marcoslucianops 12 | * 13 | */ 14 | 15 | #include 16 | #include 17 | #include 18 | #include 19 | #include 20 | 21 | inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } 22 | 23 | __global__ void gpuYoloLayer(const float* input, float* output, const uint gridSizeX, const uint gridSizeY, const uint numOutputClasses, 24 | const uint numBBoxes, const float scale_x_y) 25 | { 26 | uint x_id = blockIdx.x * blockDim.x + threadIdx.x; 27 | uint y_id = blockIdx.y * blockDim.y + threadIdx.y; 28 | uint z_id = blockIdx.z * blockDim.z + threadIdx.z; 29 | 30 | if ((x_id >= gridSizeX) || (y_id >= gridSizeY) || (z_id >= numBBoxes)) 31 | { 32 | return; 33 | } 34 | 35 | const int numGridCells = gridSizeX * gridSizeY; 36 | const int bbindex = y_id * gridSizeX + x_id; 37 | 38 | const float alpha = scale_x_y; 39 | const float beta = -0.5 * (scale_x_y - 1); 40 | 41 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] 42 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]) * alpha + beta; 43 | 44 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] 45 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]) * alpha + beta; 46 | 47 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] 48 | = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]); 49 | 50 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] 51 | = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]); 52 | 53 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] 54 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); 55 | 56 | for (uint i = 0; i < numOutputClasses; ++i) 57 | { 58 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] 59 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]); 60 | } 61 | } 62 | 63 | cudaError_t cudaYoloLayer(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 64 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, 65 | const float modelScale); 66 | 67 | cudaError_t cudaYoloLayer(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 68 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, 69 | const float modelScale) 70 | { 71 | dim3 threads_per_block(16, 16, 4); 72 | dim3 number_of_blocks((gridSizeX / threads_per_block.x) + 1, 73 | (gridSizeY / threads_per_block.y) + 1, 74 | (numBBoxes / threads_per_block.z) + 1); 75 | for (unsigned int batch = 0; batch < batchSize; ++batch) 76 | { 77 | gpuYoloLayer<<>>( 78 | reinterpret_cast(input) + (batch * outputSize), 79 | reinterpret_cast(output) + (batch * outputSize), gridSizeX, gridSizeY, numOutputClasses, 80 | numBBoxes, modelScale); 81 | } 82 | return cudaGetLastError(); 83 | } 84 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_nc.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } 13 | 14 | __global__ void gpuYoloLayer_nc(const float* input, float* output, const uint gridSizeX, const uint gridSizeY, const uint numOutputClasses, 15 | const uint numBBoxes, const float scale_x_y) 16 | { 17 | uint x_id = blockIdx.x * blockDim.x + threadIdx.x; 18 | uint y_id = blockIdx.y * blockDim.y + threadIdx.y; 19 | uint z_id = blockIdx.z * blockDim.z + threadIdx.z; 20 | 21 | if ((x_id >= gridSizeX) || (y_id >= gridSizeY) || (z_id >= numBBoxes)) 22 | { 23 | return; 24 | } 25 | 26 | const int numGridCells = gridSizeX * gridSizeY; 27 | const int bbindex = y_id * gridSizeX + x_id; 28 | 29 | const float alpha = scale_x_y; 30 | const float beta = -0.5 * (scale_x_y - 1); 31 | 32 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] 33 | = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] * alpha + beta; 34 | 35 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] 36 | = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] * alpha + beta; 37 | 38 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] 39 | = pow(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] * 2, 2); 40 | 41 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] 42 | = pow(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] * 2, 2); 43 | 44 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] 45 | = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]; 46 | 47 | for (uint i = 0; i < numOutputClasses; ++i) 48 | { 49 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] 50 | = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]; 51 | } 52 | } 53 | 54 | cudaError_t cudaYoloLayer_nc(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 55 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, 56 | const float modelScale); 57 | 58 | cudaError_t cudaYoloLayer_nc(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 59 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, 60 | const float modelScale) 61 | { 62 | dim3 threads_per_block(16, 16, 4); 63 | dim3 number_of_blocks((gridSizeX / threads_per_block.x) + 1, 64 | (gridSizeY / threads_per_block.y) + 1, 65 | (numBBoxes / threads_per_block.z) + 1); 66 | for (unsigned int batch = 0; batch < batchSize; ++batch) 67 | { 68 | gpuYoloLayer_nc<<>>( 69 | reinterpret_cast(input) + (batch * outputSize), 70 | reinterpret_cast(output) + (batch * outputSize), gridSizeX, gridSizeY, numOutputClasses, 71 | numBBoxes, modelScale); 72 | } 73 | return cudaGetLastError(); 74 | } 75 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_nc.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_nc.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_r.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } 13 | 14 | __global__ void gpuYoloLayer_r(const float* input, float* output, const uint gridSizeX, const uint gridSizeY, const uint numOutputClasses, 15 | const uint numBBoxes, const float scale_x_y) 16 | { 17 | uint x_id = blockIdx.x * blockDim.x + threadIdx.x; 18 | uint y_id = blockIdx.y * blockDim.y + threadIdx.y; 19 | uint z_id = blockIdx.z * blockDim.z + threadIdx.z; 20 | 21 | if ((x_id >= gridSizeX) || (y_id >= gridSizeY) || (z_id >= numBBoxes)) 22 | { 23 | return; 24 | } 25 | 26 | const int numGridCells = gridSizeX * gridSizeY; 27 | const int bbindex = y_id * gridSizeX + x_id; 28 | 29 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] 30 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]) * 2.0 - 0.5; 31 | 32 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] 33 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]) * 2.0 - 0.5; 34 | 35 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] 36 | = pow(sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]) * 2, 2); 37 | 38 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] 39 | = pow(sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]) * 2, 2); 40 | 41 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] 42 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); 43 | 44 | for (uint i = 0; i < numOutputClasses; ++i) 45 | { 46 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] 47 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]); 48 | } 49 | } 50 | 51 | cudaError_t cudaYoloLayer_r(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 52 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, 53 | const float modelScale); 54 | 55 | cudaError_t cudaYoloLayer_r(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 56 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, 57 | const float modelScale) 58 | { 59 | dim3 threads_per_block(16, 16, 4); 60 | dim3 number_of_blocks((gridSizeX / threads_per_block.x) + 1, 61 | (gridSizeY / threads_per_block.y) + 1, 62 | (numBBoxes / threads_per_block.z) + 1); 63 | for (unsigned int batch = 0; batch < batchSize; ++batch) 64 | { 65 | gpuYoloLayer_r<<>>( 66 | reinterpret_cast(input) + (batch * outputSize), 67 | reinterpret_cast(output) + (batch * outputSize), gridSizeX, gridSizeY, numOutputClasses, 68 | numBBoxes, modelScale); 69 | } 70 | return cudaGetLastError(); 71 | } 72 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_r.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_r.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_v2.cu: -------------------------------------------------------------------------------- 1 | /* 2 | * Created by Marcos Luciano 3 | * https://www.github.com/marcoslucianops 4 | */ 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | 12 | inline __device__ float sigmoidGPU(const float& x) { return 1.0f / (1.0f + __expf(-x)); } 13 | 14 | __global__ void gpuRegionLayer(const float* input, float* output, const uint gridSizeX, const uint gridSizeY, const uint numOutputClasses, 15 | const uint numBBoxes) 16 | { 17 | uint x_id = blockIdx.x * blockDim.x + threadIdx.x; 18 | uint y_id = blockIdx.y * blockDim.y + threadIdx.y; 19 | uint z_id = blockIdx.z * blockDim.z + threadIdx.z; 20 | 21 | if ((x_id >= gridSizeX) || (y_id >= gridSizeY) || (z_id >= numBBoxes)) 22 | { 23 | return; 24 | } 25 | 26 | const int numGridCells = gridSizeX * gridSizeY; 27 | const int bbindex = y_id * gridSizeX + x_id; 28 | 29 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)] 30 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 0)]); 31 | 32 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)] 33 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 1)]); 34 | 35 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)] 36 | = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 2)]); 37 | 38 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)] 39 | = __expf(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 3)]); 40 | 41 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)] 42 | = sigmoidGPU(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + 4)]); 43 | 44 | float temp = 1.0; 45 | int i; 46 | float sum = 0; 47 | float largest = -INFINITY; 48 | for(i = 0; i < numOutputClasses; ++i){ 49 | int val = input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))]; 50 | largest = (val>largest) ? val : largest; 51 | } 52 | for(i = 0; i < numOutputClasses; ++i){ 53 | float e = exp(input[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] / temp - largest / temp); 54 | sum += e; 55 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] = e; 56 | } 57 | for(i = 0; i < numOutputClasses; ++i){ 58 | output[bbindex + numGridCells * (z_id * (5 + numOutputClasses) + (5 + i))] /= sum; 59 | } 60 | } 61 | 62 | cudaError_t cudaYoloLayer_v2(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 63 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream); 64 | 65 | cudaError_t cudaYoloLayer_v2(const void* input, void* output, const uint& batchSize, const uint& gridSizeX, const uint& gridSizeY, 66 | const uint& numOutputClasses, const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream) 67 | { 68 | dim3 threads_per_block(16, 16, 4); 69 | dim3 number_of_blocks((gridSizeX / threads_per_block.x) + 1, 70 | (gridSizeY / threads_per_block.y) + 1, 71 | (numBBoxes / threads_per_block.z) + 1); 72 | for (unsigned int batch = 0; batch < batchSize; ++batch) 73 | { 74 | gpuRegionLayer<<>>( 75 | reinterpret_cast(input) + (batch * outputSize), 76 | reinterpret_cast(output) + (batch * outputSize), gridSizeX, gridSizeY, numOutputClasses, 77 | numBBoxes); 78 | } 79 | return cudaGetLastError(); 80 | } 81 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_v2.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloForward_v2.o -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #include "yoloPlugins.h" 27 | #include "NvInferPlugin.h" 28 | #include 29 | #include 30 | #include 31 | 32 | int kNUM_CLASSES; 33 | float kBETA_NMS; 34 | std::vector kANCHORS; 35 | std::vector> kMASK; 36 | 37 | namespace { 38 | template 39 | void write(char*& buffer, const T& val) 40 | { 41 | *reinterpret_cast(buffer) = val; 42 | buffer += sizeof(T); 43 | } 44 | 45 | template 46 | void read(const char*& buffer, T& val) 47 | { 48 | val = *reinterpret_cast(buffer); 49 | buffer += sizeof(T); 50 | } 51 | } 52 | 53 | cudaError_t cudaYoloLayer ( 54 | const void* input, void* output, const uint& batchSize, 55 | const uint& gridSizeX, const uint& gridSizeY, const uint& numOutputClasses, 56 | const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, const float modelScale); 57 | 58 | cudaError_t cudaYoloLayer_v2 ( 59 | const void* input, void* output, const uint& batchSize, 60 | const uint& gridSizeX, const uint& gridSizeY, const uint& numOutputClasses, 61 | const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream); 62 | 63 | cudaError_t cudaYoloLayer_nc ( 64 | const void* input, void* output, const uint& batchSize, 65 | const uint& gridSizeX, const uint& gridSizeY, const uint& numOutputClasses, 66 | const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, const float modelScale); 67 | 68 | cudaError_t cudaYoloLayer_r ( 69 | const void* input, void* output, const uint& batchSize, 70 | const uint& gridSizeX, const uint& gridSizeY, const uint& numOutputClasses, 71 | const uint& numBBoxes, uint64_t outputSize, cudaStream_t stream, const float modelScale); 72 | 73 | YoloLayer::YoloLayer (const void* data, size_t length) 74 | { 75 | const char *d = static_cast(data); 76 | read(d, m_NumBoxes); 77 | read(d, m_NumClasses); 78 | read(d, m_GridSizeX); 79 | read(d, m_GridSizeY); 80 | read(d, m_OutputSize); 81 | 82 | read(d, m_type); 83 | read(d, m_new_coords); 84 | read(d, m_scale_x_y); 85 | read(d, m_beta_nms); 86 | uint anchorsSize; 87 | read(d, anchorsSize); 88 | for (uint i = 0; i < anchorsSize; i++) { 89 | float result; 90 | read(d, result); 91 | m_Anchors.push_back(result); 92 | } 93 | uint maskSize; 94 | read(d, maskSize); 95 | for (uint i = 0; i < maskSize; i++) { 96 | uint nMask; 97 | read(d, nMask); 98 | std::vector pMask; 99 | for (uint f = 0; f < nMask; f++) { 100 | int result; 101 | read(d, result); 102 | pMask.push_back(result); 103 | } 104 | m_Mask.push_back(pMask); 105 | } 106 | kNUM_CLASSES = m_NumClasses; 107 | kBETA_NMS = m_beta_nms; 108 | kANCHORS = m_Anchors; 109 | kMASK = m_Mask; 110 | }; 111 | 112 | YoloLayer::YoloLayer ( 113 | const uint& numBoxes, const uint& numClasses, const uint& gridSizeX, const uint& gridSizeY, const uint model_type, const uint new_coords, const float scale_x_y, const float beta_nms, const std::vector anchors, std::vector> mask) : 114 | m_NumBoxes(numBoxes), 115 | m_NumClasses(numClasses), 116 | m_GridSizeX(gridSizeX), 117 | m_GridSizeY(gridSizeY), 118 | m_type(model_type), 119 | m_new_coords(new_coords), 120 | m_scale_x_y(scale_x_y), 121 | m_beta_nms(beta_nms), 122 | m_Anchors(anchors), 123 | m_Mask(mask) 124 | { 125 | assert(m_NumBoxes > 0); 126 | assert(m_NumClasses > 0); 127 | assert(m_GridSizeX > 0); 128 | assert(m_GridSizeY > 0); 129 | m_OutputSize = m_GridSizeX * m_GridSizeY * (m_NumBoxes * (4 + 1 + m_NumClasses)); 130 | }; 131 | 132 | nvinfer1::Dims 133 | YoloLayer::getOutputDimensions( 134 | int index, const nvinfer1::Dims* inputs, int nbInputDims) noexcept 135 | { 136 | assert(index == 0); 137 | assert(nbInputDims == 1); 138 | return inputs[0]; 139 | } 140 | 141 | bool YoloLayer::supportsFormat ( 142 | nvinfer1::DataType type, nvinfer1::PluginFormat format) const noexcept { 143 | return (type == nvinfer1::DataType::kFLOAT && 144 | format == nvinfer1::PluginFormat::kLINEAR); 145 | } 146 | 147 | void 148 | YoloLayer::configureWithFormat ( 149 | const nvinfer1::Dims* inputDims, int nbInputs, 150 | const nvinfer1::Dims* outputDims, int nbOutputs, 151 | nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) noexcept 152 | { 153 | assert(nbInputs == 1); 154 | assert (format == nvinfer1::PluginFormat::kLINEAR); 155 | assert(inputDims != nullptr); 156 | } 157 | 158 | int YoloLayer::enqueue( 159 | int batchSize, void const* const* inputs, void* const* outputs, void* workspace, 160 | cudaStream_t stream) noexcept 161 | { 162 | if (m_type == 2) { // YOLOR incorrect param 163 | CHECK(cudaYoloLayer_r( 164 | inputs[0], outputs[0], batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, 165 | m_OutputSize, stream, m_scale_x_y)); 166 | } 167 | else if (m_type == 1) { 168 | if (m_new_coords) { 169 | CHECK(cudaYoloLayer_nc( 170 | inputs[0], outputs[0], batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, 171 | m_OutputSize, stream, m_scale_x_y)); 172 | } 173 | else { 174 | CHECK(cudaYoloLayer( 175 | inputs[0], outputs[0], batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, 176 | m_OutputSize, stream, m_scale_x_y)); 177 | } 178 | } 179 | else { 180 | CHECK(cudaYoloLayer_v2( 181 | inputs[0], outputs[0], batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, 182 | m_OutputSize, stream)); 183 | } 184 | return 0; 185 | } 186 | 187 | size_t YoloLayer::getSerializationSize() const noexcept 188 | { 189 | int anchorsSum = 1; 190 | for (uint i = 0; i < m_Anchors.size(); i++) { 191 | anchorsSum += 1; 192 | } 193 | int maskSum = 1; 194 | for (uint i = 0; i < m_Mask.size(); i++) { 195 | maskSum += 1; 196 | for (uint f = 0; f < m_Mask[i].size(); f++) { 197 | maskSum += 1; 198 | } 199 | } 200 | 201 | return sizeof(m_NumBoxes) + sizeof(m_NumClasses) + sizeof(m_GridSizeX) + sizeof(m_GridSizeY) + sizeof(m_OutputSize) + sizeof(m_type) 202 | + sizeof(m_new_coords) + sizeof(m_scale_x_y) + sizeof(m_beta_nms) + anchorsSum * sizeof(float) + maskSum * sizeof(int); 203 | } 204 | 205 | void YoloLayer::serialize(void* buffer) const noexcept 206 | { 207 | char *d = static_cast(buffer); 208 | write(d, m_NumBoxes); 209 | write(d, m_NumClasses); 210 | write(d, m_GridSizeX); 211 | write(d, m_GridSizeY); 212 | write(d, m_OutputSize); 213 | 214 | write(d, m_type); 215 | write(d, m_new_coords); 216 | write(d, m_scale_x_y); 217 | write(d, m_beta_nms); 218 | uint anchorsSize = m_Anchors.size(); 219 | write(d, anchorsSize); 220 | for (uint i = 0; i < anchorsSize; i++) { 221 | write(d, m_Anchors[i]); 222 | } 223 | uint maskSize = m_Mask.size(); 224 | write(d, maskSize); 225 | for (uint i = 0; i < maskSize; i++) { 226 | uint pMaskSize = m_Mask[i].size(); 227 | write(d, pMaskSize); 228 | for (uint f = 0; f < pMaskSize; f++) { 229 | write(d, m_Mask[i][f]); 230 | } 231 | } 232 | kNUM_CLASSES = m_NumClasses; 233 | kBETA_NMS = m_beta_nms; 234 | kANCHORS = m_Anchors; 235 | kMASK = m_Mask; 236 | } 237 | 238 | nvinfer1::IPluginV2* YoloLayer::clone() const noexcept 239 | { 240 | return new YoloLayer (m_NumBoxes, m_NumClasses, m_GridSizeX, m_GridSizeY, m_type, m_new_coords, m_scale_x_y, m_beta_nms, m_Anchors, m_Mask); 241 | } 242 | 243 | REGISTER_TENSORRT_PLUGIN(YoloLayerPluginCreator); -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloPlugins.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019-2021, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a 5 | * copy of this software and associated documentation files (the "Software"), 6 | * to deal in the Software without restriction, including without limitation 7 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, 8 | * and/or sell copies of the Software, and to permit persons to whom the 9 | * Software is furnished to do so, subject to the following conditions: 10 | * 11 | * The above copyright notice and this permission notice shall be included in 12 | * all copies or substantial portions of the Software. 13 | * 14 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 15 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 16 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 17 | * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 18 | * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 19 | * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 20 | * DEALINGS IN THE SOFTWARE. 21 | 22 | * Edited by Marcos Luciano 23 | * https://www.github.com/marcoslucianops 24 | */ 25 | 26 | #ifndef __YOLO_PLUGINS__ 27 | #define __YOLO_PLUGINS__ 28 | 29 | #include 30 | #include 31 | #include 32 | #include 33 | #include 34 | 35 | #include 36 | 37 | #include "NvInferPlugin.h" 38 | 39 | #define CHECK(status) \ 40 | { \ 41 | if (status != 0) \ 42 | { \ 43 | std::cout << "CUDA failure: " << cudaGetErrorString(status) << " in file " << __FILE__ \ 44 | << " at line " << __LINE__ << std::endl; \ 45 | abort(); \ 46 | } \ 47 | } 48 | 49 | namespace 50 | { 51 | const char* YOLOLAYER_PLUGIN_VERSION {"1"}; 52 | const char* YOLOLAYER_PLUGIN_NAME {"YoloLayer_TRT"}; 53 | } // namespace 54 | 55 | class YoloLayer : public nvinfer1::IPluginV2 56 | { 57 | public: 58 | YoloLayer (const void* data, size_t length); 59 | YoloLayer (const uint& numBoxes, const uint& numClasses, const uint& gridSizeX, const uint& gridSizeY, 60 | const uint model_type, const uint new_coords, const float scale_x_y, const float beta_nms, 61 | const std::vector anchors, const std::vector> mask); 62 | const char* getPluginType () const noexcept override { return YOLOLAYER_PLUGIN_NAME; } 63 | const char* getPluginVersion () const noexcept override { return YOLOLAYER_PLUGIN_VERSION; } 64 | int getNbOutputs () const noexcept override { return 1; } 65 | 66 | nvinfer1::Dims getOutputDimensions ( 67 | int index, const nvinfer1::Dims* inputs, 68 | int nbInputDims) noexcept override; 69 | 70 | bool supportsFormat ( 71 | nvinfer1::DataType type, nvinfer1::PluginFormat format) const noexcept override; 72 | 73 | void configureWithFormat ( 74 | const nvinfer1::Dims* inputDims, int nbInputs, 75 | const nvinfer1::Dims* outputDims, int nbOutputs, 76 | nvinfer1::DataType type, nvinfer1::PluginFormat format, int maxBatchSize) noexcept override; 77 | 78 | int initialize () noexcept override { return 0; } 79 | void terminate () noexcept override {} 80 | size_t getWorkspaceSize (int maxBatchSize) const noexcept override { return 0; } 81 | int enqueue ( 82 | int batchSize, void const* const* inputs, void* const* outputs, 83 | void* workspace, cudaStream_t stream) noexcept override; 84 | size_t getSerializationSize() const noexcept override; 85 | void serialize (void* buffer) const noexcept override; 86 | void destroy () noexcept override { delete this; } 87 | nvinfer1::IPluginV2* clone() const noexcept override; 88 | 89 | void setPluginNamespace (const char* pluginNamespace) noexcept override { 90 | m_Namespace = pluginNamespace; 91 | } 92 | virtual const char* getPluginNamespace () const noexcept override { 93 | return m_Namespace.c_str(); 94 | } 95 | 96 | private: 97 | uint m_NumBoxes {0}; 98 | uint m_NumClasses {0}; 99 | uint m_GridSizeX {0}; 100 | uint m_GridSizeY {0}; 101 | uint64_t m_OutputSize {0}; 102 | std::string m_Namespace {""}; 103 | 104 | uint m_type {0}; 105 | uint m_new_coords {0}; 106 | float m_scale_x_y {0}; 107 | float m_beta_nms {0}; 108 | std::vector m_Anchors; 109 | std::vector> m_Mask; 110 | }; 111 | 112 | class YoloLayerPluginCreator : public nvinfer1::IPluginCreator 113 | { 114 | public: 115 | YoloLayerPluginCreator () {} 116 | ~YoloLayerPluginCreator () {} 117 | 118 | const char* getPluginName () const noexcept override { return YOLOLAYER_PLUGIN_NAME; } 119 | const char* getPluginVersion () const noexcept override { return YOLOLAYER_PLUGIN_VERSION; } 120 | 121 | const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override { 122 | std::cerr<< "YoloLayerPluginCreator::getFieldNames is not implemented" << std::endl; 123 | return nullptr; 124 | } 125 | 126 | nvinfer1::IPluginV2* createPlugin ( 127 | const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept override 128 | { 129 | std::cerr<< "YoloLayerPluginCreator::getFieldNames is not implemented"; 130 | return nullptr; 131 | } 132 | 133 | nvinfer1::IPluginV2* deserializePlugin ( 134 | const char* name, const void* serialData, size_t serialLength) noexcept override 135 | { 136 | std::cout << "Deserialize yoloLayer plugin: " << name << std::endl; 137 | return new YoloLayer(serialData, serialLength); 138 | } 139 | 140 | void setPluginNamespace(const char* libNamespace) noexcept override { 141 | m_Namespace = libNamespace; 142 | } 143 | const char* getPluginNamespace() const noexcept override { 144 | return m_Namespace.c_str(); 145 | } 146 | 147 | private: 148 | std::string m_Namespace {""}; 149 | }; 150 | 151 | extern int kNUM_CLASSES; 152 | extern float kBETA_NMS; 153 | extern std::vector kANCHORS; 154 | extern std::vector> kMASK; 155 | 156 | #endif // __YOLO_PLUGINS__ 157 | -------------------------------------------------------------------------------- /Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloPlugins.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Deepstream-app/nvdsinfer_custom_impl_Yolo/yoloPlugins.o -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Clinton Oduor 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 |
3 | 4 | # PCB-Defect-Detection-using-Deepstream 5 |
6 |
7 | 8 | Manual visual inspection is one of the most complex and expensive tasks for PCB manufacturing companies. Over the years, Printed Circuit Boards have become much smaller and more densely packed with components making manual visual inspection less scalable. With increased demands from the electronics industry, many defects go unnoticed which may lead to poor company reputation ,and reduced business. The figure below shows an image of a PCB automated optical inspection procedure at PCBWay. 9 |
10 |
11 | 12 |
13 |

14 | 15 | 16 |

17 |
18 |
19 |
20 | image source: https://www.etechnophiles.com/pcb-testing-methods-and-properties-7-ways-to-ensure-best-quality-pcbs/ 21 |
22 | 23 |
24 | 25 | This project uses deeplearning with object detection to automatically detect 6 PCB Defects using YoloV5 and Deepstream SDK running on an Nvidia jetson AGX Xavier. 26 | 27 | 28 | 29 | The defects include: 30 | 31 | * Missing hole 32 | * Mouse bite 33 | * Open circuit 34 | * Short circuit 35 | * Spurious copper 36 | * spur 37 | 38 | 39 | ### Steps 40 | 41 | 1. [Model training](https://github.com/clintonoduor/PCB-Defect-Detection-using-Deepstream#model-training) 42 | 43 | 2. [TensorRT engine generation](https://github.com/clintonoduor/PCB-Defect-Detection-using-Deepstream#generating-wts--cfg-files-for-tensorrt-engine) 44 | 3. [Configuring Deepstream app](https://github.com/clintonoduor/PCB-Defect-Detection-using-Deepstream#configuring-deepstream) 45 | 4. [Optimizing Deepstream app for perfomance](https://github.com/clintonoduor/PCB-Defect-Detection-using-Deepstream#improving-perfomance-tricks) 46 | 47 | ### Model Training 48 | 49 | The model was trained using YoloV5 on Google Colab. The dataset used is hosted on Roboflow to make it easier for others to be able to train and reproduce the model in Colab without the hussle of re-uploading the data. 50 | 51 | YoloV5 🚀 is a family of object detection architectures and models pretrained on the COCO dataset, and represents Ultralytics open-source research into future vision AI methods, incorporating lessons learned and best practices evolved over thousands of hours of research and development. 52 | 53 | The choice to use YoloV5 was based on its inherent fast speed, smaller size, and its ability to learn to auto tune anchor boxes based on the distribution of the training set. 54 | 55 | The full code used to train the model for this project can be found by clicking the Colab icon below: 56 | 57 |
58 | 59 | 60 | 61 |
62 |
63 | 64 | The dataset used was sourced from the Open Lab on Human Robot Interaction of Peking University. 65 | 66 | #### Training Results 67 | 68 | Metrics were closely monitored and tracked using Weights & Biases for 1083 epochs as shown in the images below: 69 | 70 |
71 | 72 |
73 |

74 | 75 | 76 |

77 |
78 | 79 |
80 | 81 |
82 |

83 | 84 | 85 |

86 |
87 |
88 | 89 | 90 | 91 |
92 |

93 | 94 | 95 |

96 |
97 | 98 | ### Generating wts & cfg files for TensorRT engine 99 | 100 | [NVIDIA® TensorRT™](https://developer.nvidia.com/tensorrt), is a machine learning framework and SDK for high-performance deep learning inference for GPU accelerated workloads such as intelligent video analytics on jetson devices. It includes a deep learning inference optimizer and runtime that delivers low latency and high throughput for inference. 101 | 102 | Since YoloV5 is based on Pytorch, we will first need to convert our trained model to a format that can be easily converted to TensorRT engine when starting the deepstream application on the Jetson Xavier. To do this,we first put the gen_wts_yoloV5.py file to the YoloV5 folder and run the following script: 103 | 104 | ``` 105 | !python3 gen_wts_yoloV5.py -w /content/yolov5/runs/train/yolov5s_results/weights/best.pt -c /content/yolov5/models/custom_yolov5s.yaml 106 | ``` 107 | ``` 108 | # -w is the path to our best.pt model while -c is the file to our input yaml file. 109 | ``` 110 | The script will output a wts and cfg file that we will put in our Deepstream-Yolo folder to build the TensorRT engine. 111 | 112 | ### Configuring Deepstream 113 | The [NVIDIA DeepStream SDK](https://developer.nvidia.com/deepstream-sdk) delivers a complete streaming analytics toolkit for AI based video and image understanding and multi-sensor processing. DeepStream SDK features hardware-accelerated building blocks called plugins, that bring deep neural networks and other complex processing tasks into a stream processing pipeline. 114 | ##### Deepstream Installation 115 | First install deepstream 6.0 using these [instructions](https://docs.nvidia.com/metropolis/deepstream/dev-guide/text/DS_Quickstart.html). 116 | ##### Editing Deepstream Config Files 117 | ``` 118 | In the Config_infer_primary_yoloV5.txt, we replace the model-file and custom-network-config with the paths to our generated weight and config files respectively. We also change the num-detected-classes with the number of our classes which is 6 as shown below. 119 | 120 | [property] 121 | gpu-id=0 122 | net-scale-factor=0.0039215697906911373 123 | model-color-format=0 124 | custom-network-config=best.cfg 125 | model-file=best.wts 126 | model-engine-file=model_b1_gpu0_fp32.engine 127 | #int8-calib-file=calib.table 128 | labelfile-path=labels.txt 129 | batch-size=1 130 | network-mode=0 131 | num-detected-classes=6 132 | interval=0 133 | gie-unique-id=1 134 | process-mode=1 135 | network-type=0 136 | cluster-mode=4 137 | maintain-aspect-ratio=1 138 | parse-bbox-func-name=NvDsInferParseYolo 139 | custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so 140 | engine-create-func-name=NvDsInferYoloCudaEngineGet 141 | 142 | [class-attrs-all] 143 | pre-cluster-threshold=0.5 144 | ``` 145 | ##### Running the application 146 | 147 | First compile the application using: 148 | 149 | ``` 150 | CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo 151 | ``` 152 | Then boost your Jetson clocks using the following commands: 153 | ``` 154 | $ sudo nvpmodel -m 0 155 | $ sudo jetson_clocks 156 | ``` 157 | 158 | ###### Running application using a sample video (.MP4) 159 | 160 | 1. Go to **deepstream_app_config.txt > [Source 0]** then edit the URI with the path of your sample MP4 video as shown below: 161 | 162 | 163 | ``` 164 | [source0] 165 | enable=1 166 | type=3 167 | uri=file:///opt/nvidia/deepstream/deepstream/samples/streams/pcb2.mp4 168 | num-sources=1 169 | gpu-id=0 170 | cudadec-memtype=0 171 | ``` 172 | 173 | 2. Run the command below on your terminal 174 | 175 | ``` 176 | deepstream-app -c deepstream_app_config.txt 177 | ``` 178 | 179 | **NB:** This dataset used for training this model was synthetic hence the trained model cannot be validated using real world data. However the instructions below could be used to run the model using a USB cam or CSI camera: 180 | 181 | ###### Running application using a Webcam 182 | 183 | To run the application using a USB cam, run the below command on your terminal: 184 | ``` 185 | deepstream-app -c deepstream_app_config_USB.txt 186 | ``` 187 | 188 | ###### Running application using a CSI cam 189 | 190 | To run the application using a CSI cam, run the command below on your terminal: 191 | ``` 192 | deepstream-app -c deepstream_app_config(CSI).txt 193 | ``` 194 | 195 | The application outputs a tiled display with on screen bounding boxes of the detected defects as shown in the image below: 196 | 197 |
198 | 199 | [![Video demo on youtube](https://github.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/blob/main/assets/pcbscreenshot.png?raw=true)](https://www.youtube.com/watch?v=op_TjAQFLfs) 200 | 201 |
202 |
203 | The model runs at an average of 51 FPS on a Jetson AGX Xavier, hence one xavier could be configured to handle approximately 10 streams running in parallel at 5 FPS. 204 | 205 | ### Optimizing Deepstream app for Perfomance 206 | 207 | There are several techniques that could be used to increase the perfomance of the deepstream app: 208 | 209 | ##### 1. Increasing the number of interval in the infer config file 210 | 211 | ``` 212 | network-mode=0 213 | num-detected-classes=6 214 | interval=1 215 | gie-unique-id=1 216 | process-mode=1 217 | network-type=0 218 | ``` 219 | ##### 2. Disable tiles & on screen perfomance when not needed 220 | 221 | To achieve this, you just need to go to the deepstream-app_config.txtfile and set enable = 0 for [tiled-display] [osd] as shown below: 222 | 223 | ``` 224 | [tiled-display] 225 | enable=0 226 | rows=1 227 | columns=1 228 | width=1280 229 | height=720 230 | gpu-id=0 231 | nvbuf-memory-type=0 232 | ``` 233 | ``` 234 | [osd] 235 | enable=0 236 | gpu-id=0 237 | border-width=1 238 | text-size=15 239 | text-color=1;1;1;1; 240 | text-bg-color=0.3;0.3;0.3;1 241 | font=Serif 242 | show-clock=0 243 | clock-x-offset=800 244 | clock-y-offset=820 245 | clock-text-size=12 246 | clock-color=1;0;0;0 247 | nvbuf-memory-type=0 248 | ``` 249 | ##### 3. Reducing model precision 250 | 251 | Change network precision to fp16 by changing network-mode=0 to network-mode=2 in the deepstream config infer file as shown below: 252 | 253 | ``` 254 | model-engine-file=model_b1_gpu0_fp32.engine 255 | #int8-calib-file=calib.table 256 | labelfile-path=labels.txt 257 | batch-size=1 258 | network-mode=2 259 | num-detected-classes=6 260 | ``` 261 | 262 | ##### 4. Always make sure the stream mux output width and height to the size of the input stream resolution 263 | 264 | ``` 265 | [streammux] 266 | ##Boolean property to inform muxer that sources are live 267 | live-source=1 268 | batch-size=1 269 | ##time out in usec, to wait after the first buffer is available 270 | ##to push the batch even if the complete batch is not formed 271 | batched-push-timeout=40000 272 | ## Set muxer output width and height 273 | width=1280 274 | height=720 275 | ``` 276 |
277 | 278 | ### Acknowledgement 279 | 280 |
281 | 282 | 1. Open Lab on Human Robot Interaction of Peking University for use of dataset. 283 | 2. Passing custom YoloV5 models for deepstream applications: https://github.com/marcoslucianops/DeepStream-Yolo 284 | 285 | 286 |
287 | 288 | ### Future works 289 | 290 |
291 | 292 | 1. Implement the project using deepstream python bindings 293 | 2. Performing analytics using deepstream analytics 294 | 3. Train using real PCB images with defects 295 | 296 | -------------------------------------------------------------------------------- /Weights/best.cfg: -------------------------------------------------------------------------------- 1 | [net] 2 | width=640 3 | height=640 4 | channels=3 5 | 6 | [convolutional] 7 | batch_normalize=1 8 | filters=64 9 | size=3 10 | stride=2 11 | pad=1 12 | activation=silu 13 | 14 | # C3 15 | 16 | [convolutional] 17 | batch_normalize=1 18 | filters=32 19 | size=1 20 | stride=1 21 | pad=1 22 | activation=silu 23 | 24 | [route] 25 | layers=-2 26 | 27 | [convolutional] 28 | batch_normalize=1 29 | filters=32 30 | size=1 31 | stride=1 32 | pad=1 33 | activation=silu 34 | 35 | [convolutional] 36 | batch_normalize=1 37 | filters=32 38 | size=1 39 | stride=1 40 | pad=1 41 | activation=silu 42 | 43 | [convolutional] 44 | batch_normalize=1 45 | filters=32 46 | size=3 47 | stride=1 48 | pad=1 49 | activation=silu 50 | 51 | [shortcut] 52 | from=-3 53 | activation=linear 54 | 55 | [route] 56 | layers=-1, -6 57 | 58 | [convolutional] 59 | batch_normalize=1 60 | filters=64 61 | size=1 62 | stride=1 63 | pad=1 64 | activation=silu 65 | 66 | ########## 67 | 68 | [convolutional] 69 | batch_normalize=1 70 | filters=128 71 | size=3 72 | stride=2 73 | pad=1 74 | activation=silu 75 | 76 | # C3 77 | 78 | [convolutional] 79 | batch_normalize=1 80 | filters=64 81 | size=1 82 | stride=1 83 | pad=1 84 | activation=silu 85 | 86 | [route] 87 | layers=-2 88 | 89 | [convolutional] 90 | batch_normalize=1 91 | filters=64 92 | size=1 93 | stride=1 94 | pad=1 95 | activation=silu 96 | 97 | [convolutional] 98 | batch_normalize=1 99 | filters=64 100 | size=1 101 | stride=1 102 | pad=1 103 | activation=silu 104 | 105 | [convolutional] 106 | batch_normalize=1 107 | filters=64 108 | size=3 109 | stride=1 110 | pad=1 111 | activation=silu 112 | 113 | [shortcut] 114 | from=-3 115 | activation=linear 116 | 117 | [convolutional] 118 | batch_normalize=1 119 | filters=64 120 | size=1 121 | stride=1 122 | pad=1 123 | activation=silu 124 | 125 | [convolutional] 126 | batch_normalize=1 127 | filters=64 128 | size=3 129 | stride=1 130 | pad=1 131 | activation=silu 132 | 133 | [shortcut] 134 | from=-3 135 | activation=linear 136 | 137 | [convolutional] 138 | batch_normalize=1 139 | filters=64 140 | size=1 141 | stride=1 142 | pad=1 143 | activation=silu 144 | 145 | [convolutional] 146 | batch_normalize=1 147 | filters=64 148 | size=3 149 | stride=1 150 | pad=1 151 | activation=silu 152 | 153 | [shortcut] 154 | from=-3 155 | activation=linear 156 | 157 | [route] 158 | layers=-1, -12 159 | 160 | [convolutional] 161 | batch_normalize=1 162 | filters=128 163 | size=1 164 | stride=1 165 | pad=1 166 | activation=silu 167 | 168 | ########## 169 | 170 | [convolutional] 171 | batch_normalize=1 172 | filters=256 173 | size=3 174 | stride=2 175 | pad=1 176 | activation=silu 177 | 178 | # C3 179 | 180 | [convolutional] 181 | batch_normalize=1 182 | filters=128 183 | size=1 184 | stride=1 185 | pad=1 186 | activation=silu 187 | 188 | [route] 189 | layers=-2 190 | 191 | [convolutional] 192 | batch_normalize=1 193 | filters=128 194 | size=1 195 | stride=1 196 | pad=1 197 | activation=silu 198 | 199 | [convolutional] 200 | batch_normalize=1 201 | filters=128 202 | size=1 203 | stride=1 204 | pad=1 205 | activation=silu 206 | 207 | [convolutional] 208 | batch_normalize=1 209 | filters=128 210 | size=3 211 | stride=1 212 | pad=1 213 | activation=silu 214 | 215 | [shortcut] 216 | from=-3 217 | activation=linear 218 | 219 | [convolutional] 220 | batch_normalize=1 221 | filters=128 222 | size=1 223 | stride=1 224 | pad=1 225 | activation=silu 226 | 227 | [convolutional] 228 | batch_normalize=1 229 | filters=128 230 | size=3 231 | stride=1 232 | pad=1 233 | activation=silu 234 | 235 | [shortcut] 236 | from=-3 237 | activation=linear 238 | 239 | [convolutional] 240 | batch_normalize=1 241 | filters=128 242 | size=1 243 | stride=1 244 | pad=1 245 | activation=silu 246 | 247 | [convolutional] 248 | batch_normalize=1 249 | filters=128 250 | size=3 251 | stride=1 252 | pad=1 253 | activation=silu 254 | 255 | [shortcut] 256 | from=-3 257 | activation=linear 258 | 259 | [route] 260 | layers=-1, -12 261 | 262 | [convolutional] 263 | batch_normalize=1 264 | filters=256 265 | size=1 266 | stride=1 267 | pad=1 268 | activation=silu 269 | 270 | ########## 271 | 272 | [convolutional] 273 | batch_normalize=1 274 | filters=512 275 | size=3 276 | stride=2 277 | pad=1 278 | activation=silu 279 | 280 | # C3 281 | 282 | [convolutional] 283 | batch_normalize=1 284 | filters=256 285 | size=1 286 | stride=1 287 | pad=1 288 | activation=silu 289 | 290 | [route] 291 | layers=-2 292 | 293 | [convolutional] 294 | batch_normalize=1 295 | filters=256 296 | size=1 297 | stride=1 298 | pad=1 299 | activation=silu 300 | 301 | [convolutional] 302 | batch_normalize=1 303 | filters=256 304 | size=1 305 | stride=1 306 | pad=1 307 | activation=silu 308 | 309 | [convolutional] 310 | batch_normalize=1 311 | filters=256 312 | size=3 313 | stride=1 314 | pad=1 315 | activation=silu 316 | 317 | [route] 318 | layers=-1, -5 319 | 320 | [convolutional] 321 | batch_normalize=1 322 | filters=512 323 | size=1 324 | stride=1 325 | pad=1 326 | activation=silu 327 | 328 | ########## 329 | 330 | [convolutional] 331 | batch_normalize=1 332 | filters=256 333 | size=1 334 | stride=1 335 | pad=1 336 | activation=silu 337 | 338 | [upsample] 339 | stride=2 340 | 341 | # Concat 342 | 343 | [route] 344 | layers=-1, 39 345 | 346 | ########## 347 | 348 | # C3 349 | 350 | [convolutional] 351 | batch_normalize=1 352 | filters=128 353 | size=1 354 | stride=1 355 | pad=1 356 | activation=silu 357 | 358 | [route] 359 | layers=-2 360 | 361 | [convolutional] 362 | batch_normalize=1 363 | filters=128 364 | size=1 365 | stride=1 366 | pad=1 367 | activation=silu 368 | 369 | [convolutional] 370 | batch_normalize=1 371 | filters=128 372 | size=1 373 | stride=1 374 | pad=1 375 | activation=silu 376 | 377 | [convolutional] 378 | batch_normalize=1 379 | filters=128 380 | size=3 381 | stride=1 382 | pad=1 383 | activation=silu 384 | 385 | [route] 386 | layers=-1, -5 387 | 388 | [convolutional] 389 | batch_normalize=1 390 | filters=256 391 | size=1 392 | stride=1 393 | pad=1 394 | activation=silu 395 | 396 | ########## 397 | 398 | [convolutional] 399 | batch_normalize=1 400 | filters=128 401 | size=1 402 | stride=1 403 | pad=1 404 | activation=silu 405 | 406 | [upsample] 407 | stride=2 408 | 409 | # Concat 410 | 411 | [route] 412 | layers=-1, 24 413 | 414 | ########## 415 | 416 | # C3 417 | 418 | [convolutional] 419 | batch_normalize=1 420 | filters=64 421 | size=1 422 | stride=1 423 | pad=1 424 | activation=silu 425 | 426 | [route] 427 | layers=-2 428 | 429 | [convolutional] 430 | batch_normalize=1 431 | filters=64 432 | size=1 433 | stride=1 434 | pad=1 435 | activation=silu 436 | 437 | [convolutional] 438 | batch_normalize=1 439 | filters=64 440 | size=1 441 | stride=1 442 | pad=1 443 | activation=silu 444 | 445 | [convolutional] 446 | batch_normalize=1 447 | filters=64 448 | size=3 449 | stride=1 450 | pad=1 451 | activation=silu 452 | 453 | [route] 454 | layers=-1, -5 455 | 456 | [convolutional] 457 | batch_normalize=1 458 | filters=128 459 | size=1 460 | stride=1 461 | pad=1 462 | activation=silu 463 | 464 | ########## 465 | 466 | [convolutional] 467 | batch_normalize=1 468 | filters=128 469 | size=3 470 | stride=2 471 | pad=1 472 | activation=silu 473 | 474 | # Concat 475 | 476 | [route] 477 | layers=-1, 59 478 | 479 | ########## 480 | 481 | # C3 482 | 483 | [convolutional] 484 | batch_normalize=1 485 | filters=128 486 | size=1 487 | stride=1 488 | pad=1 489 | activation=silu 490 | 491 | [route] 492 | layers=-2 493 | 494 | [convolutional] 495 | batch_normalize=1 496 | filters=128 497 | size=1 498 | stride=1 499 | pad=1 500 | activation=silu 501 | 502 | [convolutional] 503 | batch_normalize=1 504 | filters=128 505 | size=1 506 | stride=1 507 | pad=1 508 | activation=silu 509 | 510 | [convolutional] 511 | batch_normalize=1 512 | filters=128 513 | size=3 514 | stride=1 515 | pad=1 516 | activation=silu 517 | 518 | [route] 519 | layers=-1, -5 520 | 521 | [convolutional] 522 | batch_normalize=1 523 | filters=256 524 | size=1 525 | stride=1 526 | pad=1 527 | activation=silu 528 | 529 | ########## 530 | 531 | [convolutional] 532 | batch_normalize=1 533 | filters=256 534 | size=3 535 | stride=2 536 | pad=1 537 | activation=silu 538 | 539 | # Concat 540 | 541 | [route] 542 | layers=-1, 49 543 | 544 | ########## 545 | 546 | # C3 547 | 548 | [convolutional] 549 | batch_normalize=1 550 | filters=256 551 | size=1 552 | stride=1 553 | pad=1 554 | activation=silu 555 | 556 | [route] 557 | layers=-2 558 | 559 | [convolutional] 560 | batch_normalize=1 561 | filters=256 562 | size=1 563 | stride=1 564 | pad=1 565 | activation=silu 566 | 567 | [convolutional] 568 | batch_normalize=1 569 | filters=256 570 | size=1 571 | stride=1 572 | pad=1 573 | activation=silu 574 | 575 | [convolutional] 576 | batch_normalize=1 577 | filters=256 578 | size=3 579 | stride=1 580 | pad=1 581 | activation=silu 582 | 583 | [route] 584 | layers=-1, -5 585 | 586 | [convolutional] 587 | batch_normalize=1 588 | filters=512 589 | size=1 590 | stride=1 591 | pad=1 592 | activation=silu 593 | 594 | ########## 595 | 596 | # Detect 597 | 598 | [route] 599 | layers=68 600 | 601 | [convolutional] 602 | size=1 603 | stride=1 604 | pad=1 605 | filters=33 606 | activation=logistic 607 | 608 | [yolo] 609 | mask=0, 1, 2 610 | anchors=10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326 611 | classes=6 612 | num=9 613 | scale_x_y=2.0 614 | beta_nms=0.6 615 | new_coords=1 616 | 617 | ########## 618 | 619 | # Detect 620 | 621 | [route] 622 | layers=77 623 | 624 | [convolutional] 625 | size=1 626 | stride=1 627 | pad=1 628 | filters=33 629 | activation=logistic 630 | 631 | [yolo] 632 | mask=3, 4, 5 633 | anchors=10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326 634 | classes=6 635 | num=9 636 | scale_x_y=2.0 637 | beta_nms=0.6 638 | new_coords=1 639 | 640 | ########## 641 | 642 | # Detect 643 | 644 | [route] 645 | layers=90 646 | 647 | [convolutional] 648 | size=1 649 | stride=1 650 | pad=1 651 | filters=33 652 | activation=logistic 653 | 654 | [yolo] 655 | mask=6, 7, 8 656 | anchors=10, 13, 16, 30, 33, 23, 30, 61, 62, 45, 59, 119, 116, 90, 156, 198, 373, 326 657 | classes=6 658 | num=9 659 | scale_x_y=2.0 660 | beta_nms=0.6 661 | new_coords=1 662 | 663 | ########## 664 | -------------------------------------------------------------------------------- /Weights/best.pt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/Weights/best.pt -------------------------------------------------------------------------------- /YoloV5/gen_wts_yoloV5.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import yaml 3 | import math 4 | import os 5 | import struct 6 | import torch 7 | from utils.torch_utils import select_device 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser(description="PyTorch YOLOv5 conversion") 12 | parser.add_argument("-w", "--weights", required=True, help="Input weights (.pt) file path (required)") 13 | parser.add_argument("-c", "--yaml", help="Input cfg (.yaml) file path") 14 | parser.add_argument("-mw", "--width", help="Model width (default = 640 / 1280 [P6])") 15 | parser.add_argument("-mh", "--height", help="Model height (default = 640 / 1280 [P6])") 16 | parser.add_argument("-mc", "--channels", help="Model channels (default = 3)") 17 | parser.add_argument("--p6", action="store_true", help="P6 model") 18 | args = parser.parse_args() 19 | if not os.path.isfile(args.weights): 20 | raise SystemExit("Invalid weights file") 21 | if not args.yaml: 22 | args.yaml = "" 23 | if not args.width: 24 | args.width = 1280 if args.p6 else 640 25 | if not args.height: 26 | args.height = 1280 if args.p6 else 640 27 | if not args.channels: 28 | args.channels = 3 29 | return args.weights, args.yaml, args.width, args.height, args.channels, args.p6 30 | 31 | 32 | def get_width(x, gw, divisor=8): 33 | return int(math.ceil((x * gw) / divisor)) * divisor 34 | 35 | 36 | def get_depth(x, gd): 37 | if x == 1: 38 | return 1 39 | r = int(round(x * gd)) 40 | if x * gd - int(x * gd) == 0.5 and int(x * gd) % 2 == 0: 41 | r -= 1 42 | return max(r, 1) 43 | 44 | 45 | pt_file, yaml_file, model_width, model_height, model_channels, p6 = parse_args() 46 | 47 | model_name = pt_file.split(".pt")[0] 48 | wts_file = model_name + ".wts" 49 | cfg_file = model_name + ".cfg" 50 | 51 | if yaml_file == "": 52 | yaml_file = "models/" + model_name + ".yaml" 53 | if not os.path.isfile(yaml_file): 54 | yaml_file = "models/hub/" + model_name + ".yaml" 55 | if not os.path.isfile(yaml_file): 56 | raise SystemExit("YAML file not found") 57 | elif not os.path.isfile(yaml_file): 58 | raise SystemExit("Invalid YAML file") 59 | 60 | device = select_device("cpu") 61 | model = torch.load(pt_file, map_location=device)["model"].float() 62 | model.to(device).eval() 63 | 64 | with open(wts_file, "w") as f: 65 | wts_write = "" 66 | conv_count = 0 67 | cv1 = "" 68 | cv3 = "" 69 | cv3_idx = 0 70 | sppf_idx = 11 if p6 else 9 71 | for k, v in model.state_dict().items(): 72 | if not "num_batches_tracked" in k and not "anchors" in k and not "anchor_grid" in k: 73 | vr = v.reshape(-1).cpu().numpy() 74 | idx = int(k.split(".")[1]) 75 | if ".cv1." in k and not ".m." in k and idx != sppf_idx: 76 | cv1 += "{} {} ".format(k, len(vr)) 77 | for vv in vr: 78 | cv1 += " " 79 | cv1 += struct.pack(">f", float(vv)).hex() 80 | cv1 += "\n" 81 | conv_count += 1 82 | elif cv1 != "" and ".m." in k: 83 | wts_write += cv1 84 | cv1 = "" 85 | if ".cv3." in k: 86 | cv3 += "{} {} ".format(k, len(vr)) 87 | for vv in vr: 88 | cv3 += " " 89 | cv3 += struct.pack(">f", float(vv)).hex() 90 | cv3 += "\n" 91 | cv3_idx = idx 92 | conv_count += 1 93 | elif cv3 != "" and cv3_idx != idx: 94 | wts_write += cv3 95 | cv3 = "" 96 | cv3_idx = 0 97 | if not ".cv3." in k and not (".cv1." in k and not ".m." in k and idx != sppf_idx): 98 | wts_write += "{} {} ".format(k, len(vr)) 99 | for vv in vr: 100 | wts_write += " " 101 | wts_write += struct.pack(">f", float(vv)).hex() 102 | wts_write += "\n" 103 | conv_count += 1 104 | f.write("{}\n".format(conv_count)) 105 | f.write(wts_write) 106 | 107 | with open(cfg_file, "w") as c: 108 | with open(yaml_file, "r") as f: 109 | nc = 0 110 | depth_multiple = 0 111 | width_multiple = 0 112 | anchors = "" 113 | masks = [] 114 | num = 0 115 | detections = [] 116 | layers = [] 117 | f = yaml.load(f,Loader=yaml.FullLoader) 118 | c.write("[net]\n") 119 | c.write("width=%d\n" % model_width) 120 | c.write("height=%d\n" % model_height) 121 | c.write("channels=%d\n" % model_channels) 122 | for l in f: 123 | if l == "nc": 124 | nc = f[l] 125 | elif l == "depth_multiple": 126 | depth_multiple = f[l] 127 | elif l == "width_multiple": 128 | width_multiple = f[l] 129 | elif l == "anchors": 130 | a = [] 131 | for v in f[l]: 132 | a.extend(v) 133 | mask = [] 134 | for _ in range(int(len(v) / 2)): 135 | mask.append(num) 136 | num += 1 137 | masks.append(mask) 138 | anchors = str(a)[1:-1] 139 | elif l == "backbone" or l == "head": 140 | for v in f[l]: 141 | if v[2] == "Conv": 142 | layer = "" 143 | blocks = 0 144 | layer += "\n[convolutional]\n" 145 | layer += "batch_normalize=1\n" 146 | layer += "filters=%d\n" % get_width(v[3][0], width_multiple) 147 | layer += "size=%d\n" % v[3][1] 148 | layer += "stride=%d\n" % v[3][2] 149 | layer += "pad=1\n" 150 | layer += "activation=silu\n" 151 | blocks += 1 152 | layers.append([layer, blocks]) 153 | elif v[2] == "C3": 154 | layer = "" 155 | blocks = 0 156 | layer += "\n# C3\n" 157 | # SPLIT 158 | layer += "\n[convolutional]\n" 159 | layer += "batch_normalize=1\n" 160 | layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple) 161 | layer += "size=1\n" 162 | layer += "stride=1\n" 163 | layer += "pad=1\n" 164 | layer += "activation=silu\n" 165 | blocks += 1 166 | layer += "\n[route]\n" 167 | layer += "layers=-2\n" 168 | blocks += 1 169 | layer += "\n[convolutional]\n" 170 | layer += "batch_normalize=1\n" 171 | layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple) 172 | layer += "size=1\n" 173 | layer += "stride=1\n" 174 | layer += "pad=1\n" 175 | layer += "activation=silu\n" 176 | blocks += 1 177 | # Residual Block 178 | if len(v[3]) == 1 or v[3][1] == True: 179 | for _ in range(get_depth(v[1], depth_multiple)): 180 | layer += "\n[convolutional]\n" 181 | layer += "batch_normalize=1\n" 182 | layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple) 183 | layer += "size=1\n" 184 | layer += "stride=1\n" 185 | layer += "pad=1\n" 186 | layer += "activation=silu\n" 187 | blocks += 1 188 | layer += "\n[convolutional]\n" 189 | layer += "batch_normalize=1\n" 190 | layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple) 191 | layer += "size=3\n" 192 | layer += "stride=1\n" 193 | layer += "pad=1\n" 194 | layer += "activation=silu\n" 195 | blocks += 1 196 | layer += "\n[shortcut]\n" 197 | layer += "from=-3\n" 198 | layer += "activation=linear\n" 199 | blocks += 1 200 | # Merge 201 | layer += "\n[route]\n" 202 | layer += "layers=-1, -%d\n" % (3 * get_depth(v[1], depth_multiple) + 3) 203 | blocks += 1 204 | else: 205 | for _ in range(get_depth(v[1], depth_multiple)): 206 | layer += "\n[convolutional]\n" 207 | layer += "batch_normalize=1\n" 208 | layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple) 209 | layer += "size=1\n" 210 | layer += "stride=1\n" 211 | layer += "pad=1\n" 212 | layer += "activation=silu\n" 213 | blocks += 1 214 | layer += "\n[convolutional]\n" 215 | layer += "batch_normalize=1\n" 216 | layer += "filters=%d\n" % get_width(v[3][0] / 2, width_multiple) 217 | layer += "size=3\n" 218 | layer += "stride=1\n" 219 | layer += "pad=1\n" 220 | layer += "activation=silu\n" 221 | blocks += 1 222 | # Merge 223 | layer += "\n[route]\n" 224 | layer += "layers=-1, -%d\n" % (2 * get_depth(v[1], depth_multiple) + 3) 225 | blocks += 1 226 | # Transition 227 | layer += "\n[convolutional]\n" 228 | layer += "batch_normalize=1\n" 229 | layer += "filters=%d\n" % get_width(v[3][0], width_multiple) 230 | layer += "size=1\n" 231 | layer += "stride=1\n" 232 | layer += "pad=1\n" 233 | layer += "activation=silu\n" 234 | layer += "\n##########\n" 235 | blocks += 1 236 | layers.append([layer, blocks]) 237 | elif v[2] == "SPPF": 238 | layer = "" 239 | blocks = 0 240 | layer += "\n# SPPF\n" 241 | layer += "\n[convolutional]\n" 242 | layer += "batch_normalize=1\n" 243 | layer += "filters=%d\n" % (get_width(v[3][0], width_multiple) / 2) 244 | layer += "size=1\n" 245 | layer += "stride=1\n" 246 | layer += "pad=1\n" 247 | layer += "activation=silu\n" 248 | blocks += 1 249 | layer += "\n[maxpool]\n" 250 | layer += "stride=1\n" 251 | layer += "size=%d\n" % v[3][1] 252 | blocks += 1 253 | layer += "\n[maxpool]\n" 254 | layer += "stride=1\n" 255 | layer += "size=%d\n" % v[3][1] 256 | blocks += 1 257 | layer += "\n[maxpool]\n" 258 | layer += "stride=1\n" 259 | layer += "size=%d\n" % v[3][1] 260 | blocks += 1 261 | layer += "\n[route]\n" 262 | layer += "layers=-4, -3, -2, -1\n" 263 | blocks += 1 264 | layer += "\n[convolutional]\n" 265 | layer += "batch_normalize=1\n" 266 | layer += "filters=%d\n" % get_width(v[3][0], width_multiple) 267 | layer += "size=1\n" 268 | layer += "stride=1\n" 269 | layer += "pad=1\n" 270 | layer += "activation=silu\n" 271 | layer += "\n##########\n" 272 | blocks += 1 273 | layers.append([layer, blocks]) 274 | elif v[2] == "nn.Upsample": 275 | layer = "" 276 | blocks = 0 277 | layer += "\n[upsample]\n" 278 | layer += "stride=%d\n" % v[3][1] 279 | blocks += 1 280 | layers.append([layer, blocks]) 281 | elif v[2] == "Concat": 282 | layer = "" 283 | blocks = 0 284 | route = v[0][1] 285 | r = 0 286 | if route > 0: 287 | for i, item in enumerate(layers): 288 | if i <= route: 289 | r += item[1] 290 | else: 291 | break 292 | else: 293 | route = len(layers) + route 294 | for i, item in enumerate(layers): 295 | if i <= route: 296 | r += item[1] 297 | else: 298 | break 299 | layer += "\n# Concat\n" 300 | layer += "\n[route]\n" 301 | layer += "layers=-1, %d\n" % (r - 1) 302 | layer += "\n##########\n" 303 | blocks += 1 304 | layers.append([layer, blocks]) 305 | elif v[2] == "Detect": 306 | for i, n in enumerate(v[0]): 307 | layer = "" 308 | blocks = 0 309 | r = 0 310 | for j, item in enumerate(layers): 311 | if j <= n: 312 | r += item[1] 313 | else: 314 | break 315 | layer += "\n# Detect\n" 316 | layer += "\n[route]\n" 317 | layer += "layers=%d\n" % (r - 1) 318 | blocks += 1 319 | layer += "\n[convolutional]\n" 320 | layer += "size=1\n" 321 | layer += "stride=1\n" 322 | layer += "pad=1\n" 323 | layer += "filters=%d\n" % ((nc + 5) * 3) 324 | layer += "activation=logistic\n" 325 | blocks += 1 326 | layer += "\n[yolo]\n" 327 | layer += "mask=%s\n" % str(masks[i])[1:-1] 328 | layer += "anchors=%s\n" % anchors 329 | layer += "classes=%d\n" % nc 330 | layer += "num=%d\n" % num 331 | layer += "scale_x_y=2.0\n" 332 | layer += "beta_nms=0.6\n" 333 | layer += "new_coords=1\n" 334 | layer += "\n##########\n" 335 | blocks += 1 336 | layers.append([layer, blocks]) 337 | for layer in layers: 338 | c.write(layer[0]) 339 | -------------------------------------------------------------------------------- /assets/map.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/assets/map.png -------------------------------------------------------------------------------- /assets/optical.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/assets/optical.jpg -------------------------------------------------------------------------------- /assets/pcbinsp.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/assets/pcbinsp.jpg -------------------------------------------------------------------------------- /assets/pcbscreenshot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/assets/pcbscreenshot.png -------------------------------------------------------------------------------- /assets/precision.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/assets/precision.png -------------------------------------------------------------------------------- /assets/recall.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/clintonoduor/PCB-Defect-Detection-using-Deepstream/1ab1450c776f25943f22b0b822e96145c227b3db/assets/recall.png --------------------------------------------------------------------------------