├── yolov5 ├── model.cpp ├── utils │ ├── __init__.py │ ├── __pycache__ │ │ ├── loss.cpython-37.pyc │ │ ├── plots.cpython-37.pyc │ │ ├── __init__.cpython-36.pyc │ │ ├── __init__.cpython-37.pyc │ │ ├── datasets.cpython-37.pyc │ │ ├── general.cpython-37.pyc │ │ ├── metrics.cpython-37.pyc │ │ ├── autoanchor.cpython-37.pyc │ │ ├── experimental.cpython-37.pyc │ │ ├── google_utils.cpython-37.pyc │ │ ├── torch_utils.cpython-36.pyc │ │ └── torch_utils.cpython-37.pyc │ ├── google_app_engine │ │ ├── additional_requirements.txt │ │ ├── app.yaml │ │ └── Dockerfile │ ├── activations.py │ ├── google_utils.py │ ├── autoanchor.py │ ├── metrics.py │ ├── loss.py │ ├── torch_utils.py │ ├── plots.py │ └── general.py ├── gen_wts.py ├── yolov5.h ├── Loader.h ├── CMakeLists.txt ├── simple_camera.cpp ├── utils.h ├── Loader.cpp ├── yololayer.h ├── yolov5.cpp ├── yololayer.cu ├── common.hpp └── logging.h ├── example.gif ├── sort ├── include │ ├── utils.h │ ├── track.h │ ├── tracker.h │ ├── kalman_filter.h │ ├── matrix.h │ ├── matrix.cpp │ └── munkres.h ├── CMakeLists.txt └── src │ ├── munkres.cpp │ ├── kalman_filter.cpp │ ├── track.cpp │ └── tracker.cpp ├── .gitignore ├── LICENSE ├── CMakeLists.txt ├── README.md └── Track.cpp /yolov5/model.cpp: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /yolov5/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /example.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/example.gif -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/loss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/loss.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/plots.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/plots.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/datasets.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/datasets.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/general.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/general.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/metrics.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/metrics.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/autoanchor.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/autoanchor.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/experimental.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/experimental.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/google_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/google_utils.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/torch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/torch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/torch_utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/SanftMonster/tensorrt_yolov5_tracker/HEAD/yolov5/utils/__pycache__/torch_utils.cpython-37.pyc -------------------------------------------------------------------------------- /yolov5/utils/google_app_engine/additional_requirements.txt: -------------------------------------------------------------------------------- 1 | # add these requirements in your app on top of the existing ones 2 | pip==18.1 3 | Flask==1.0.2 4 | gunicorn==19.9.0 5 | -------------------------------------------------------------------------------- /sort/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | constexpr int kNumColors = 32; 4 | 5 | constexpr int kMaxCoastCycles = 1; 6 | 7 | constexpr int kMinHits = 3; 8 | 9 | // Set threshold to 0 to accept all detections 10 | constexpr float kMinConfidence = 0.6; -------------------------------------------------------------------------------- /yolov5/utils/google_app_engine/app.yaml: -------------------------------------------------------------------------------- 1 | runtime: custom 2 | env: flex 3 | 4 | service: yolov5app 5 | 6 | liveness_check: 7 | initial_delay_sec: 600 8 | 9 | manual_scaling: 10 | instances: 1 11 | resources: 12 | cpu: 1 13 | memory_gb: 4 14 | disk_size_gb: 20 -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.gif 2 | # Prerequisites 3 | *.d 4 | 5 | # Compiled Object files 6 | *.slo 7 | *.lo 8 | *.o 9 | *.obj 10 | 11 | # Precompiled Headers 12 | *.gch 13 | *.pch 14 | 15 | # Compiled Dynamic libraries 16 | *.so 17 | *.dylib 18 | *.dll 19 | 20 | # Fortran module files 21 | *.mod 22 | *.smod 23 | 24 | # Compiled Static libraries 25 | *.lai 26 | *.la 27 | *.a 28 | *.lib 29 | 30 | # Executables 31 | *.exe 32 | *.out 33 | *.app 34 | -------------------------------------------------------------------------------- /sort/include/track.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include "kalman_filter.h" 5 | 6 | class Track { 7 | public: 8 | // Constructor 9 | Track(); 10 | 11 | // Destructor 12 | ~Track() = default; 13 | 14 | void Init(const cv::Rect& bbox); 15 | void Predict(); 16 | void Update(const cv::Rect& bbox); 17 | cv::Rect GetStateAsBbox() const; 18 | float GetNIS() const; 19 | 20 | int coast_cycles_ = 0, hit_streak_ = 0; 21 | 22 | private: 23 | Eigen::VectorXd ConvertBboxToObservation(const cv::Rect& bbox) const; 24 | cv::Rect ConvertStateToBbox(const Eigen::VectorXd &state) const; 25 | 26 | KalmanFilter kf_; 27 | }; 28 | -------------------------------------------------------------------------------- /yolov5/utils/google_app_engine/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM gcr.io/google-appengine/python 2 | 3 | # Create a virtualenv for dependencies. This isolates these packages from 4 | # system-level packages. 5 | # Use -p python3 or -p python3.7 to select python version. Default is version 2. 6 | RUN virtualenv /env -p python3 7 | 8 | # Setting these environment variables are the same as running 9 | # source /env/bin/activate. 10 | ENV VIRTUAL_ENV /env 11 | ENV PATH /env/bin:$PATH 12 | 13 | RUN apt-get update && apt-get install -y python-opencv 14 | 15 | # Copy the application's requirements.txt and run pip to install all 16 | # dependencies into the virtualenv. 17 | ADD requirements.txt /app/requirements.txt 18 | RUN pip install -r /app/requirements.txt 19 | 20 | # Add the application source code. 21 | ADD . /app 22 | 23 | # Run a WSGI server to serve the application. gunicorn must be declared as 24 | # a dependency in requirements.txt. 25 | CMD gunicorn -b :$PORT main:app 26 | -------------------------------------------------------------------------------- /sort/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | project(sort_cpp) 3 | 4 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall") 5 | set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/bin) 6 | 7 | # Try to find OpenCV 8 | find_package(OpenCV REQUIRED) 9 | if (OpenCV_FOUND) 10 | # If the package has been found, several variables will 11 | # be set, you can find the full list with descriptions 12 | # in the OpenCVConfig.cmake file. 13 | # Print some message showing some of them 14 | message(STATUS "OpenCV library status:") 15 | message(STATUS " version: ${OpenCV_VERSION}") 16 | message(STATUS " include path: ${OpenCV_INCLUDE_DIRS}") 17 | else () 18 | message(FATAL_ERROR "Could not locate OpenCV") 19 | endif() 20 | 21 | include_directories(${PROJECT_SOURCE_DIR}/include) 22 | 23 | file(GLOB SOURCE_FILES src/*.cpp) 24 | 25 | add_library(sort_cpp ${SOURCE_FILES}) 26 | 27 | target_link_libraries ( 28 | sort_cpp 29 | ${OpenCV_LIBS} 30 | ) 31 | -------------------------------------------------------------------------------- /yolov5/gen_wts.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import struct 3 | from utils.torch_utils import select_device 4 | import argparse 5 | 6 | if __name__ == '__main__': 7 | args = argparse.ArgumentParser() 8 | args.add_argument('--input',type=str,required=True,help='The path of input model.') 9 | args.add_argument('--output',type=str,required=True,help='The path you want to out put the model.') 10 | opt = args.parse_args() 11 | 12 | # Initialize 13 | device = select_device('cpu') 14 | # Load model 15 | model = torch.load(opt.input, map_location=device) 16 | model = model['model'].float() # load to FP32 17 | model.to(device).eval() 18 | 19 | f = open(opt.output, 'w') 20 | f.write('{}\n'.format(len(model.state_dict().keys()))) 21 | for k, v in model.state_dict().items(): 22 | vr = v.reshape(-1).cpu().numpy() 23 | f.write('{} {} '.format(k, len(vr))) 24 | for vv in vr: 25 | f.write(' ') 26 | f.write(struct.pack('>f', float(vv)).hex()) 27 | f.write('\n') 28 | -------------------------------------------------------------------------------- /sort/src/munkres.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2007 John Weaver 3 | * Copyright (c) 2015 Miroslav Krajicek 4 | * 5 | * This program is free software; you can redistribute it and/or modify 6 | * it under the terms of the GNU General Public License as published by 7 | * the Free Software Foundation; either version 2 of the License, or 8 | * (at your option) any later version. 9 | * 10 | * This program is distributed in the hope that it will be useful, 11 | * but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 | * GNU General Public License for more details. 14 | * 15 | * You should have received a copy of the GNU General Public License 16 | * along with this program; if not, write to the Free Software 17 | * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 18 | */ 19 | 20 | #include "munkres.h" 21 | 22 | template class Munkres; 23 | template class Munkres; 24 | template class Munkres; 25 | 26 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 AsakusaRinne 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 | -------------------------------------------------------------------------------- /yolov5/yolov5.h: -------------------------------------------------------------------------------- 1 | #ifndef TRACKCOUNTER_YOLOV5_H_ 2 | #define TRACKCOUNTER_YOLOV5_H_ 3 | 4 | #include 5 | #include "yololayer.h" 6 | #include "NvInfer.h" 7 | 8 | #define USE_FP16 // comment out this if want to use FP32 9 | #define DEVICE 0 // GPU id 10 | #define NMS_THRESH 0.3 11 | #define CONF_THRESH 0.45 12 | #define BATCH_SIZE 16 13 | #define WEIGHT "../yolov5/weights/yolov5.wts" 14 | 15 | #define NET s // s m l x 16 | #define NETSTRUCT(str) createEngine_##str 17 | #define CREATENET(net) NETSTRUCT(net) 18 | #define STR1(x) #x 19 | #define STR2(x) STR1(x) 20 | 21 | #endif 22 | 23 | using namespace nvinfer1; 24 | 25 | int yolov5_build_engine(std::string weightFile, std::string engineFile, float depth_multiple, float width_multiple); 26 | 27 | class yolov5_detector { 28 | private: 29 | IExecutionContext* context; 30 | cudaStream_t stream; 31 | void* buffers[2]; 32 | static float prob[BATCH_SIZE * (Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1)]; 33 | ICudaEngine* engine; 34 | IRuntime* runtime; 35 | 36 | 37 | public: 38 | yolov5_detector(std::string engine_name); 39 | std::vector> detect(float* data); 40 | void release(); 41 | }; 42 | -------------------------------------------------------------------------------- /yolov5/Loader.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #ifndef YOLOV5_LOADER_H_ 3 | #define YOLOV5_LOADER_H_ 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include "yolov5.h" 9 | 10 | #endif 11 | 12 | #define FPS 30 13 | #define INPUT_H Yolo::INPUT_H 14 | #define INPUT_W Yolo::INPUT_W 15 | 16 | std::string gstreamer_pipeline(int capture_width, int capture_height, int display_width, int display_height, int framerate, int flip_method); 17 | void preProcessImg(cv::Mat* img, float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]); 18 | cv::Mat preprocess_img(cv::Mat& img); 19 | 20 | class LoadVideo { 21 | public: 22 | 23 | cv::VideoCapture cap; 24 | int index = 0; 25 | int maxIndex; 26 | 27 | LoadVideo(std::string filename); 28 | 29 | float* getBatch(); 30 | 31 | int getLoops(); 32 | 33 | void release() { 34 | this->cap.release(); 35 | } 36 | }; 37 | 38 | class WebCam { 39 | private: 40 | static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; 41 | cv::VideoCapture cap; 42 | public: 43 | 44 | WebCam(int port); 45 | 46 | float* getBatch(); 47 | 48 | cv::Mat getFrame(); 49 | 50 | void release() { 51 | this->cap.release(); 52 | } 53 | }; -------------------------------------------------------------------------------- /yolov5/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.6) 2 | 3 | add_definitions(-std=c++11) 4 | 5 | option(CUDA_USE_STATIC_CUDA_RUNTIME OFF) 6 | set(CMAKE_CXX_STANDARD 11) 7 | set(CMAKE_BUILD_TYPE Debug) 8 | 9 | find_package(CUDA REQUIRED) 10 | 11 | include_directories(${PROJECT_SOURCE_DIR}/../include) 12 | # include and link dirs of cuda and tensorrt, you need adapt them if yours are different 13 | # cuda 14 | include_directories(/usr/local/cuda/include) 15 | link_directories(/usr/local/cuda/lib64) 16 | # tensorrt 17 | include_directories(/usr/include/x86_64-linux-gnu/) 18 | link_directories(/usr/lib/x86_64-linux-gnu/) 19 | 20 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") 21 | 22 | cuda_add_library(myplugins SHARED yololayer.cu) 23 | target_link_libraries(myplugins nvinfer cudart) 24 | 25 | find_package(OpenCV) 26 | include_directories(OpenCV_INCLUDE_DIRS) 27 | 28 | add_library(yolov5 yolov5.cpp yolov5.h Loader.cpp) 29 | 30 | target_link_libraries(yolov5 nvinfer) 31 | target_link_libraries(yolov5 cudart) 32 | target_link_libraries(yolov5 myplugins) 33 | target_link_libraries(yolov5 ${OpenCV_LIBS}) 34 | 35 | add_definitions(-O2 -pthread) 36 | 37 | 38 | 39 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # CMakeList.txt: 顶层 CMake 项目文件，在此处执行全局配置 2 | # 并包含子项目。 3 | # 4 | cmake_minimum_required (VERSION 2.6) 5 | 6 | add_definitions(-std=c++11) 7 | 8 | project (TrackCounter) 9 | # 将当前目录所有源文件保存在变量DIR_SRCS中 10 | aux_source_directory(. DIR_SRCS) 11 | 12 | # 包含子项目。 13 | add_subdirectory ("yolov5") 14 | add_subdirectory ("sort") 15 | 16 | find_package(CUDA REQUIRED) 17 | # cuda 18 | include_directories(/usr/local/cuda/include) 19 | link_directories(/usr/local/cuda/lib64) 20 | # tensorrt 21 | include_directories(/usr/include/x86_64-linux-gnu/) 22 | link_directories(/usr/lib/x86_64-linux-gnu/) 23 | 24 | 25 | add_executable(TrackCounter ${DIR_SRCS} ) 26 | # 添加链接库，MyMath是子目录中的链接库名 27 | target_link_libraries(TrackCounter yolov5) 28 | target_link_libraries(TrackCounter sort_cpp) 29 | 30 | find_package(CURL) 31 | if(CURL_FOUND) 32 | include_directories(${CURL_INCLUDE_DIR}) 33 | target_link_libraries(TrackCounter ${CURL_LIBRARY}) 34 | else(CURL_FOUND) 35 | message(FATAL_ERROR "CURL library not found") 36 | endif(CURL_FOUND) 37 | 38 | include_directories(${PROJECT_SOURCE_DIR}/include) 39 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -pthread -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED") 40 | 41 | set(CMAKE_CXX_STANDARD 11) 42 | set(CMAKE_BUILD_TYPE Debug) 43 | 44 | 45 | add_definitions(-O2 -pthread) 46 | 47 | 48 | -------------------------------------------------------------------------------- /sort/include/tracker.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | #include "track.h" 7 | #include "munkres.h" 8 | #include "utils.h" 9 | 10 | class Tracker { 11 | public: 12 | Tracker(); 13 | ~Tracker() = default; 14 | 15 | static float CalculateIou(const cv::Rect& det, const Track& track); 16 | 17 | static void HungarianMatching(const std::vector>& iou_matrix, 18 | size_t nrows, size_t ncols, 19 | std::vector>& association); 20 | 21 | /** 22 | * Assigns detections to tracked object (both represented as bounding boxes) 23 | * Returns 2 lists of matches, unmatched_detections 24 | * @param detection 25 | * @param tracks 26 | * @param matched 27 | * @param unmatched_det 28 | * @param iou_threshold 29 | */ 30 | static void AssociateDetectionsToTrackers(const std::vector& detection, 31 | std::map& tracks, 32 | std::map& matched, 33 | std::vector& unmatched_det, 34 | float iou_threshold = 0.1); 35 | 36 | void Run(const std::vector& detections); 37 | 38 | std::map GetTracks(); 39 | 40 | private: 41 | // Hash-map between ID and corresponding tracker 42 | std::map tracks_; 43 | 44 | // Assigned ID for each bounding box 45 | int id_; 46 | }; -------------------------------------------------------------------------------- /sort/include/kalman_filter.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | // abstract class for Kalman filter 7 | // implementation could be KF/EKF/UKF... 8 | class KalmanFilter { 9 | public: 10 | /** 11 | * user need to define H matrix & R matrix 12 | * @param num_states 13 | * @param num_obs 14 | */ 15 | // constructor 16 | explicit KalmanFilter(unsigned int num_states, unsigned int num_obs); 17 | 18 | // destructor 19 | virtual ~KalmanFilter() = default; 20 | 21 | /** 22 | * Coast state and state covariance using the process model 23 | * User can use this function without change the internal 24 | * tracking state x_ 25 | */ 26 | virtual void Coast(); 27 | 28 | /** 29 | * Predict without measurement update 30 | */ 31 | void Predict(); 32 | 33 | /** 34 | * This function maps the true state space into the observed space 35 | * using the observation model 36 | * User can implement their own method for more complicated models 37 | */ 38 | virtual Eigen::VectorXd PredictionToObservation(const Eigen::VectorXd &state); 39 | 40 | /** 41 | * Updates the state by using Extended Kalman Filter equations 42 | * @param z The measurement at k+1 43 | */ 44 | virtual void Update(const Eigen::VectorXd &z); 45 | 46 | /** 47 | * Calculate marginal log-likelihood to evaluate different parameter choices 48 | */ 49 | float CalculateLogLikelihood(const Eigen::VectorXd& y, const Eigen::MatrixXd& S); 50 | 51 | // State vector 52 | Eigen::VectorXd x_, x_predict_; 53 | 54 | // Error covariance matrix 55 | Eigen::MatrixXd P_, P_predict_; 56 | 57 | // State transition matrix 58 | Eigen::MatrixXd F_; 59 | 60 | // Covariance matrix of process noise 61 | Eigen::MatrixXd Q_; 62 | 63 | // measurement matrix 64 | Eigen::MatrixXd H_; 65 | 66 | // covariance matrix of observation noise 67 | Eigen::MatrixXd R_; 68 | 69 | unsigned int num_states_, num_obs_; 70 | 71 | float log_likelihood_delta_; 72 | 73 | float NIS_; 74 | }; -------------------------------------------------------------------------------- /yolov5/simple_camera.cpp: -------------------------------------------------------------------------------- 1 | // simple_camera.cpp 2 | // MIT License 3 | // Copyright (c) 2019 JetsonHacks 4 | // See LICENSE for OpenCV license and additional information 5 | // Using a CSI camera (such as the Raspberry Pi Version 2) connected to a 6 | // NVIDIA Jetson Nano Developer Kit using OpenCV 7 | // Drivers for the camera and OpenCV are included in the base image 8 | 9 | // #include 10 | #include 11 | // #include 12 | // #include 13 | 14 | std::string gstreamer_pipeline (int capture_width, int capture_height, int display_width, int display_height, int framerate, int flip_method) { 15 | return "nvarguscamerasrc ! video/x-raw(memory:NVMM), width=(int)" + std::to_string(capture_width) + ", height=(int)" + 16 | std::to_string(capture_height) + ", format=(string)NV12, framerate=(fraction)" + std::to_string(framerate) + 17 | "/1 ! nvvidconv flip-method=" + std::to_string(flip_method) + " ! video/x-raw, width=(int)" + std::to_string(display_width) + ", height=(int)" + 18 | std::to_string(display_height) + ", format=(string)BGRx ! videoconvert ! video/x-raw, format=(string)BGR ! appsink"; 19 | } 20 | 21 | int main() 22 | { 23 | int capture_width = 1280 ; 24 | int capture_height = 720 ; 25 | int display_width = 1280 ; 26 | int display_height = 720 ; 27 | int framerate = 60 ; 28 | int flip_method = 0 ; 29 | 30 | std::string pipeline = gstreamer_pipeline(capture_width, 31 | capture_height, 32 | display_width, 33 | display_height, 34 | framerate, 35 | flip_method); 36 | std::cout << "Using pipeline: \n\t" << pipeline << "\n"; 37 | 38 | cv::VideoCapture cap(pipeline, cv::CAP_GSTREAMER); 39 | if(!cap.isOpened()) { 40 | std::cout<<"Failed to open camera."< 23 | #include 24 | #include 25 | 26 | template 27 | class Matrix { 28 | public: 29 | Matrix(); 30 | Matrix(const size_t rows, const size_t columns); 31 | Matrix(const std::initializer_list> init); 32 | Matrix(const Matrix &other); 33 | Matrix & operator= (const Matrix &other); 34 | ~Matrix(); 35 | // all operations modify the matrix in-place. 36 | void resize(const size_t rows, const size_t columns, const T default_value = 0); 37 | void clear(); 38 | T& operator () (const size_t x, const size_t y); 39 | const T& operator () (const size_t x, const size_t y) const; 40 | const T min() const; 41 | const T max() const; 42 | inline size_t minsize() { return ((m_rows < m_columns) ? m_rows : m_columns); } 43 | inline size_t columns() const { return m_columns;} 44 | inline size_t rows() const { return m_rows;} 45 | 46 | friend std::ostream& operator<<(std::ostream& os, const Matrix &matrix) 47 | { 48 | os << "Matrix:" << std::endl; 49 | for (size_t row = 0 ; row < matrix.rows() ; row++ ) 50 | { 51 | for (size_t col = 0 ; col < matrix.columns() ; col++ ) 52 | { 53 | os.width(8); 54 | os << matrix(row, col) << ","; 55 | } 56 | os << std::endl; 57 | } 58 | return os; 59 | } 60 | 61 | private: 62 | T **m_matrix; 63 | size_t m_rows; 64 | size_t m_columns; 65 | }; 66 | 67 | #ifndef USE_EXPORT_KEYWORD 68 | #include "matrix.cpp" 69 | //#define export /*export*/ 70 | #endif 71 | 72 | #endif /* !defined(_MATRIX_H_) */ 73 | 74 | -------------------------------------------------------------------------------- /yolov5/utils/activations.py: -------------------------------------------------------------------------------- 1 | # Activation functions 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | # SiLU https://arxiv.org/pdf/1905.02244.pdf ---------------------------------------------------------------------------- 9 | class SiLU(nn.Module): # export-friendly version of nn.SiLU() 10 | @staticmethod 11 | def forward(x): 12 | return x * torch.sigmoid(x) 13 | 14 | 15 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish() 16 | @staticmethod 17 | def forward(x): 18 | # return x * F.hardsigmoid(x) # for torchscript and CoreML 19 | return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX 20 | 21 | 22 | class MemoryEfficientSwish(nn.Module): 23 | class F(torch.autograd.Function): 24 | @staticmethod 25 | def forward(ctx, x): 26 | ctx.save_for_backward(x) 27 | return x * torch.sigmoid(x) 28 | 29 | @staticmethod 30 | def backward(ctx, grad_output): 31 | x = ctx.saved_tensors[0] 32 | sx = torch.sigmoid(x) 33 | return grad_output * (sx * (1 + x * (1 - sx))) 34 | 35 | def forward(self, x): 36 | return self.F.apply(x) 37 | 38 | 39 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 40 | class Mish(nn.Module): 41 | @staticmethod 42 | def forward(x): 43 | return x * F.softplus(x).tanh() 44 | 45 | 46 | class MemoryEfficientMish(nn.Module): 47 | class F(torch.autograd.Function): 48 | @staticmethod 49 | def forward(ctx, x): 50 | ctx.save_for_backward(x) 51 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 52 | 53 | @staticmethod 54 | def backward(ctx, grad_output): 55 | x = ctx.saved_tensors[0] 56 | sx = torch.sigmoid(x) 57 | fx = F.softplus(x).tanh() 58 | return grad_output * (fx + x * sx * (1 - fx * fx)) 59 | 60 | def forward(self, x): 61 | return self.F.apply(x) 62 | 63 | 64 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 65 | class FReLU(nn.Module): 66 | def __init__(self, c1, k=3): # ch_in, kernel 67 | super().__init__() 68 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False) 69 | self.bn = nn.BatchNorm2d(c1) 70 | 71 | def forward(self, x): 72 | return torch.max(x, self.bn(self.conv(x))) 73 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # tensorrt_yolov5_tracker 2 | 3 | This is a project to deploy object tracking algorithm with yolov5 and TensorRT. Sort and Deep-sort algorithm are used to track the objects. 4 | Thanks for the contribution of [tensorrtx](https://github.com/wang-xinyu/tensorrtx) and [sort-cpp](https://github.com/yasenh/sort-cpp)! 5 | Now this project implements the tracki algorithm with yolov5 and sort(c++ version). **Deep-sort, multi-thread acceleration and python implemention** will be updated soon. Please give a star if this project helps ^_^ 6 | 7 | ![example result of object tracking](https://github.com/AsakusaRinne/tensorrt_yolov5_tracker/blob/main/example.gif) 8 | ## Table of Contents 9 | 10 | - [ToDo](#ToDo) 11 | - [Install](#install) 12 | - [Usage](#usage) 13 | - [Reference](#Reference) 14 | - [License](#license) 15 | 16 | ## ToDo 17 | ***1.Deep-sort tracker.*** 18 | 19 | ***2.Multi-thread acceleration for c++.*** 20 | 21 | ***3.Python implemention.*** 22 | 23 | 24 | ## Install 25 | 26 | ``` 27 | $ git https://github.com/AsakusaRinne/tensorrt_yolov5_tracker.git 28 | $ cd tensorrt_yolov5_tracker 29 | $ mkdir build 30 | $ cd build 31 | $ cmake .. 32 | $ make 33 | ``` 34 | 35 | ## Usage 36 | You could change the batch size in ``yolov5/yolov5.h`` 37 | You could change the input image size in ``yolov5/yololayer.h`` 38 | 39 | We provide the support for building models of yolov5-3.1 and yolov5-4.0 in our project. The default branch is corresponding to yolov5-4.0. Please change to ``c269f65`` if you want to use yolov5-3.1 models to build engines. 40 | 41 | If other versions are wanted, please kindly refer to [tensorrtx](https://github.com/wang-xinyu/tensorrtx) to build the engine and then copy to the folder of this project. 42 | 43 | 44 | ### Build engine 45 | At first, please put the ``yolov5/gen_wts.py`` in the folder of corresponding version of [**Yolov5**](https://github.com/ultralytics/yolov5). For example, if you use the model of yolov5-4.0, please download the release of yolov5-4.0 and put the ``gen_wts.py`` in its folder. Then run the code to convert model from ``.pt`` to ``.wts``. Then please use the following instructions to build your engines. 46 | ``` 47 | $ sudo ./Tracker -s [.wts filename] [.engine filename] [s, m, l, or x] # build yolov5 model of [s, m, l, or x] 48 | 49 | $ sudo ./Tracker -s [.wts filename] [.engine filename] [depth] [width] # build yolov5 model of scales which are not s, m, l, or x 50 | **for example** 51 | $ sudo ./Tracker -s ../yolov5/weights/yolov5.wts ../engines/yolov5.engine 0.17 0.25 52 | ``` 53 | ### Process video 54 | ``` 55 | $ sudo ./Tracker -v [video filename] [engine filename] 56 | ``` 57 | The output video will be in the folder ``output`` 58 | Note that the track time includes the time of drawing boxes on images and saving them to the disk, which makes it a little long. You can delete the codes for drawing and saving if they are not needed to accelerate the process. 59 | 60 | ## Reference 61 | [**tensorrtx**](https://github.com/wang-xinyu/tensorrtx) 62 | 63 | [**Yolov5**](https://github.com/ultralytics/yolov5) 64 | 65 | [**sort-cpp**](https://github.com/yasenh/sort-cpp) 66 | 67 | ## License 68 | [MIT license](https://mit-license.org/) 69 | -------------------------------------------------------------------------------- /yolov5/utils.h: -------------------------------------------------------------------------------- 1 | #ifndef __TRT_UTILS_H_ 2 | #define __TRT_UTILS_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | #ifndef CUDA_CHECK 10 | 11 | #define CUDA_CHECK(callstr) \ 12 | { \ 13 | cudaError_t error_code = callstr; \ 14 | if (error_code != cudaSuccess) { \ 15 | std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ 16 | assert(0); \ 17 | } \ 18 | } 19 | 20 | #endif 21 | 22 | namespace Tn 23 | { 24 | class Profiler : public nvinfer1::IProfiler 25 | { 26 | public: 27 | void printLayerTimes(int itrationsTimes) 28 | { 29 | float totalTime = 0; 30 | for (size_t i = 0; i < mProfile.size(); i++) 31 | { 32 | printf("%-40.40s %4.3fms\n", mProfile[i].first.c_str(), mProfile[i].second / itrationsTimes); 33 | totalTime += mProfile[i].second; 34 | } 35 | printf("Time over all layers: %4.3f\n", totalTime / itrationsTimes); 36 | } 37 | private: 38 | typedef std::pair Record; 39 | std::vector mProfile; 40 | 41 | virtual void reportLayerTime(const char* layerName, float ms) 42 | { 43 | auto record = std::find_if(mProfile.begin(), mProfile.end(), [&](const Record& r){ return r.first == layerName; }); 44 | if (record == mProfile.end()) 45 | mProfile.push_back(std::make_pair(layerName, ms)); 46 | else 47 | record->second += ms; 48 | } 49 | }; 50 | 51 | //Logger for TensorRT info/warning/errors 52 | class Logger : public nvinfer1::ILogger 53 | { 54 | public: 55 | 56 | Logger(): Logger(Severity::kWARNING) {} 57 | 58 | Logger(Severity severity): reportableSeverity(severity) {} 59 | 60 | void log(Severity severity, const char* msg) override 61 | { 62 | // suppress messages with severity enum value greater than the reportable 63 | if (severity > reportableSeverity) return; 64 | 65 | switch (severity) 66 | { 67 | case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break; 68 | case Severity::kERROR: std::cerr << "ERROR: "; break; 69 | case Severity::kWARNING: std::cerr << "WARNING: "; break; 70 | case Severity::kINFO: std::cerr << "INFO: "; break; 71 | default: std::cerr << "UNKNOWN: "; break; 72 | } 73 | std::cerr << msg << std::endl; 74 | } 75 | 76 | Severity reportableSeverity{Severity::kWARNING}; 77 | }; 78 | 79 | template 80 | void write(char*& buffer, const T& val) 81 | { 82 | *reinterpret_cast(buffer) = val; 83 | buffer += sizeof(T); 84 | } 85 | 86 | template 87 | void read(const char*& buffer, T& val) 88 | { 89 | val = *reinterpret_cast(buffer); 90 | buffer += sizeof(T); 91 | } 92 | } 93 | 94 | #endif -------------------------------------------------------------------------------- /sort/src/kalman_filter.cpp: -------------------------------------------------------------------------------- 1 | #include "kalman_filter.h" 2 | 3 | 4 | KalmanFilter::KalmanFilter(unsigned int num_states, unsigned int num_obs) : 5 | num_states_(num_states), num_obs_(num_obs) { 6 | /*** Predict ***/ 7 | // State vector 8 | x_ = Eigen::VectorXd::Zero(num_states); 9 | // Predicted(a prior) state vector 10 | x_predict_ = Eigen::VectorXd::Zero(num_states); 11 | 12 | // State transition matrix F_ 13 | F_ = Eigen::MatrixXd::Zero(num_states, num_states); 14 | 15 | // Error covariance matrix P 16 | P_ = Eigen::MatrixXd::Zero(num_states, num_states); 17 | // Predicted(a prior) error covariance matrix 18 | P_predict_ = Eigen::MatrixXd::Zero(num_states, num_states); 19 | 20 | // Covariance matrix of process noise 21 | Q_ = Eigen::MatrixXd::Zero(num_states, num_states); 22 | 23 | /*** Update ***/ 24 | // Observation matrix 25 | H_ = Eigen::MatrixXd::Zero(num_obs, num_states); 26 | 27 | // Covariance matrix of observation noise 28 | R_ = Eigen::MatrixXd::Zero(num_obs, num_obs); 29 | 30 | log_likelihood_delta_ = 0.0; 31 | NIS_ = 0.0; 32 | } 33 | 34 | 35 | void KalmanFilter::Coast() { 36 | x_predict_ = F_ * x_; 37 | P_predict_ = F_ * P_ * F_.transpose() + Q_; 38 | } 39 | 40 | 41 | void KalmanFilter::Predict() { 42 | Coast(); 43 | x_ = x_predict_; 44 | P_ = P_predict_; 45 | } 46 | 47 | 48 | Eigen::VectorXd KalmanFilter::PredictionToObservation(const Eigen::VectorXd &state) { 49 | return (H_*state); 50 | } 51 | 52 | 53 | void KalmanFilter::Update(const Eigen::VectorXd& z) { 54 | Eigen::VectorXd z_predict = PredictionToObservation(x_predict_); 55 | 56 | // y - innovation, z - real observation, z_predict - predicted observation 57 | Eigen::VectorXd y = z - z_predict; 58 | 59 | Eigen::MatrixXd Ht = H_.transpose(); 60 | 61 | // S - innovation covariance 62 | Eigen::MatrixXd S = H_ * P_predict_ * Ht + R_; 63 | 64 | NIS_ = y.transpose() * S.inverse() * y; 65 | 66 | // std::cout << std::endl; 67 | // std::cout << "P_predict = " << std::endl; 68 | // std::cout << P_predict_ << std::endl; 69 | // 70 | // 71 | // std::cout << "Z = " << std::endl; 72 | // std::cout << z << std::endl; 73 | // 74 | // std::cout << "Z_pred = " << std::endl; 75 | // std::cout << z_predict << std::endl; 76 | // 77 | // std::cout << "y = " << std::endl; 78 | // std::cout << y << std::endl; 79 | // 80 | // std::cout << "S = " << std::endl; 81 | // std::cout << S << std::endl; 82 | // 83 | // std::cout << "NIS = " << NIS_ << std::endl; 84 | 85 | 86 | // K - Kalman gain 87 | Eigen::MatrixXd K = P_predict_ * Ht * S.inverse(); 88 | 89 | // Updated state estimation 90 | x_ = x_predict_ + K * y; 91 | 92 | Eigen::MatrixXd I = Eigen::MatrixXd::Identity(num_states_, num_states_); 93 | // Joseph form 94 | //P_ = (I - K * H_) * P_predict_ * (I - K * H_).transpose() + K * R_ * K.transpose(); 95 | // Optimal gain 96 | P_ = (I - K * H_) * P_predict_; 97 | } 98 | 99 | 100 | float KalmanFilter::CalculateLogLikelihood(const Eigen::VectorXd& y, const Eigen::MatrixXd& S) { 101 | float log_likelihood; 102 | 103 | // Note: Computing log(M.determinant()) in Eigen C++ is risky for large matrices since it may overflow or underflow. 104 | // compute the Cholesky decomposition of the innovation covariance matrix, because it is symmetric 105 | // S = L * L^T = U^T * U 106 | // then retrieve factor L in the decomposition 107 | auto& L = S.llt().matrixL(); 108 | 109 | // find log determinant of innovation covariance matrix 110 | float log_determinant = 0; 111 | for (unsigned int i = 0; i < S.rows(); i++) 112 | log_determinant += log(L(i, i)); 113 | log_determinant *= 2; 114 | 115 | // log-likelihood expression for current iteration 116 | log_likelihood = -0.5 * (y.transpose() * S.inverse() * y + num_obs_ * log(2 * M_PI) + log_determinant); 117 | 118 | if (std::isnan(log_likelihood)) { 119 | log_likelihood = -1e50; 120 | } 121 | 122 | return log_likelihood; 123 | } -------------------------------------------------------------------------------- /sort/src/track.cpp: -------------------------------------------------------------------------------- 1 | #include "track.h" 2 | 3 | 4 | Track::Track() : kf_(8, 4) { 5 | 6 | /*** Define constant velocity model ***/ 7 | // state - center_x, center_y, width, height, v_cx, v_cy, v_width, v_height 8 | kf_.F_ << 9 | 1, 0, 0, 0, 1, 0, 0, 0, 10 | 0, 1, 0, 0, 0, 1, 0, 0, 11 | 0, 0, 1, 0, 0, 0, 1, 0, 12 | 0, 0, 0, 1, 0, 0, 0, 1, 13 | 0, 0, 0, 0, 1, 0, 0, 0, 14 | 0, 0, 0, 0, 0, 1, 0, 0, 15 | 0, 0, 0, 0, 0, 0, 1, 0, 16 | 0, 0, 0, 0, 0, 0, 0, 1; 17 | 18 | // Give high uncertainty to the unobservable initial velocities 19 | kf_.P_ << 20 | 10, 0, 0, 0, 0, 0, 0, 0, 21 | 0, 10, 0, 0, 0, 0, 0, 0, 22 | 0, 0, 10, 0, 0, 0, 0, 0, 23 | 0, 0, 0, 10, 0, 0, 0, 0, 24 | 0, 0, 0, 0, 10000, 0, 0, 0, 25 | 0, 0, 0, 0, 0, 10000, 0, 0, 26 | 0, 0, 0, 0, 0, 0, 10000, 0, 27 | 0, 0, 0, 0, 0, 0, 0, 10000; 28 | 29 | 30 | kf_.H_ << 31 | 1, 0, 0, 0, 0, 0, 0, 0, 32 | 0, 1, 0, 0, 0, 0, 0, 0, 33 | 0, 0, 1, 0, 0, 0, 0, 0, 34 | 0, 0, 0, 1, 0, 0, 0, 0; 35 | 36 | kf_.Q_ << 37 | 1, 0, 0, 0, 0, 0, 0, 0, 38 | 0, 1, 0, 0, 0, 0, 0, 0, 39 | 0, 0, 1, 0, 0, 0, 0, 0, 40 | 0, 0, 0, 1, 0, 0, 0, 0, 41 | 0, 0, 0, 0, 0.01, 0, 0, 0, 42 | 0, 0, 0, 0, 0, 0.01, 0, 0, 43 | 0, 0, 0, 0, 0, 0, 0.0001, 0, 44 | 0, 0, 0, 0, 0, 0, 0, 0.0001; 45 | 46 | kf_.R_ << 47 | 1, 0, 0, 0, 48 | 0, 1, 0, 0, 49 | 0, 0, 10, 0, 50 | 0, 0, 0, 10; 51 | } 52 | 53 | 54 | // Get predicted locations from existing trackers 55 | // dt is time elapsed between the current and previous measurements 56 | void Track::Predict() { 57 | kf_.Predict(); 58 | 59 | // hit streak count will be reset 60 | if (coast_cycles_ > 0) { 61 | hit_streak_ = 0; 62 | } 63 | // accumulate coast cycle count 64 | coast_cycles_++; 65 | } 66 | 67 | 68 | // Update matched trackers with assigned detections 69 | void Track::Update(const cv::Rect& bbox) { 70 | 71 | // get measurement update, reset coast cycle count 72 | coast_cycles_ = 0; 73 | // accumulate hit streak count 74 | hit_streak_++; 75 | 76 | // observation - center_x, center_y, area, ratio 77 | Eigen::VectorXd observation = ConvertBboxToObservation(bbox); 78 | kf_.Update(observation); 79 | 80 | 81 | } 82 | 83 | 84 | // Create and initialize new trackers for unmatched detections, with initial bounding box 85 | void Track::Init(const cv::Rect &bbox) { 86 | kf_.x_.head(4) << ConvertBboxToObservation(bbox); 87 | hit_streak_++; 88 | } 89 | 90 | 91 | /** 92 | * Returns the current bounding box estimate 93 | * @return 94 | */ 95 | cv::Rect Track::GetStateAsBbox() const { 96 | return ConvertStateToBbox(kf_.x_); 97 | } 98 | 99 | 100 | float Track::GetNIS() const { 101 | return kf_.NIS_; 102 | } 103 | 104 | 105 | /** 106 | * Takes a bounding box in the form [x, y, width, height] and returns z in the form 107 | * [x, y, s, r] where x,y is the centre of the box and s is the scale/area and r is 108 | * the aspect ratio 109 | * 110 | * @param bbox 111 | * @return 112 | */ 113 | Eigen::VectorXd Track::ConvertBboxToObservation(const cv::Rect& bbox) const{ 114 | Eigen::VectorXd observation = Eigen::VectorXd::Zero(4); 115 | auto width = static_cast(bbox.width); 116 | auto height = static_cast(bbox.height); 117 | float center_x = bbox.x + width / 2; 118 | float center_y = bbox.y + height / 2; 119 | observation << center_x, center_y, width, height; 120 | return observation; 121 | } 122 | 123 | 124 | /** 125 | * Takes a bounding box in the centre form [x,y,s,r] and returns it in the form 126 | * [x1,y1,x2,y2] where x1,y1 is the top left and x2,y2 is the bottom right 127 | * 128 | * @param state 129 | * @return 130 | */ 131 | cv::Rect Track::ConvertStateToBbox(const Eigen::VectorXd &state) const { 132 | // state - center_x, center_y, width, height, v_cx, v_cy, v_width, v_height 133 | auto width = static_cast(state[2]); 134 | auto height = static_cast(state[3]); 135 | auto tl_x = static_cast(state[0] - width / 2.0); 136 | auto tl_y = static_cast(state[1] - height / 2.0); 137 | cv::Rect rect(cv::Point(tl_x, tl_y), cv::Size(width, height)); 138 | return rect; 139 | } -------------------------------------------------------------------------------- /yolov5/Loader.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "Loader.h" 3 | #include "yololayer.h" 4 | 5 | using namespace std; 6 | using namespace cv; 7 | 8 | std::string gstreamer_pipeline(int capture_width, int capture_height, int display_width, int display_height, int framerate, int flip_method) { 9 | return "nvarguscamerasrc ! video/x-raw(memory:NVMM), width=(int)" + std::to_string(capture_width) + ", height=(int)" + 10 | std::to_string(capture_height) + ", format=(string)NV12, framerate=(fraction)" + std::to_string(framerate) + 11 | "/1 ! nvvidconv flip-method=" + std::to_string(flip_method) + " ! video/x-raw, " + "format=(string)BGRx ! videoconvert ! video/x-raw, format=(string)BGR ! appsink"; 12 | } 13 | 14 | 15 | cv::Mat preprocess_img(cv::Mat& img) { 16 | int w, h, x, y; 17 | float r_w = INPUT_W / (img.cols*1.0); 18 | float r_h = INPUT_H / (img.rows*1.0); 19 | if (r_h > r_w) { 20 | w = INPUT_W; 21 | h = r_w * img.rows; 22 | x = 0; 23 | y = (INPUT_H - h) / 2; 24 | } else { 25 | w = r_h * img.cols; 26 | h = INPUT_H; 27 | x = (INPUT_W - w) / 2; 28 | y = 0; 29 | } 30 | cv::Mat re(h, w, CV_8UC3); 31 | cv::resize(img, re, re.size(), 0, 0, cv::INTER_LINEAR); 32 | cv::Mat out(INPUT_H, INPUT_W, CV_8UC3, cv::Scalar(128, 128, 128)); 33 | re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); 34 | return out; 35 | } 36 | 37 | 38 | LoadVideo::LoadVideo(string filename) { 39 | std::cout<<"Loading video: "<cap = cv::VideoCapture(pipeline, cv::CAP_GSTREAMER); 86 | 87 | //this->cap.open(port); 88 | cout << "01"<< endl; 89 | 90 | if (!this->cap.isOpened()) { 91 | std::cout << "Failed to open camera." << std::endl; 92 | } 93 | 94 | } 95 | 96 | cv::Mat WebCam::getFrame() { 97 | cv::Mat img; 98 | while (!this->cap.read(img)); 99 | return img; 100 | } 101 | 102 | float* WebCam::getBatch() { 103 | static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; 104 | cv::Mat img; 105 | for (int b = 0; b < BATCH_SIZE; b++) { 106 | if (cap.read(img)) { 107 | cv::Mat pr_img = preprocess_img(img); 108 | int i = 0; 109 | for (int row = 0; row < INPUT_H; ++row) { 110 | uchar* uc_pixel = pr_img.data + row * pr_img.step; 111 | for (int col = 0; col < INPUT_W; ++col) { 112 | data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0; 113 | data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0; 114 | data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0; 115 | uc_pixel += 3; 116 | ++i; 117 | } 118 | } 119 | } 120 | } 121 | return data; 122 | } 123 | 124 | 125 | 126 | void preProcessImg(cv::Mat img[BATCH_SIZE], float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]) { 127 | for (int b = 0; b < BATCH_SIZE; b++) { 128 | cv::Mat pr_img = preprocess_img(img[b]); 129 | int i = 0; 130 | for (int row = 0; row < INPUT_H; ++row) { 131 | uchar* uc_pixel = pr_img.data + row * pr_img.step; 132 | for (int col = 0; col < INPUT_W; ++col) { 133 | data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0; 134 | data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0; 135 | data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0; 136 | uc_pixel += 3; 137 | ++i; 138 | } 139 | } 140 | } 141 | } 142 | -------------------------------------------------------------------------------- /yolov5/yololayer.h: -------------------------------------------------------------------------------- 1 | #ifndef _YOLO_LAYER_H 2 | #define _YOLO_LAYER_H 3 | 4 | #include 5 | #include 6 | #include "NvInfer.h" 7 | 8 | namespace Yolo 9 | { 10 | static constexpr int CHECK_COUNT = 3; 11 | static constexpr float IGNORE_THRESH = 0.1f; 12 | struct YoloKernel 13 | { 14 | int width; 15 | int height; 16 | float anchors[CHECK_COUNT * 2]; 17 | }; 18 | static constexpr int MAX_OUTPUT_BBOX_COUNT = 100; 19 | static constexpr int CLASS_NUM = 1; 20 | static constexpr int INPUT_H = 480; 21 | static constexpr int INPUT_W = 640; 22 | static constexpr int ORIGINAL_H = 1080; 23 | static constexpr int ORIGINAL_W = 1920; 24 | 25 | static constexpr int LOCATIONS = 4; 26 | struct alignas(float) Detection { 27 | //center_x center_y w h 28 | float bbox[LOCATIONS]; 29 | float conf; // bbox_conf * cls_conf 30 | float class_id; 31 | }; 32 | } 33 | 34 | namespace nvinfer1 35 | { 36 | class YoloLayerPlugin : public IPluginV2IOExt 37 | { 38 | public: 39 | YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel); 40 | YoloLayerPlugin(const void* data, size_t length); 41 | ~YoloLayerPlugin(); 42 | 43 | int getNbOutputs() const override 44 | { 45 | return 1; 46 | } 47 | 48 | Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; 49 | 50 | int initialize() override; 51 | 52 | virtual void terminate() override {}; 53 | 54 | virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0; } 55 | 56 | virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override; 57 | 58 | virtual size_t getSerializationSize() const override; 59 | 60 | virtual void serialize(void* buffer) const override; 61 | 62 | bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override { 63 | return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT; 64 | } 65 | 66 | const char* getPluginType() const override; 67 | 68 | const char* getPluginVersion() const override; 69 | 70 | void destroy() override; 71 | 72 | IPluginV2IOExt* clone() const override; 73 | 74 | void setPluginNamespace(const char* pluginNamespace) override; 75 | 76 | const char* getPluginNamespace() const override; 77 | 78 | DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; 79 | 80 | bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override; 81 | 82 | bool canBroadcastInputAcrossBatch(int inputIndex) const override; 83 | 84 | void attachToContext( 85 | cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override; 86 | 87 | void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override; 88 | 89 | void detachFromContext() override; 90 | 91 | private: 92 | void forwardGpu(const float *const * inputs, float * output, cudaStream_t stream, int batchSize = 1); 93 | int mThreadCount = 256; 94 | const char* mPluginNamespace; 95 | int mKernelCount; 96 | int mClassCount; 97 | int mYoloV5NetWidth; 98 | int mYoloV5NetHeight; 99 | int mMaxOutObject; 100 | std::vector mYoloKernel; 101 | void** mAnchor; 102 | }; 103 | 104 | class YoloPluginCreator : public IPluginCreator 105 | { 106 | public: 107 | YoloPluginCreator(); 108 | 109 | ~YoloPluginCreator() override = default; 110 | 111 | const char* getPluginName() const override; 112 | 113 | const char* getPluginVersion() const override; 114 | 115 | const PluginFieldCollection* getFieldNames() override; 116 | 117 | IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override; 118 | 119 | IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; 120 | 121 | void setPluginNamespace(const char* libNamespace) override 122 | { 123 | mNamespace = libNamespace; 124 | } 125 | 126 | const char* getPluginNamespace() const override 127 | { 128 | return mNamespace.c_str(); 129 | } 130 | 131 | private: 132 | std::string mNamespace; 133 | static PluginFieldCollection mFC; 134 | static std::vector mPluginAttributes; 135 | }; 136 | REGISTER_TENSORRT_PLUGIN(YoloPluginCreator); 137 | }; 138 | 139 | #endif 140 | -------------------------------------------------------------------------------- /yolov5/utils/google_utils.py: -------------------------------------------------------------------------------- 1 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries 2 | 3 | import os 4 | import platform 5 | import subprocess 6 | import time 7 | from pathlib import Path 8 | 9 | import torch 10 | 11 | 12 | def gsutil_getsize(url=''): 13 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du 14 | s = subprocess.check_output('gsutil du %s' % url, shell=True).decode('utf-8') 15 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes 16 | 17 | 18 | def attempt_download(weights): 19 | # Attempt to download pretrained weights if not found locally 20 | weights = str(weights).strip().replace("'", '') 21 | file = Path(weights).name.lower() 22 | 23 | msg = weights + ' missing, try downloading from https://github.com/ultralytics/yolov5/releases/' 24 | models = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'] # available models 25 | redundant = False # offer second download option 26 | 27 | if file in models and not os.path.isfile(weights): 28 | # Google Drive 29 | # d = {'yolov5s.pt': '1R5T6rIyy3lLwgFXNms8whc-387H0tMQO', 30 | # 'yolov5m.pt': '1vobuEExpWQVpXExsJ2w-Mbf3HJjWkQJr', 31 | # 'yolov5l.pt': '1hrlqD1Wdei7UT4OgT785BEk1JwnSvNEV', 32 | # 'yolov5x.pt': '1mM8aZJlWTxOg7BZJvNUMrTnA2AbeCVzS'} 33 | # r = gdrive_download(id=d[file], name=weights) if file in d else 1 34 | # if r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6: # check 35 | # return 36 | 37 | try: # GitHub 38 | url = 'https://github.com/ultralytics/yolov5/releases/download/v3.1/' + file 39 | print('Downloading %s to %s...' % (url, weights)) 40 | torch.hub.download_url_to_file(url, weights) 41 | assert os.path.exists(weights) and os.path.getsize(weights) > 1E6 # check 42 | except Exception as e: # GCP 43 | print('Download error: %s' % e) 44 | assert redundant, 'No secondary mirror' 45 | url = 'https://storage.googleapis.com/ultralytics/yolov5/ckpt/' + file 46 | print('Downloading %s to %s...' % (url, weights)) 47 | r = os.system('curl -L %s -o %s' % (url, weights)) # torch.hub.download_url_to_file(url, weights) 48 | finally: 49 | if not (os.path.exists(weights) and os.path.getsize(weights) > 1E6): # check 50 | os.remove(weights) if os.path.exists(weights) else None # remove partial downloads 51 | print('ERROR: Download failure: %s' % msg) 52 | print('') 53 | return 54 | 55 | 56 | def gdrive_download(id='1uH2BylpFxHKEGXKL6wJJlsgMU2YEjxuc', name='tmp.zip'): 57 | # Downloads a file from Google Drive. from utils.google_utils import *; gdrive_download() 58 | t = time.time() 59 | 60 | print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='') 61 | os.remove(name) if os.path.exists(name) else None # remove existing 62 | os.remove('cookie') if os.path.exists('cookie') else None 63 | 64 | # Attempt file download 65 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 66 | os.system('curl -c ./cookie -s -L "drive.google.com/uc?export=download&id=%s" > %s ' % (id, out)) 67 | if os.path.exists('cookie'): # large file 68 | s = 'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm=%s&id=%s" -o %s' % (get_token(), id, name) 69 | else: # small file 70 | s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id) 71 | r = os.system(s) # execute, capture return 72 | os.remove('cookie') if os.path.exists('cookie') else None 73 | 74 | # Error check 75 | if r != 0: 76 | os.remove(name) if os.path.exists(name) else None # remove partial 77 | print('Download error ') # raise Exception('Download error') 78 | return r 79 | 80 | # Unzip if archive 81 | if name.endswith('.zip'): 82 | print('unzipping... ', end='') 83 | os.system('unzip -q %s' % name) # unzip 84 | os.remove(name) # remove zip to free space 85 | 86 | print('Done (%.1fs)' % (time.time() - t)) 87 | return r 88 | 89 | 90 | def get_token(cookie="./cookie"): 91 | with open(cookie) as f: 92 | for line in f: 93 | if "download" in line: 94 | return line.split()[-1] 95 | return "" 96 | 97 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 98 | # # Uploads a file to a bucket 99 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 100 | # 101 | # storage_client = storage.Client() 102 | # bucket = storage_client.get_bucket(bucket_name) 103 | # blob = bucket.blob(destination_blob_name) 104 | # 105 | # blob.upload_from_filename(source_file_name) 106 | # 107 | # print('File {} uploaded to {}.'.format( 108 | # source_file_name, 109 | # destination_blob_name)) 110 | # 111 | # 112 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 113 | # # Uploads a blob from a bucket 114 | # storage_client = storage.Client() 115 | # bucket = storage_client.get_bucket(bucket_name) 116 | # blob = bucket.blob(source_blob_name) 117 | # 118 | # blob.download_to_filename(destination_file_name) 119 | # 120 | # print('Blob {} downloaded to {}.'.format( 121 | # source_blob_name, 122 | # destination_file_name)) 123 | -------------------------------------------------------------------------------- /sort/src/tracker.cpp: -------------------------------------------------------------------------------- 1 | #include "tracker.h" 2 | 3 | 4 | Tracker::Tracker() { 5 | id_ = 0; 6 | } 7 | 8 | float Tracker::CalculateIou(const cv::Rect& det, const Track& track) { 9 | auto trk = track.GetStateAsBbox(); 10 | // get min/max points 11 | auto xx1 = std::max(det.tl().x, trk.tl().x); 12 | auto yy1 = std::max(det.tl().y, trk.tl().y); 13 | auto xx2 = std::min(det.br().x, trk.br().x); 14 | auto yy2 = std::min(det.br().y, trk.br().y); 15 | auto w = std::max(0, xx2 - xx1); 16 | auto h = std::max(0, yy2 - yy1); 17 | 18 | // calculate area of intersection and union 19 | float det_area = det.area(); 20 | float trk_area = trk.area(); 21 | auto intersection_area = w * h; 22 | float union_area = det_area + trk_area - intersection_area; 23 | auto iou = intersection_area / union_area; 24 | return iou; 25 | } 26 | 27 | 28 | void Tracker::HungarianMatching(const std::vector>& iou_matrix, 29 | size_t nrows, size_t ncols, 30 | std::vector>& association) { 31 | Matrix matrix(nrows, ncols); 32 | // Initialize matrix with IOU values 33 | for (size_t i = 0 ; i < nrows ; i++) { 34 | for (size_t j = 0 ; j < ncols ; j++) { 35 | // Multiply by -1 to find max cost 36 | if (iou_matrix[i][j] != 0) { 37 | matrix(i, j) = -iou_matrix[i][j]; 38 | } 39 | else { 40 | // TODO: figure out why we have to assign value to get correct result 41 | matrix(i, j) = 1.0f; 42 | } 43 | } 44 | } 45 | 46 | // // Display begin matrix state. 47 | // for (size_t row = 0 ; row < nrows ; row++) { 48 | // for (size_t col = 0 ; col < ncols ; col++) { 49 | // std::cout.width(10); 50 | // std::cout << matrix(row,col) << ","; 51 | // } 52 | // std::cout << std::endl; 53 | // } 54 | // std::cout << std::endl; 55 | 56 | 57 | // Apply Kuhn-Munkres algorithm to matrix. 58 | Munkres m; 59 | m.solve(matrix); 60 | 61 | // // Display solved matrix. 62 | // for (size_t row = 0 ; row < nrows ; row++) { 63 | // for (size_t col = 0 ; col < ncols ; col++) { 64 | // std::cout.width(2); 65 | // std::cout << matrix(row,col) << ","; 66 | // } 67 | // std::cout << std::endl; 68 | // } 69 | // std::cout << std::endl; 70 | 71 | for (size_t i = 0 ; i < nrows ; i++) { 72 | for (size_t j = 0 ; j < ncols ; j++) { 73 | association[i][j] = matrix(i, j); 74 | } 75 | } 76 | } 77 | 78 | 79 | void Tracker::AssociateDetectionsToTrackers(const std::vector& detection, 80 | std::map& tracks, 81 | std::map& matched, 82 | std::vector& unmatched_det, 83 | float iou_threshold) { 84 | 85 | // Set all detection as unmatched if no tracks existing 86 | if (tracks.empty()) { 87 | for (const auto& det : detection) { 88 | unmatched_det.push_back(det); 89 | } 90 | return; 91 | } 92 | 93 | std::vector> iou_matrix; 94 | // resize IOU matrix based on number of detection and tracks 95 | iou_matrix.resize(detection.size(), std::vector(tracks.size())); 96 | 97 | std::vector> association; 98 | // resize association matrix based on number of detection and tracks 99 | association.resize(detection.size(), std::vector(tracks.size())); 100 | 101 | 102 | // row - detection, column - tracks 103 | for (size_t i = 0; i < detection.size(); i++) { 104 | size_t j = 0; 105 | for (const auto& trk : tracks) { 106 | iou_matrix[i][j] = CalculateIou(detection[i], trk.second); 107 | j++; 108 | } 109 | } 110 | 111 | // Find association 112 | HungarianMatching(iou_matrix, detection.size(), tracks.size(), association); 113 | 114 | for (size_t i = 0; i < detection.size(); i++) { 115 | bool matched_flag = false; 116 | size_t j = 0; 117 | for (const auto& trk : tracks) { 118 | if (0 == association[i][j]) { 119 | // Filter out matched with low IOU 120 | if (iou_matrix[i][j] >= iou_threshold) { 121 | matched[trk.first] = detection[i]; 122 | matched_flag = true; 123 | } 124 | // It builds 1 to 1 association, so we can break from here 125 | break; 126 | } 127 | j++; 128 | } 129 | // if detection cannot match with any tracks 130 | if (!matched_flag) { 131 | unmatched_det.push_back(detection[i]); 132 | } 133 | } 134 | } 135 | 136 | 137 | void Tracker::Run(const std::vector& detections) { 138 | 139 | /*** Predict internal tracks from previous frame ***/ 140 | for (auto &track : tracks_) { 141 | track.second.Predict(); 142 | } 143 | 144 | // Hash-map between track ID and associated detection bounding box 145 | std::map matched; 146 | // vector of unassociated detections 147 | std::vector unmatched_det; 148 | 149 | // return values - matched, unmatched_det 150 | if (!detections.empty()) { 151 | AssociateDetectionsToTrackers(detections, tracks_, matched, unmatched_det); 152 | } 153 | 154 | /*** Update tracks with associated bbox ***/ 155 | for (const auto &match : matched) { 156 | const auto &ID = match.first; 157 | tracks_[ID].Update(match.second); 158 | } 159 | 160 | /*** Create new tracks for unmatched detections ***/ 161 | for (const auto &det : unmatched_det) { 162 | Track tracker; 163 | tracker.Init(det); 164 | // Create new track and generate new ID 165 | tracks_[id_++] = tracker; 166 | } 167 | 168 | /*** Delete lose tracked tracks ***/ 169 | for (auto it = tracks_.begin(); it != tracks_.end();) { 170 | if (it->second.coast_cycles_ > kMaxCoastCycles) { 171 | it = tracks_.erase(it); 172 | } else { 173 | it++; 174 | } 175 | } 176 | } 177 | 178 | 179 | std::map Tracker::GetTracks() { 180 | return tracks_; 181 | } -------------------------------------------------------------------------------- /Track.cpp: -------------------------------------------------------------------------------- 1 | // TrackCounter.cpp: 定义应用程序的入口点。 2 | // 3 | #include 4 | #include 5 | #include 6 | #include "sort/include/tracker.h" 7 | #include "yolov5/Loader.h" 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include "time.h" 15 | using namespace std; 16 | 17 | #define INPUT_H Yolo::INPUT_H 18 | #define INPUT_W Yolo::INPUT_W 19 | #define ORIGINAL_W Yolo::ORIGINAL_W 20 | #define ORIGINAL_H Yolo::ORIGINAL_H 21 | 22 | static string engineFile; 23 | 24 | struct sharedData { 25 | std::queue* dataQueue; 26 | std::queue>>* detQueue; 27 | std::queue* imgQueue; 28 | }; 29 | 30 | cv::Rect get_rect(int w, int h, float bbox[4]) { 31 | int l, r, t, b; 32 | float r_w = INPUT_W / (w * 1.0); 33 | float r_h = INPUT_H / (h * 1.0); 34 | if (r_h > r_w) { 35 | l = bbox[0] - bbox[2] / 2.f; 36 | r = bbox[0] + bbox[2] / 2.f; 37 | t = bbox[1] - bbox[3] / 2.f - (INPUT_H - r_w * h) / 2; 38 | b = bbox[1] + bbox[3] / 2.f - (INPUT_H - r_w * h) / 2; 39 | l = l / r_w; 40 | r = r / r_w; 41 | t = t / r_w; 42 | b = b / r_w; 43 | } 44 | else { 45 | l = bbox[0] - bbox[2] / 2.f - (INPUT_W - r_h * w) / 2; 46 | r = bbox[0] + bbox[2] / 2.f - (INPUT_W - r_h * w) / 2; 47 | t = bbox[1] - bbox[3] / 2.f; 48 | b = bbox[1] + bbox[3] / 2.f; 49 | l = l / r_h; 50 | r = r / r_h; 51 | t = t / r_h; 52 | b = b / r_h; 53 | } 54 | return cv::Rect(l, t, r - l, b - t); 55 | } 56 | 57 | 58 | int main(int argc, char** argv) { 59 | if (argc == 6 && std::string(argv[1]) == "-s") { 60 | yolov5_build_engine(argv[2],argv[3],atof(argv[4]),atof(argv[5])); 61 | } 62 | else if (argc == 5 && std::string(argv[1]) == "-s" && std::string(argv[1]) == "s") { 63 | yolov5_build_engine(argv[2], argv[3], 0.33, 0.50); 64 | } 65 | else if (argc == 5 && std::string(argv[1]) == "-s" && std::string(argv[1]) == "m") { 66 | yolov5_build_engine(argv[2], argv[3], 0.67, 0.75); 67 | } 68 | else if (argc == 5 && std::string(argv[1]) == "-s" && std::string(argv[1]) == "l") { 69 | yolov5_build_engine(argv[2], argv[3], 1.0, 1.0); 70 | } 71 | else if (argc == 5 && std::string(argv[1]) == "-s" && std::string(argv[1]) == "x") { 72 | yolov5_build_engine(argv[2], argv[3], 1.25, 1.33); 73 | } 74 | else if (argc == 4 && std::string(argv[1]) == "-v") { 75 | //initialize capture, writer, detector, and tracker 76 | std::string filename = argv[2]; 77 | cv::VideoCapture cap; 78 | cap.open(filename); 79 | if (!cap.isOpened()) { 80 | std::cout << "failed to open video " << filename << std::endl; 81 | return -1; 82 | } 83 | int frameCount = cap.get(cv::CAP_PROP_FRAME_COUNT); 84 | int maxIndex = int(frameCount / BATCH_SIZE); 85 | 86 | yolov5_detector detector(argv[3]); 87 | Tracker sort; 88 | std::ofstream resultFile("testSingleThread.txt", std::ios::out); 89 | std::map tracks; 90 | cv::VideoWriter writer; 91 | 92 | auto code = cv::VideoWriter::fourcc('X', 'V', 'I', 'D'); 93 | auto fps = cap.get(cv::CAP_PROP_FPS); 94 | auto size = cv::Size(cap.get(cv::CAP_PROP_FRAME_WIDTH), cap.get(cv::CAP_PROP_FRAME_HEIGHT)); 95 | writer.open("../output/result.avi", code, fps, size, true); 96 | 97 | //get a batch of images and preprocess them 98 | static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; 99 | for (int f = 0; f < maxIndex; f++) { 100 | auto start = std::chrono::system_clock::now(); 101 | vector imgs; 102 | for (int b = 0; b < BATCH_SIZE; b++) { 103 | cv::Mat img; 104 | cap.read(img); 105 | //cv::Mat copy_img; 106 | //img.copyTo(copy_img); 107 | auto pr_img = preprocess_img(img); 108 | imgs.push_back(img); 109 | int i = 0; 110 | for (int row = 0; row < INPUT_H; ++row) { 111 | uchar* uc_pixel = pr_img.data + row * pr_img.step; 112 | for (int col = 0; col < INPUT_W; ++col) { 113 | data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0; 114 | data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0; 115 | data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0; 116 | uc_pixel += 3; 117 | ++i; 118 | } 119 | } 120 | } 121 | 122 | //detect 123 | auto end = std::chrono::system_clock::now(); 124 | std::cout << "Preprocess Time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; 125 | start = std::chrono::system_clock::now(); 126 | std::vector> results = detector.detect(data); 127 | end = std::chrono::system_clock::now(); 128 | std::cout << "Detect Time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; 129 | start = std::chrono::system_clock::now(); 130 | 131 | //track 132 | for (int b = 0; b < BATCH_SIZE; b++) { 133 | auto& result = results[b]; 134 | std::vector dets; 135 | for (size_t j = 0; j < result.size(); j++) { 136 | dets.push_back(get_rect(ORIGINAL_W, ORIGINAL_H, result[j].bbox)); 137 | } 138 | sort.Run(dets); 139 | std::map tracks = sort.GetTracks(); 140 | 141 | auto im = imgs[b]; 142 | for (auto& trk : tracks) { 143 | const auto& bbox = trk.second.GetStateAsBbox(); 144 | /*Note that we will not export coasted tracks 145 | If we export coasted tracks, the total number of false negative will decrease (and maybe ID switch) 146 | However, the total number of false positive will increase more (from experiments), 147 | which leads to MOTA decrease 148 | Developer can export coasted cycles if false negative tracks is critical in the system*/ 149 | if (trk.second.coast_cycles_ < kMaxCoastCycles 150 | && (trk.second.hit_streak_ >= kMinHits || f * BATCH_SIZE + b < kMinHits)) { 151 | // Print to terminal for debugging 152 | 153 | resultFile << f * BATCH_SIZE + b << "," << trk.first << "," << bbox.tl().x << "," << bbox.tl().y 154 | << "," << bbox.width << "," << bbox.height << ",1,-1,-1,-1\n"; 155 | cv::rectangle(im, cv::Rect(bbox.tl().x, bbox.tl().y, bbox.width, bbox.height), cv::Scalar(0x27, 0xC1, 0x36), 2); 156 | cv::putText(im, std::to_string(trk.first), cv::Point(bbox.tl().x, bbox.tl().y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); 157 | } 158 | } 159 | writer.write(im); 160 | } 161 | end = std::chrono::system_clock::now(); 162 | std::cout << "Track Time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; 163 | imgs.clear(); 164 | } 165 | detector.release(); 166 | resultFile.close(); 167 | cap.release(); 168 | writer.release(); 169 | } 170 | return 0; 171 | } 172 | -------------------------------------------------------------------------------- /sort/include/matrix.cpp: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2007 John Weaver 3 | * 4 | * This program is free software; you can redistribute it and/or modify 5 | * it under the terms of the GNU General Public License as published by 6 | * the Free Software Foundation; either version 2 of the License, or 7 | * (at your option) any later version. 8 | * 9 | * This program is distributed in the hope that it will be useful, 10 | * but WITHOUT ANY WARRANTY; without even the implied warranty of 11 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 | * GNU General Public License for more details. 13 | * 14 | * You should have received a copy of the GNU General Public License 15 | * along with this program; if not, write to the Free Software 16 | * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 17 | */ 18 | 19 | #include "matrix.h" 20 | 21 | #include 22 | #include 23 | #include 24 | 25 | /*export*/ template 26 | Matrix::Matrix() { 27 | m_rows = 0; 28 | m_columns = 0; 29 | m_matrix = nullptr; 30 | } 31 | 32 | 33 | /*export*/ template 34 | Matrix::Matrix(const std::initializer_list> init) { 35 | m_matrix = nullptr; 36 | m_rows = init.size(); 37 | if ( m_rows == 0 ) { 38 | m_columns = 0; 39 | } else { 40 | m_columns = init.begin()->size(); 41 | if ( m_columns > 0 ) { 42 | resize(m_rows, m_columns); 43 | } 44 | } 45 | 46 | size_t i = 0, j; 47 | for ( auto row = init.begin() ; row != init.end() ; ++row, ++i ) { 48 | assert ( row->size() == m_columns && "All rows must have the same number of columns." ); 49 | j = 0; 50 | for ( auto value = row->begin() ; value != row->end() ; ++value, ++j ) { 51 | m_matrix[i][j] = *value; 52 | } 53 | } 54 | } 55 | 56 | /*export*/ template 57 | Matrix::Matrix(const Matrix &other) { 58 | if ( other.m_matrix != nullptr ) { 59 | // copy arrays 60 | m_matrix = nullptr; 61 | resize(other.m_rows, other.m_columns); 62 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 63 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 64 | m_matrix[i][j] = other.m_matrix[i][j]; 65 | } 66 | } 67 | } else { 68 | m_matrix = nullptr; 69 | m_rows = 0; 70 | m_columns = 0; 71 | } 72 | } 73 | 74 | /*export*/ template 75 | Matrix::Matrix(const size_t rows, const size_t columns) { 76 | m_matrix = nullptr; 77 | resize(rows, columns); 78 | } 79 | 80 | /*export*/ template 81 | Matrix & 82 | Matrix::operator= (const Matrix &other) { 83 | if ( other.m_matrix != nullptr ) { 84 | // copy arrays 85 | resize(other.m_rows, other.m_columns); 86 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 87 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 88 | m_matrix[i][j] = other.m_matrix[i][j]; 89 | } 90 | } 91 | } else { 92 | // free arrays 93 | for ( size_t i = 0 ; i < m_columns ; i++ ) { 94 | delete [] m_matrix[i]; 95 | } 96 | 97 | delete [] m_matrix; 98 | 99 | m_matrix = nullptr; 100 | m_rows = 0; 101 | m_columns = 0; 102 | } 103 | 104 | return *this; 105 | } 106 | 107 | /*export*/ template 108 | Matrix::~Matrix() { 109 | if ( m_matrix != nullptr ) { 110 | // free arrays 111 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 112 | delete [] m_matrix[i]; 113 | } 114 | 115 | delete [] m_matrix; 116 | } 117 | m_matrix = nullptr; 118 | } 119 | 120 | /*export*/ template 121 | void 122 | Matrix::resize(const size_t rows, const size_t columns, const T default_value) { 123 | assert ( rows > 0 && columns > 0 && "Columns and rows must exist." ); 124 | 125 | if ( m_matrix == nullptr ) { 126 | // alloc arrays 127 | m_matrix = new T*[rows]; // rows 128 | for ( size_t i = 0 ; i < rows ; i++ ) { 129 | m_matrix[i] = new T[columns]; // columns 130 | } 131 | 132 | m_rows = rows; 133 | m_columns = columns; 134 | clear(); 135 | } else { 136 | // save array pointer 137 | T **new_matrix; 138 | // alloc new arrays 139 | new_matrix = new T*[rows]; // rows 140 | for ( size_t i = 0 ; i < rows ; i++ ) { 141 | new_matrix[i] = new T[columns]; // columns 142 | for ( size_t j = 0 ; j < columns ; j++ ) { 143 | new_matrix[i][j] = default_value; 144 | } 145 | } 146 | 147 | // copy data from saved pointer to new arrays 148 | size_t minrows = std::min(rows, m_rows); 149 | size_t mincols = std::min(columns, m_columns); 150 | for ( size_t x = 0 ; x < minrows ; x++ ) { 151 | for ( size_t y = 0 ; y < mincols ; y++ ) { 152 | new_matrix[x][y] = m_matrix[x][y]; 153 | } 154 | } 155 | 156 | // delete old arrays 157 | if ( m_matrix != nullptr ) { 158 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 159 | delete [] m_matrix[i]; 160 | } 161 | 162 | delete [] m_matrix; 163 | } 164 | 165 | m_matrix = new_matrix; 166 | } 167 | 168 | m_rows = rows; 169 | m_columns = columns; 170 | } 171 | 172 | /*export*/ template 173 | void 174 | Matrix::clear() { 175 | assert( m_matrix != nullptr ); 176 | 177 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 178 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 179 | m_matrix[i][j] = 0; 180 | } 181 | } 182 | } 183 | 184 | /*export*/ template 185 | inline T& 186 | Matrix::operator ()(const size_t x, const size_t y) { 187 | assert ( x < m_rows ); 188 | assert ( y < m_columns ); 189 | assert ( m_matrix != nullptr ); 190 | return m_matrix[x][y]; 191 | } 192 | 193 | 194 | /*export*/ template 195 | inline const T& 196 | Matrix::operator ()(const size_t x, const size_t y) const { 197 | assert ( x < m_rows ); 198 | assert ( y < m_columns ); 199 | assert ( m_matrix != nullptr ); 200 | return m_matrix[x][y]; 201 | } 202 | 203 | 204 | /*export*/ template 205 | const T 206 | Matrix::min() const { 207 | assert( m_matrix != nullptr ); 208 | assert ( m_rows > 0 ); 209 | assert ( m_columns > 0 ); 210 | T min = m_matrix[0][0]; 211 | 212 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 213 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 214 | min = std::min(min, m_matrix[i][j]); 215 | } 216 | } 217 | 218 | return min; 219 | } 220 | 221 | 222 | /*export*/ template 223 | const T 224 | Matrix::max() const { 225 | assert( m_matrix != nullptr ); 226 | assert ( m_rows > 0 ); 227 | assert ( m_columns > 0 ); 228 | T max = m_matrix[0][0]; 229 | 230 | for ( size_t i = 0 ; i < m_rows ; i++ ) { 231 | for ( size_t j = 0 ; j < m_columns ; j++ ) { 232 | max = std::max(max, m_matrix[i][j]); 233 | } 234 | } 235 | 236 | return max; 237 | } 238 | -------------------------------------------------------------------------------- /yolov5/utils/autoanchor.py: -------------------------------------------------------------------------------- 1 | # Auto-anchor utils 2 | 3 | import numpy as np 4 | import torch 5 | import yaml 6 | from scipy.cluster.vq import kmeans 7 | from tqdm import tqdm 8 | 9 | 10 | def check_anchor_order(m): 11 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary 12 | a = m.anchor_grid.prod(-1).view(-1) # anchor area 13 | da = a[-1] - a[0] # delta a 14 | ds = m.stride[-1] - m.stride[0] # delta s 15 | if da.sign() != ds.sign(): # same order 16 | print('Reversing anchor order') 17 | m.anchors[:] = m.anchors.flip(0) 18 | m.anchor_grid[:] = m.anchor_grid.flip(0) 19 | 20 | 21 | def check_anchors(dataset, model, thr=4.0, imgsz=640): 22 | # Check anchor fit to data, recompute if necessary 23 | print('\nAnalyzing anchors... ', end='') 24 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() 25 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) 26 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale 27 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh 28 | 29 | def metric(k): # compute metric 30 | r = wh[:, None] / k[None] 31 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 32 | best = x.max(1)[0] # best_x 33 | aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold 34 | bpr = (best > 1. / thr).float().mean() # best possible recall 35 | return bpr, aat 36 | 37 | bpr, aat = metric(m.anchor_grid.clone().cpu().view(-1, 2)) 38 | print('anchors/target = %.2f, Best Possible Recall (BPR) = %.4f' % (aat, bpr), end='') 39 | if bpr < 0.98: # threshold to recompute 40 | print('. Attempting to improve anchors, please wait...') 41 | na = m.anchor_grid.numel() // 2 # number of anchors 42 | new_anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) 43 | new_bpr = metric(new_anchors.reshape(-1, 2))[0] 44 | if new_bpr > bpr: # replace anchors 45 | new_anchors = torch.tensor(new_anchors, device=m.anchors.device).type_as(m.anchors) 46 | m.anchor_grid[:] = new_anchors.clone().view_as(m.anchor_grid) # for inference 47 | m.anchors[:] = new_anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss 48 | check_anchor_order(m) 49 | print('New anchors saved to model. Update model *.yaml to use these anchors in the future.') 50 | else: 51 | print('Original anchors better than new anchors. Proceeding with original anchors.') 52 | print('') # newline 53 | 54 | 55 | def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): 56 | """ Creates kmeans-evolved anchors from training dataset 57 | 58 | Arguments: 59 | path: path to dataset *.yaml, or a loaded dataset 60 | n: number of anchors 61 | img_size: image size used for training 62 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 63 | gen: generations to evolve anchors using genetic algorithm 64 | verbose: print all results 65 | 66 | Return: 67 | k: kmeans evolved anchors 68 | 69 | Usage: 70 | from utils.autoanchor import *; _ = kmean_anchors() 71 | """ 72 | thr = 1. / thr 73 | 74 | def metric(k, wh): # compute metrics 75 | r = wh[:, None] / k[None] 76 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric 77 | # x = wh_iou(wh, torch.tensor(k)) # iou metric 78 | return x, x.max(1)[0] # x, best_x 79 | 80 | def anchor_fitness(k): # mutation fitness 81 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh) 82 | return (best * (best > thr).float()).mean() # fitness 83 | 84 | def print_results(k): 85 | k = k[np.argsort(k.prod(1))] # sort small to large 86 | x, best = metric(k, wh0) 87 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr 88 | print('thr=%.2f: %.4f best possible recall, %.2f anchors past thr' % (thr, bpr, aat)) 89 | print('n=%g, img_size=%s, metric_all=%.3f/%.3f-mean/best, past_thr=%.3f-mean: ' % 90 | (n, img_size, x.mean(), best.mean(), x[x > thr].mean()), end='') 91 | for i, x in enumerate(k): 92 | print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg 93 | return k 94 | 95 | if isinstance(path, str): # *.yaml file 96 | with open(path) as f: 97 | data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict 98 | from utils.datasets import LoadImagesAndLabels 99 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) 100 | else: 101 | dataset = path # dataset 102 | 103 | # Get label wh 104 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) 105 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh 106 | 107 | # Filter 108 | i = (wh0 < 3.0).any(1).sum() 109 | if i: 110 | print('WARNING: Extremely small objects found. ' 111 | '%g of %g labels are < 3 pixels in width or height.' % (i, len(wh0))) 112 | wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels 113 | 114 | # Kmeans calculation 115 | print('Running kmeans for %g anchors on %g points...' % (n, len(wh))) 116 | s = wh.std(0) # sigmas for whitening 117 | k, dist = kmeans(wh / s, n, iter=30) # points, mean distance 118 | k *= s 119 | wh = torch.tensor(wh, dtype=torch.float32) # filtered 120 | wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered 121 | k = print_results(k) 122 | 123 | # Plot 124 | # k, d = [None] * 20, [None] * 20 125 | # for i in tqdm(range(1, 21)): 126 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance 127 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True) 128 | # ax = ax.ravel() 129 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.') 130 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh 131 | # ax[0].hist(wh[wh[:, 0]<100, 0],400) 132 | # ax[1].hist(wh[wh[:, 1]<100, 1],400) 133 | # fig.savefig('wh.png', dpi=200) 134 | 135 | # Evolve 136 | npr = np.random 137 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma 138 | pbar = tqdm(range(gen), desc='Evolving anchors with Genetic Algorithm') # progress bar 139 | for _ in pbar: 140 | v = np.ones(sh) 141 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates) 142 | v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0) 143 | kg = (k.copy() * v).clip(min=2.0) 144 | fg = anchor_fitness(kg) 145 | if fg > f: 146 | f, k = fg, kg.copy() 147 | pbar.desc = 'Evolving anchors with Genetic Algorithm: fitness = %.4f' % f 148 | if verbose: 149 | print_results(k) 150 | 151 | return print_results(k) 152 | -------------------------------------------------------------------------------- /yolov5/utils/metrics.py: -------------------------------------------------------------------------------- 1 | # Model validation metrics 2 | 3 | from pathlib import Path 4 | 5 | import matplotlib.pyplot as plt 6 | import numpy as np 7 | import torch 8 | 9 | from . import general 10 | 11 | 12 | def fitness(x): 13 | # Model fitness as a weighted combination of metrics 14 | w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95] 15 | return (x[:, :4] * w).sum(1) 16 | 17 | 18 | def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision-recall_curve.png', names=[]): 19 | """ Compute the average precision, given the recall and precision curves. 20 | Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. 21 | # Arguments 22 | tp: True positives (nparray, nx1 or nx10). 23 | conf: Objectness value from 0-1 (nparray). 24 | pred_cls: Predicted object classes (nparray). 25 | target_cls: True object classes (nparray). 26 | plot: Plot precision-recall curve at mAP@0.5 27 | save_dir: Plot save directory 28 | # Returns 29 | The average precision as computed in py-faster-rcnn. 30 | """ 31 | 32 | # Sort by objectness 33 | i = np.argsort(-conf) 34 | tp, conf, pred_cls = tp[i], conf[i], pred_cls[i] 35 | 36 | # Find unique classes 37 | unique_classes = np.unique(target_cls) 38 | 39 | # Create Precision-Recall curve and compute AP for each class 40 | px, py = np.linspace(0, 1, 1000), [] # for plotting 41 | pr_score = 0.1 # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898 42 | s = [unique_classes.shape[0], tp.shape[1]] # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95) 43 | ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s) 44 | for ci, c in enumerate(unique_classes): 45 | i = pred_cls == c 46 | n_l = (target_cls == c).sum() # number of labels 47 | n_p = i.sum() # number of predictions 48 | 49 | if n_p == 0 or n_l == 0: 50 | continue 51 | else: 52 | # Accumulate FPs and TPs 53 | fpc = (1 - tp[i]).cumsum(0) 54 | tpc = tp[i].cumsum(0) 55 | 56 | # Recall 57 | recall = tpc / (n_l + 1e-16) # recall curve 58 | r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0]) # r at pr_score, negative x, xp because xp decreases 59 | 60 | # Precision 61 | precision = tpc / (tpc + fpc) # precision curve 62 | p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0]) # p at pr_score 63 | 64 | # AP from recall-precision curve 65 | for j in range(tp.shape[1]): 66 | ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j]) 67 | if plot and (j == 0): 68 | py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5 69 | 70 | # Compute F1 score (harmonic mean of precision and recall) 71 | f1 = 2 * p * r / (p + r + 1e-16) 72 | 73 | if plot: 74 | plot_pr_curve(px, py, ap, save_dir, names) 75 | 76 | return p, r, ap, f1, unique_classes.astype('int32') 77 | 78 | 79 | def compute_ap(recall, precision): 80 | """ Compute the average precision, given the recall and precision curves 81 | # Arguments 82 | recall: The recall curve (list) 83 | precision: The precision curve (list) 84 | # Returns 85 | Average precision, precision curve, recall curve 86 | """ 87 | 88 | # Append sentinel values to beginning and end 89 | mrec = np.concatenate(([0.], recall, [recall[-1] + 0.01])) 90 | mpre = np.concatenate(([1.], precision, [0.])) 91 | 92 | # Compute the precision envelope 93 | mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) 94 | 95 | # Integrate area under curve 96 | method = 'interp' # methods: 'continuous', 'interp' 97 | if method == 'interp': 98 | x = np.linspace(0, 1, 101) # 101-point interp (COCO) 99 | ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate 100 | else: # 'continuous' 101 | i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes 102 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve 103 | 104 | return ap, mpre, mrec 105 | 106 | 107 | class ConfusionMatrix: 108 | # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix 109 | def __init__(self, nc, conf=0.25, iou_thres=0.45): 110 | self.matrix = np.zeros((nc + 1, nc + 1)) 111 | self.nc = nc # number of classes 112 | self.conf = conf 113 | self.iou_thres = iou_thres 114 | 115 | def process_batch(self, detections, labels): 116 | """ 117 | Return intersection-over-union (Jaccard index) of boxes. 118 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 119 | Arguments: 120 | detections (Array[N, 6]), x1, y1, x2, y2, conf, class 121 | labels (Array[M, 5]), class, x1, y1, x2, y2 122 | Returns: 123 | None, updates confusion matrix accordingly 124 | """ 125 | detections = detections[detections[:, 4] > self.conf] 126 | gt_classes = labels[:, 0].int() 127 | detection_classes = detections[:, 5].int() 128 | iou = general.box_iou(labels[:, 1:], detections[:, :4]) 129 | 130 | x = torch.where(iou > self.iou_thres) 131 | if x[0].shape[0]: 132 | matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() 133 | if x[0].shape[0] > 1: 134 | matches = matches[matches[:, 2].argsort()[::-1]] 135 | matches = matches[np.unique(matches[:, 1], return_index=True)[1]] 136 | matches = matches[matches[:, 2].argsort()[::-1]] 137 | matches = matches[np.unique(matches[:, 0], return_index=True)[1]] 138 | else: 139 | matches = np.zeros((0, 3)) 140 | 141 | n = matches.shape[0] > 0 142 | m0, m1, _ = matches.transpose().astype(np.int16) 143 | for i, gc in enumerate(gt_classes): 144 | j = m0 == i 145 | if n and sum(j) == 1: 146 | self.matrix[gc, detection_classes[m1[j]]] += 1 # correct 147 | else: 148 | self.matrix[gc, self.nc] += 1 # background FP 149 | 150 | if n: 151 | for i, dc in enumerate(detection_classes): 152 | if not any(m1 == i): 153 | self.matrix[self.nc, dc] += 1 # background FN 154 | 155 | def matrix(self): 156 | return self.matrix 157 | 158 | def plot(self, save_dir='', names=()): 159 | try: 160 | import seaborn as sn 161 | 162 | array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6) # normalize 163 | array[array < 0.005] = np.nan # don't annotate (would appear as 0.00) 164 | 165 | fig = plt.figure(figsize=(12, 9), tight_layout=True) 166 | sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size 167 | labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels 168 | sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, 169 | xticklabels=names + ['background FN'] if labels else "auto", 170 | yticklabels=names + ['background FP'] if labels else "auto").set_facecolor((1, 1, 1)) 171 | fig.axes[0].set_xlabel('True') 172 | fig.axes[0].set_ylabel('Predicted') 173 | fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) 174 | except Exception as e: 175 | pass 176 | 177 | def print(self): 178 | for i in range(self.nc + 1): 179 | print(' '.join(map(str, self.matrix[i]))) 180 | 181 | 182 | # Plots ---------------------------------------------------------------------------------------------------------------- 183 | 184 | def plot_pr_curve(px, py, ap, save_dir='.', names=()): 185 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) 186 | py = np.stack(py, axis=1) 187 | 188 | if 0 < len(names) < 21: # show mAP in legend if < 10 classes 189 | for i, y in enumerate(py.T): 190 | ax.plot(px, y, linewidth=1, label=f'{names[i]} %.3f' % ap[i, 0]) # plot(recall, precision) 191 | else: 192 | ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision) 193 | 194 | ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean()) 195 | ax.set_xlabel('Recall') 196 | ax.set_ylabel('Precision') 197 | ax.set_xlim(0, 1) 198 | ax.set_ylim(0, 1) 199 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") 200 | fig.savefig(Path(save_dir) / 'precision_recall_curve.png', dpi=250) 201 | -------------------------------------------------------------------------------- /yolov5/utils/loss.py: -------------------------------------------------------------------------------- 1 | # Loss functions 2 | 3 | import torch 4 | import torch.nn as nn 5 | 6 | from utils.general import bbox_iou 7 | from utils.torch_utils import is_parallel 8 | 9 | 10 | def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441 11 | # return positive, negative label smoothing BCE targets 12 | return 1.0 - 0.5 * eps, 0.5 * eps 13 | 14 | 15 | class BCEBlurWithLogitsLoss(nn.Module): 16 | # BCEwithLogitLoss() with reduced missing label effects. 17 | def __init__(self, alpha=0.05): 18 | super(BCEBlurWithLogitsLoss, self).__init__() 19 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss() 20 | self.alpha = alpha 21 | 22 | def forward(self, pred, true): 23 | loss = self.loss_fcn(pred, true) 24 | pred = torch.sigmoid(pred) # prob from logits 25 | dx = pred - true # reduce only missing label effects 26 | # dx = (pred - true).abs() # reduce missing label and false label effects 27 | alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4)) 28 | loss *= alpha_factor 29 | return loss.mean() 30 | 31 | 32 | class FocalLoss(nn.Module): 33 | # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 34 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 35 | super(FocalLoss, self).__init__() 36 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() 37 | self.gamma = gamma 38 | self.alpha = alpha 39 | self.reduction = loss_fcn.reduction 40 | self.loss_fcn.reduction = 'none' # required to apply FL to each element 41 | 42 | def forward(self, pred, true): 43 | loss = self.loss_fcn(pred, true) 44 | # p_t = torch.exp(-loss) 45 | # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability 46 | 47 | # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py 48 | pred_prob = torch.sigmoid(pred) # prob from logits 49 | p_t = true * pred_prob + (1 - true) * (1 - pred_prob) 50 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 51 | modulating_factor = (1.0 - p_t) ** self.gamma 52 | loss *= alpha_factor * modulating_factor 53 | 54 | if self.reduction == 'mean': 55 | return loss.mean() 56 | elif self.reduction == 'sum': 57 | return loss.sum() 58 | else: # 'none' 59 | return loss 60 | 61 | 62 | class QFocalLoss(nn.Module): 63 | # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5) 64 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): 65 | super(QFocalLoss, self).__init__() 66 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss() 67 | self.gamma = gamma 68 | self.alpha = alpha 69 | self.reduction = loss_fcn.reduction 70 | self.loss_fcn.reduction = 'none' # required to apply FL to each element 71 | 72 | def forward(self, pred, true): 73 | loss = self.loss_fcn(pred, true) 74 | 75 | pred_prob = torch.sigmoid(pred) # prob from logits 76 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha) 77 | modulating_factor = torch.abs(true - pred_prob) ** self.gamma 78 | loss *= alpha_factor * modulating_factor 79 | 80 | if self.reduction == 'mean': 81 | return loss.mean() 82 | elif self.reduction == 'sum': 83 | return loss.sum() 84 | else: # 'none' 85 | return loss 86 | 87 | 88 | def compute_loss(p, targets, model): # predictions, targets, model 89 | device = targets.device 90 | lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device) 91 | tcls, tbox, indices, anchors = build_targets(p, targets, model) # targets 92 | h = model.hyp # hyperparameters 93 | 94 | # Define criteria 95 | BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) # weight=model.class_weights) 96 | BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device)) 97 | 98 | # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 99 | cp, cn = smooth_BCE(eps=0.0) 100 | 101 | # Focal loss 102 | g = h['fl_gamma'] # focal loss gamma 103 | if g > 0: 104 | BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) 105 | 106 | # Losses 107 | nt = 0 # number of targets 108 | no = len(p) # number of outputs 109 | balance = [4.0, 1.0, 0.4] if no == 3 else [4.0, 1.0, 0.4, 0.1] # P3-5 or P3-6 110 | for i, pi in enumerate(p): # layer index, layer predictions 111 | b, a, gj, gi = indices[i] # image, anchor, gridy, gridx 112 | tobj = torch.zeros_like(pi[..., 0], device=device) # target obj 113 | 114 | n = b.shape[0] # number of targets 115 | if n: 116 | nt += n # cumulative targets 117 | ps = pi[b, a, gj, gi] # prediction subset corresponding to targets 118 | 119 | # Regression 120 | pxy = ps[:, :2].sigmoid() * 2. - 0.5 121 | pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i] 122 | pbox = torch.cat((pxy, pwh), 1) # predicted box 123 | iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target) 124 | lbox += (1.0 - iou).mean() # iou loss 125 | 126 | # Objectness 127 | tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio 128 | 129 | # Classification 130 | if model.nc > 1: # cls loss (only if multiple classes) 131 | t = torch.full_like(ps[:, 5:], cn, device=device) # targets 132 | t[range(n), tcls[i]] = cp 133 | lcls += BCEcls(ps[:, 5:], t) # BCE 134 | 135 | # Append targets to text file 136 | # with open('targets.txt', 'a') as file: 137 | # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)] 138 | 139 | lobj += BCEobj(pi[..., 4], tobj) * balance[i] # obj loss 140 | 141 | s = 3 / no # output count scaling 142 | lbox *= h['box'] * s 143 | lobj *= h['obj'] * s * (1.4 if no == 4 else 1.) 144 | lcls *= h['cls'] * s 145 | bs = tobj.shape[0] # batch size 146 | 147 | loss = lbox + lobj + lcls 148 | return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach() 149 | 150 | 151 | def build_targets(p, targets, model): 152 | # Build targets for compute_loss(), input targets(image,class,x,y,w,h) 153 | det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module 154 | na, nt = det.na, targets.shape[0] # number of anchors, targets 155 | tcls, tbox, indices, anch = [], [], [], [] 156 | gain = torch.ones(7, device=targets.device) # normalized to gridspace gain 157 | ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt) 158 | targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices 159 | 160 | g = 0.5 # bias 161 | off = torch.tensor([[0, 0], 162 | [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m 163 | # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm 164 | ], device=targets.device).float() * g # offsets 165 | 166 | for i in range(det.nl): 167 | anchors = det.anchors[i] 168 | gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain 169 | 170 | # Match targets to anchors 171 | t = targets * gain 172 | if nt: 173 | # Matches 174 | r = t[:, :, 4:6] / anchors[:, None] # wh ratio 175 | j = torch.max(r, 1. / r).max(2)[0] < model.hyp['anchor_t'] # compare 176 | # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) 177 | t = t[j] # filter 178 | 179 | # Offsets 180 | gxy = t[:, 2:4] # grid xy 181 | gxi = gain[[2, 3]] - gxy # inverse 182 | j, k = ((gxy % 1. < g) & (gxy > 1.)).T 183 | l, m = ((gxi % 1. < g) & (gxi > 1.)).T 184 | j = torch.stack((torch.ones_like(j), j, k, l, m)) 185 | t = t.repeat((5, 1, 1))[j] 186 | offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j] 187 | else: 188 | t = targets[0] 189 | offsets = 0 190 | 191 | # Define 192 | b, c = t[:, :2].long().T # image, class 193 | gxy = t[:, 2:4] # grid xy 194 | gwh = t[:, 4:6] # grid wh 195 | gij = (gxy - offsets).long() 196 | gi, gj = gij.T # grid xy indices 197 | 198 | # Append 199 | a = t[:, 6].long() # anchor indices 200 | indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices 201 | tbox.append(torch.cat((gxy - gij, gwh), 1)) # box 202 | anch.append(anchors[a]) # anchors 203 | tcls.append(c) # class 204 | 205 | return tcls, tbox, indices, anch 206 | -------------------------------------------------------------------------------- /yolov5/yolov5.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "cuda_runtime_api.h" 4 | #include "logging.h" 5 | #include "common.hpp" 6 | #include "yolov5.h" 7 | 8 | // stuff we know about the network and the input/output blobs 9 | static const int INPUT_H = Yolo::INPUT_H; 10 | static const int INPUT_W = Yolo::INPUT_W; 11 | static const int CLASS_NUM = Yolo::CLASS_NUM; 12 | static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1 13 | const char* INPUT_BLOB_NAME = "data"; 14 | const char* OUTPUT_BLOB_NAME = "prob"; 15 | static Logger gLogger; 16 | float yolov5_detector::prob[BATCH_SIZE * OUTPUT_SIZE]; 17 | static float gw; 18 | static float gd; 19 | 20 | int get_width(int x, int divisor = 8) { 21 | //return math.ceil(x / divisor) * divisor 22 | if (int(x * gw) % divisor == 0) { 23 | return int(x * gw); 24 | } 25 | return (int(x * gw / divisor) + 1) * divisor; 26 | } 27 | 28 | int get_depth(int x) { 29 | if (x == 1) { 30 | return 1; 31 | } 32 | else { 33 | return round(x * gd) > 1 ? round(x * gd) : 1; 34 | } 35 | } 36 | 37 | ICudaEngine* build_engine(std::string weightFile,unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { 38 | INetworkDefinition* network = builder->createNetworkV2(0U); 39 | 40 | // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME 41 | ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W }); 42 | assert(data); 43 | 44 | std::map weightMap = loadWeights(weightFile); 45 | Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; 46 | 47 | /* ------ yolov5 backbone------ */ 48 | auto focus0 = focus(network, weightMap, *data, 3, get_width(64), 3, "model.0"); 49 | auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128), 3, 2, 1, "model.1"); 50 | auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128), get_width(128), get_depth(3), true, 1, 0.5, "model.2"); 51 | auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256), 3, 2, 1, "model.3"); 52 | auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256), get_width(256), get_depth(9), true, 1, 0.5, "model.4"); 53 | auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512), 3, 2, 1, "model.5"); 54 | auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512), get_width(512), get_depth(9), true, 1, 0.5, "model.6"); 55 | auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024), 3, 2, 1, "model.7"); 56 | auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), get_width(1024), get_width(1024), 5, 9, 13, "model.8"); 57 | 58 | /* ------ yolov5 head ------ */ 59 | auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), get_width(1024), get_width(1024), get_depth(3), false, 1, 0.5, "model.9"); 60 | auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), get_width(512), 1, 1, 1, "model.10"); 61 | 62 | float* deval = reinterpret_cast(malloc(sizeof(float) * get_width(512) * 2 * 2)); 63 | for (int i = 0; i < get_width(512) * 2 * 2; i++) { 64 | deval[i] = 1.0; 65 | } 66 | Weights deconvwts11{ DataType::kFLOAT, deval, get_width(512) * 2 * 2 }; 67 | IDeconvolutionLayer* deconv11 = network->addDeconvolutionNd(*conv10->getOutput(0), get_width(512), DimsHW{ 2, 2 }, deconvwts11, emptywts); 68 | deconv11->setStrideNd(DimsHW{ 2, 2 }); 69 | deconv11->setNbGroups(get_width(512)); 70 | weightMap["deconv11"] = deconvwts11; 71 | 72 | ITensor* inputTensors12[] = { deconv11->getOutput(0), bottleneck_csp6->getOutput(0) }; 73 | auto cat12 = network->addConcatenation(inputTensors12, 2); 74 | auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024), get_width(512), get_depth(3), false, 1, 0.5, "model.13"); 75 | auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256), 1, 1, 1, "model.14"); 76 | 77 | Weights deconvwts15{ DataType::kFLOAT, deval, get_width(256) * 2 * 2 }; 78 | IDeconvolutionLayer* deconv15 = network->addDeconvolutionNd(*conv14->getOutput(0), get_width(256), DimsHW{ 2, 2 }, deconvwts15, emptywts); 79 | deconv15->setStrideNd(DimsHW{ 2, 2 }); 80 | deconv15->setNbGroups(get_width(256)); 81 | ITensor* inputTensors16[] = { deconv15->getOutput(0), bottleneck_csp4->getOutput(0) }; 82 | auto cat16 = network->addConcatenation(inputTensors16, 2); 83 | 84 | auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512), get_width(256), get_depth(3), false, 1, 0.5, "model.17"); 85 | 86 | // yolo layer 0 87 | IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); 88 | auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256), 3, 2, 1, "model.18"); 89 | ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; 90 | auto cat19 = network->addConcatenation(inputTensors19, 2); 91 | auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512), get_width(512), get_depth(3), false, 1, 0.5, "model.20"); 92 | //yolo layer 1 93 | IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); 94 | auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512), 3, 2, 1, "model.21"); 95 | ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; 96 | auto cat22 = network->addConcatenation(inputTensors22, 2); 97 | auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024), get_width(1024), get_depth(3), false, 1, 0.5, "model.23"); 98 | IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); 99 | 100 | auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); 101 | yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); 102 | network->markOutput(*yolo->getOutput(0)); 103 | 104 | // Build engine 105 | builder->setMaxBatchSize(maxBatchSize); 106 | config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB 107 | #if defined(USE_FP16) 108 | config->setFlag(BuilderFlag::kFP16); 109 | #elif defined(USE_INT8) 110 | std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; 111 | assert(builder->platformHasFastInt8()); 112 | config->setFlag(BuilderFlag::kINT8); 113 | Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); 114 | config->setInt8Calibrator(calibrator); 115 | #endif 116 | 117 | std::cout << "Building engine, please wait for a while..." << std::endl; 118 | ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); 119 | std::cout << "Build engine successfully!" << std::endl; 120 | 121 | // Don't need the network any more 122 | network->destroy(); 123 | 124 | // Release host memory 125 | for (auto& mem : weightMap) 126 | { 127 | free((void*)(mem.second.values)); 128 | } 129 | 130 | return engine; 131 | } 132 | 133 | 134 | 135 | void APIToModel(std::string weightFile, unsigned int maxBatchSize, IHostMemory** modelStream) { 136 | // Create builder 137 | IBuilder* builder = createInferBuilder(gLogger); 138 | IBuilderConfig* config = builder->createBuilderConfig(); 139 | 140 | // Create model to populate the network, then set the outputs and create an engine 141 | //ICudaEngine* engine = (CREATENET(NET))(maxBatchSize, builder, config, DataType::kFLOAT); 142 | ICudaEngine* engine = build_engine(weightFile, maxBatchSize, builder, config, DataType::kFLOAT); 143 | 144 | //ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT); 145 | assert(engine != nullptr); 146 | 147 | // Serialize the engine 148 | (*modelStream) = engine->serialize(); 149 | 150 | // Close everything down 151 | engine->destroy(); 152 | builder->destroy(); 153 | } 154 | 155 | void doInference(IExecutionContext& context, cudaStream_t& stream, void** buffers, float* input, float* output, int batchSize) { 156 | // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host 157 | CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream)); 158 | context.enqueue(batchSize, buffers, stream, nullptr); 159 | CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); 160 | cudaStreamSynchronize(stream); 161 | } 162 | 163 | yolov5_detector::yolov5_detector(std::string engine_name) { 164 | cudaSetDevice(DEVICE); 165 | // create a model using the API directly and serialize it to a stream 166 | char* trtModelStream{ nullptr }; 167 | size_t size{ 0 }; 168 | 169 | std::ifstream file(engine_name, std::ios::binary); 170 | if (file.good()) { 171 | file.seekg(0, file.end); 172 | size = file.tellg(); 173 | file.seekg(0, file.beg); 174 | trtModelStream = new char[size]; 175 | assert(trtModelStream); 176 | file.read(trtModelStream, size); 177 | file.close(); 178 | } 179 | 180 | runtime = createInferRuntime(gLogger); 181 | assert(runtime != nullptr); 182 | engine = runtime->deserializeCudaEngine(trtModelStream, size); 183 | assert(engine != nullptr); 184 | context = engine->createExecutionContext(); 185 | assert(context != nullptr); 186 | delete[] trtModelStream; 187 | assert(engine->getNbBindings() == 2); 188 | // In order to bind the buffers, we need to know the names of the input and output tensors. 189 | // Note that indices are guaranteed to be less than IEngine::getNbBindings() 190 | const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME); 191 | const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME); 192 | assert(inputIndex == 0); 193 | assert(outputIndex == 1); 194 | // Create GPU buffers on device 195 | CHECK(cudaMalloc(&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float))); 196 | CHECK(cudaMalloc(&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float))); 197 | // Create stream 198 | CHECK(cudaStreamCreate(&stream)); 199 | } 200 | 201 | void yolov5_detector::release() { 202 | // Release stream and buffers 203 | cudaStreamDestroy(stream); 204 | CHECK(cudaFree(buffers[0])); 205 | CHECK(cudaFree(buffers[1])); 206 | // Destroy the engine 207 | context->destroy(); 208 | engine->destroy(); 209 | runtime->destroy(); 210 | } 211 | 212 | std::vector> yolov5_detector::detect(float* data) { 213 | doInference(*context, stream, buffers, data, prob, BATCH_SIZE); 214 | std::vector> batch_res(BATCH_SIZE); 215 | for (int b = 0; b < BATCH_SIZE; b++) { 216 | auto& res = batch_res[b]; 217 | nms(res, &prob[b * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH); 218 | } 219 | return batch_res; 220 | } 221 | 222 | 223 | int yolov5_build_engine(std::string weightFile, std::string engineFile,float depth_multiple, float width_multiple) { 224 | gw = width_multiple; 225 | gd = depth_multiple; 226 | cudaSetDevice(DEVICE); 227 | // create a model using the API directly and serialize it to a stream 228 | char* trtModelStream{ nullptr }; 229 | size_t size{ 0 }; 230 | std::cout << -1 << std::endl; 231 | //std::string engine_name = STR2(NET); 232 | //engine_name = "yolov5" + engine_name + ".engine"; 233 | std::string engine_name = engineFile; 234 | 235 | std::cout << -2 << std::endl; 236 | IHostMemory* modelStream{ nullptr }; 237 | std::cout << -3 << std::endl; 238 | APIToModel(weightFile,BATCH_SIZE, &modelStream); 239 | std::cout << 0 << std::endl; 240 | assert(modelStream != nullptr); 241 | std::cout << 0 << std::endl; 242 | std::ofstream p(engine_name, std::ios::binary); 243 | if (!p) { 244 | std::cerr << "could not open plan output file" << std::endl; 245 | return -1; 246 | } 247 | p.write(reinterpret_cast(modelStream->data()), modelStream->size()); 248 | modelStream->destroy(); 249 | return 0; 250 | } 251 | -------------------------------------------------------------------------------- /yolov5/utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | # PyTorch utils 2 | 3 | import logging 4 | import math 5 | import os 6 | import time 7 | from contextlib import contextmanager 8 | from copy import deepcopy 9 | 10 | import torch 11 | import torch.backends.cudnn as cudnn 12 | import torch.nn as nn 13 | import torch.nn.functional as F 14 | import torchvision 15 | 16 | try: 17 | import thop # for FLOPS computation 18 | except ImportError: 19 | thop = None 20 | logger = logging.getLogger(__name__) 21 | 22 | 23 | @contextmanager 24 | def torch_distributed_zero_first(local_rank: int): 25 | """ 26 | Decorator to make all processes in distributed training wait for each local_master to do something. 27 | """ 28 | if local_rank not in [-1, 0]: 29 | torch.distributed.barrier() 30 | yield 31 | if local_rank == 0: 32 | torch.distributed.barrier() 33 | 34 | 35 | def init_torch_seeds(seed=0): 36 | # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html 37 | torch.manual_seed(seed) 38 | if seed == 0: # slower, more reproducible 39 | cudnn.deterministic = True 40 | cudnn.benchmark = False 41 | else: # faster, less reproducible 42 | cudnn.deterministic = False 43 | cudnn.benchmark = True 44 | 45 | 46 | def select_device(device='', batch_size=None): 47 | # device = 'cpu' or '0' or '0,1,2,3' 48 | cpu_request = device.lower() == 'cpu' 49 | if device and not cpu_request: # if device requested other than 'cpu' 50 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable 51 | assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availablity 52 | 53 | cuda = False if cpu_request else torch.cuda.is_available() 54 | if cuda: 55 | c = 1024 ** 2 # bytes to MB 56 | ng = torch.cuda.device_count() 57 | if ng > 1 and batch_size: # check that batch_size is compatible with device_count 58 | assert batch_size % ng == 0, f'batch-size {batch_size} not multiple of GPU count {ng}' 59 | x = [torch.cuda.get_device_properties(i) for i in range(ng)] 60 | s = f'Using torch {torch.__version__} ' 61 | for i, d in enumerate((device or '0').split(',')): 62 | if i == 1: 63 | s = ' ' * len(s) 64 | logger.info(f"{s}CUDA:{d} ({x[i].name}, {x[i].total_memory / c}MB)") 65 | else: 66 | logger.info(f'Using torch {torch.__version__} CPU') 67 | 68 | logger.info('') # skip a line 69 | return torch.device('cuda:0' if cuda else 'cpu') 70 | 71 | 72 | def time_synchronized(): 73 | # pytorch-accurate time 74 | torch.cuda.synchronize() if torch.cuda.is_available() else None 75 | return time.time() 76 | 77 | 78 | def profile(x, ops, n=100, device=None): 79 | # profile a pytorch module or list of modules. Example usage: 80 | # x = torch.randn(16, 3, 640, 640) # input 81 | # m1 = lambda x: x * torch.sigmoid(x) 82 | # m2 = nn.SiLU() 83 | # profile(x, [m1, m2], n=100) # profile speed over 100 iterations 84 | 85 | device = device or torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 86 | x = x.to(device) 87 | x.requires_grad = True 88 | print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '') 89 | print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}") 90 | for m in ops if isinstance(ops, list) else [ops]: 91 | m = m.to(device) if hasattr(m, 'to') else m # device 92 | m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m # type 93 | dtf, dtb, t = 0., 0., [0., 0., 0.] # dt forward, backward 94 | try: 95 | flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPS 96 | except: 97 | flops = 0 98 | 99 | for _ in range(n): 100 | t[0] = time_synchronized() 101 | y = m(x) 102 | t[1] = time_synchronized() 103 | try: 104 | _ = y.sum().backward() 105 | t[2] = time_synchronized() 106 | except: # no backward method 107 | t[2] = float('nan') 108 | dtf += (t[1] - t[0]) * 1000 / n # ms per op forward 109 | dtb += (t[2] - t[1]) * 1000 / n # ms per op backward 110 | 111 | s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' 112 | s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list' 113 | p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters 114 | print(f'{p:12.4g}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}') 115 | 116 | 117 | def is_parallel(model): 118 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 119 | 120 | 121 | def intersect_dicts(da, db, exclude=()): 122 | # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values 123 | return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape} 124 | 125 | 126 | def initialize_weights(model): 127 | for m in model.modules(): 128 | t = type(m) 129 | if t is nn.Conv2d: 130 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 131 | elif t is nn.BatchNorm2d: 132 | m.eps = 1e-3 133 | m.momentum = 0.03 134 | elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 135 | m.inplace = True 136 | 137 | 138 | def find_modules(model, mclass=nn.Conv2d): 139 | # Finds layer indices matching module class 'mclass' 140 | return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)] 141 | 142 | 143 | def sparsity(model): 144 | # Return global model sparsity 145 | a, b = 0., 0. 146 | for p in model.parameters(): 147 | a += p.numel() 148 | b += (p == 0).sum() 149 | return b / a 150 | 151 | 152 | def prune(model, amount=0.3): 153 | # Prune model to requested global sparsity 154 | import torch.nn.utils.prune as prune 155 | print('Pruning model... ', end='') 156 | for name, m in model.named_modules(): 157 | if isinstance(m, nn.Conv2d): 158 | prune.l1_unstructured(m, name='weight', amount=amount) # prune 159 | prune.remove(m, 'weight') # make permanent 160 | print(' %.3g global sparsity' % sparsity(model)) 161 | 162 | 163 | def fuse_conv_and_bn(conv, bn): 164 | # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/ 165 | fusedconv = nn.Conv2d(conv.in_channels, 166 | conv.out_channels, 167 | kernel_size=conv.kernel_size, 168 | stride=conv.stride, 169 | padding=conv.padding, 170 | groups=conv.groups, 171 | bias=True).requires_grad_(False).to(conv.weight.device) 172 | 173 | # prepare filters 174 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 175 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 176 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) 177 | 178 | # prepare spatial bias 179 | b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias 180 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) 181 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) 182 | 183 | return fusedconv 184 | 185 | 186 | def model_info(model, verbose=False, img_size=640): 187 | # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320] 188 | n_p = sum(x.numel() for x in model.parameters()) # number parameters 189 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients 190 | if verbose: 191 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) 192 | for i, (name, p) in enumerate(model.named_parameters()): 193 | name = name.replace('module_list.', '') 194 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' % 195 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) 196 | 197 | try: # FLOPS 198 | from thop import profile 199 | stride = int(model.stride.max()) if hasattr(model, 'stride') else 32 200 | img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input 201 | flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS 202 | img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float 203 | fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPS 204 | except (ImportError, Exception): 205 | fs = '' 206 | 207 | logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}") 208 | 209 | 210 | def load_classifier(name='resnet101', n=2): 211 | # Loads a pretrained model reshaped to n-class output 212 | model = torchvision.models.__dict__[name](pretrained=True) 213 | 214 | # ResNet model properties 215 | # input_size = [3, 224, 224] 216 | # input_space = 'RGB' 217 | # input_range = [0, 1] 218 | # mean = [0.485, 0.456, 0.406] 219 | # std = [0.229, 0.224, 0.225] 220 | 221 | # Reshape output to n classes 222 | filters = model.fc.weight.shape[1] 223 | model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True) 224 | model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True) 225 | model.fc.out_features = n 226 | return model 227 | 228 | 229 | def scale_img(img, ratio=1.0, same_shape=False): # img(16,3,256,416), r=ratio 230 | # scales img(bs,3,y,x) by ratio 231 | if ratio == 1.0: 232 | return img 233 | else: 234 | h, w = img.shape[2:] 235 | s = (int(h * ratio), int(w * ratio)) # new size 236 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize 237 | if not same_shape: # pad/crop img 238 | gs = 32 # (pixels) grid size 239 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] 240 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean 241 | 242 | 243 | def copy_attr(a, b, include=(), exclude=()): 244 | # Copy attributes from b to a, options to only include [...] and to exclude [...] 245 | for k, v in b.__dict__.items(): 246 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 247 | continue 248 | else: 249 | setattr(a, k, v) 250 | 251 | 252 | class ModelEMA: 253 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 254 | Keep a moving average of everything in the model state_dict (parameters and buffers). 255 | This is intended to allow functionality like 256 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 257 | A smoothed version of the weights is necessary for some training schemes to perform well. 258 | This class is sensitive where it is initialized in the sequence of model init, 259 | GPU assignment and distributed training wrappers. 260 | """ 261 | 262 | def __init__(self, model, decay=0.9999, updates=0): 263 | # Create EMA 264 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA 265 | # if next(model.parameters()).device.type != 'cpu': 266 | # self.ema.half() # FP16 EMA 267 | self.updates = updates # number of EMA updates 268 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) 269 | for p in self.ema.parameters(): 270 | p.requires_grad_(False) 271 | 272 | def update(self, model): 273 | # Update EMA parameters 274 | with torch.no_grad(): 275 | self.updates += 1 276 | d = self.decay(self.updates) 277 | 278 | msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict 279 | for k, v in self.ema.state_dict().items(): 280 | if v.dtype.is_floating_point: 281 | v *= d 282 | v += (1. - d) * msd[k].detach() 283 | 284 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): 285 | # Update EMA attributes 286 | copy_attr(self.ema, model, include, exclude) 287 | -------------------------------------------------------------------------------- /yolov5/yololayer.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include "yololayer.h" 3 | #include "utils.h" 4 | 5 | using namespace Yolo; 6 | 7 | namespace nvinfer1 8 | { 9 | YoloLayerPlugin::YoloLayerPlugin(int classCount, int netWidth, int netHeight, int maxOut, const std::vector& vYoloKernel) 10 | { 11 | mClassCount = classCount; 12 | mYoloV5NetWidth = netWidth; 13 | mYoloV5NetHeight = netHeight; 14 | mMaxOutObject = maxOut; 15 | mYoloKernel = vYoloKernel; 16 | mKernelCount = vYoloKernel.size(); 17 | 18 | CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); 19 | size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; 20 | for (int ii = 0; ii < mKernelCount; ii++) 21 | { 22 | CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); 23 | const auto& yolo = mYoloKernel[ii]; 24 | CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); 25 | } 26 | } 27 | YoloLayerPlugin::~YoloLayerPlugin() 28 | { 29 | for (int ii = 0; ii < mKernelCount; ii++) 30 | { 31 | CUDA_CHECK(cudaFree(mAnchor[ii])); 32 | } 33 | CUDA_CHECK(cudaFreeHost(mAnchor)); 34 | } 35 | 36 | // create the plugin at runtime from a byte stream 37 | YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length) 38 | { 39 | using namespace Tn; 40 | const char *d = reinterpret_cast(data), *a = d; 41 | read(d, mClassCount); 42 | read(d, mThreadCount); 43 | read(d, mKernelCount); 44 | read(d, mYoloV5NetWidth); 45 | read(d, mYoloV5NetHeight); 46 | read(d, mMaxOutObject); 47 | mYoloKernel.resize(mKernelCount); 48 | auto kernelSize = mKernelCount * sizeof(YoloKernel); 49 | memcpy(mYoloKernel.data(), d, kernelSize); 50 | d += kernelSize; 51 | CUDA_CHECK(cudaMallocHost(&mAnchor, mKernelCount * sizeof(void*))); 52 | size_t AnchorLen = sizeof(float)* CHECK_COUNT * 2; 53 | for (int ii = 0; ii < mKernelCount; ii++) 54 | { 55 | CUDA_CHECK(cudaMalloc(&mAnchor[ii], AnchorLen)); 56 | const auto& yolo = mYoloKernel[ii]; 57 | CUDA_CHECK(cudaMemcpy(mAnchor[ii], yolo.anchors, AnchorLen, cudaMemcpyHostToDevice)); 58 | } 59 | assert(d == a + length); 60 | } 61 | 62 | void YoloLayerPlugin::serialize(void* buffer) const 63 | { 64 | using namespace Tn; 65 | char* d = static_cast(buffer), *a = d; 66 | write(d, mClassCount); 67 | write(d, mThreadCount); 68 | write(d, mKernelCount); 69 | write(d, mYoloV5NetWidth); 70 | write(d, mYoloV5NetHeight); 71 | write(d, mMaxOutObject); 72 | auto kernelSize = mKernelCount * sizeof(YoloKernel); 73 | memcpy(d, mYoloKernel.data(), kernelSize); 74 | d += kernelSize; 75 | 76 | assert(d == a + getSerializationSize()); 77 | } 78 | 79 | size_t YoloLayerPlugin::getSerializationSize() const 80 | { 81 | return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size() + sizeof(mYoloV5NetWidth) + sizeof(mYoloV5NetHeight) + sizeof(mMaxOutObject); 82 | } 83 | 84 | int YoloLayerPlugin::initialize() 85 | { 86 | return 0; 87 | } 88 | 89 | Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) 90 | { 91 | //output the result to channel 92 | int totalsize = mMaxOutObject * sizeof(Detection) / sizeof(float); 93 | 94 | return Dims3(totalsize + 1, 1, 1); 95 | } 96 | 97 | // Set plugin namespace 98 | void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) 99 | { 100 | mPluginNamespace = pluginNamespace; 101 | } 102 | 103 | const char* YoloLayerPlugin::getPluginNamespace() const 104 | { 105 | return mPluginNamespace; 106 | } 107 | 108 | // Return the DataType of the plugin output at the requested index 109 | DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const 110 | { 111 | return DataType::kFLOAT; 112 | } 113 | 114 | // Return true if output tensor is broadcast across a batch. 115 | bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const 116 | { 117 | return false; 118 | } 119 | 120 | // Return true if plugin can use input that is broadcast across batch without replication. 121 | bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const 122 | { 123 | return false; 124 | } 125 | 126 | void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) 127 | { 128 | } 129 | 130 | // Attach the plugin object to an execution context and grant the plugin the access to some context resource. 131 | void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) 132 | { 133 | } 134 | 135 | // Detach the plugin object from its execution context. 136 | void YoloLayerPlugin::detachFromContext() {} 137 | 138 | const char* YoloLayerPlugin::getPluginType() const 139 | { 140 | return "YoloLayer_TRT"; 141 | } 142 | 143 | const char* YoloLayerPlugin::getPluginVersion() const 144 | { 145 | return "1"; 146 | } 147 | 148 | void YoloLayerPlugin::destroy() 149 | { 150 | delete this; 151 | } 152 | 153 | // Clone the plugin 154 | IPluginV2IOExt* YoloLayerPlugin::clone() const 155 | { 156 | YoloLayerPlugin* p = new YoloLayerPlugin(mClassCount, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, mYoloKernel); 157 | p->setPluginNamespace(mPluginNamespace); 158 | return p; 159 | } 160 | 161 | __device__ float Logist(float data) { return 1.0f / (1.0f + expf(-data)); }; 162 | 163 | __global__ void CalDetection(const float *input, float *output, int noElements, 164 | const int netwidth, const int netheight, int maxoutobject, int yoloWidth, int yoloHeight, const float anchors[CHECK_COUNT * 2], int classes, int outputElem) 165 | { 166 | 167 | int idx = threadIdx.x + blockDim.x * blockIdx.x; 168 | if (idx >= noElements) return; 169 | 170 | int total_grid = yoloWidth * yoloHeight; 171 | int bnIdx = idx / total_grid; 172 | idx = idx - total_grid * bnIdx; 173 | int info_len_i = 5 + classes; 174 | const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT); 175 | 176 | for (int k = 0; k < 3; ++k) { 177 | float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); 178 | if (box_prob < IGNORE_THRESH) continue; 179 | int class_id = 0; 180 | float max_cls_prob = 0.0; 181 | for (int i = 5; i < info_len_i; ++i) { 182 | float p = Logist(curInput[idx + k * info_len_i * total_grid + i * total_grid]); 183 | if (p > max_cls_prob) { 184 | max_cls_prob = p; 185 | class_id = i - 5; 186 | } 187 | } 188 | float *res_count = output + bnIdx * outputElem; 189 | int count = (int)atomicAdd(res_count, 1); 190 | if (count >= maxoutobject) return; 191 | char* data = (char *)res_count + sizeof(float) + count * sizeof(Detection); 192 | Detection* det = (Detection*)(data); 193 | 194 | int row = idx / yoloWidth; 195 | int col = idx % yoloWidth; 196 | 197 | //Location 198 | // pytorch: 199 | // y = x[i].sigmoid() 200 | // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 201 | // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 202 | // X: (sigmoid(tx) + cx)/FeaturemapW * netwidth 203 | det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; 204 | det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; 205 | 206 | // W: (Pw * e^tw) / FeaturemapW * netwidth 207 | // v5: https://github.com/ultralytics/yolov5/issues/471 208 | det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); 209 | det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; 210 | det->bbox[3] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]); 211 | det->bbox[3] = det->bbox[3] * det->bbox[3] * anchors[2 * k + 1]; 212 | det->conf = box_prob * max_cls_prob; 213 | det->class_id = class_id; 214 | } 215 | } 216 | 217 | void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) 218 | { 219 | int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float); 220 | for (int idx = 0; idx < batchSize; ++idx) { 221 | CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float))); 222 | } 223 | int numElem = 0; 224 | for (unsigned int i = 0; i < mYoloKernel.size(); ++i) 225 | { 226 | const auto& yolo = mYoloKernel[i]; 227 | numElem = yolo.width*yolo.height*batchSize; 228 | if (numElem < mThreadCount) 229 | mThreadCount = numElem; 230 | 231 | //printf("Net: %d %d \n", mYoloV5NetWidth, mYoloV5NetHeight); 232 | CalDetection << < (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount >> > 233 | (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float *)mAnchor[i], mClassCount, outputElem); 234 | } 235 | } 236 | 237 | 238 | int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) 239 | { 240 | forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize); 241 | return 0; 242 | } 243 | 244 | PluginFieldCollection YoloPluginCreator::mFC{}; 245 | std::vector YoloPluginCreator::mPluginAttributes; 246 | 247 | YoloPluginCreator::YoloPluginCreator() 248 | { 249 | mPluginAttributes.clear(); 250 | 251 | mFC.nbFields = mPluginAttributes.size(); 252 | mFC.fields = mPluginAttributes.data(); 253 | } 254 | 255 | const char* YoloPluginCreator::getPluginName() const 256 | { 257 | return "YoloLayer_TRT"; 258 | } 259 | 260 | const char* YoloPluginCreator::getPluginVersion() const 261 | { 262 | return "1"; 263 | } 264 | 265 | const PluginFieldCollection* YoloPluginCreator::getFieldNames() 266 | { 267 | return &mFC; 268 | } 269 | 270 | IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) 271 | { 272 | int class_count = 80; 273 | int input_w = 416; 274 | int input_h = 416; 275 | int max_output_object_count = 1000; 276 | std::vector yolo_kernels(3); 277 | 278 | const PluginField* fields = fc->fields; 279 | for (int i = 0; i < fc->nbFields; i++) { 280 | if (strcmp(fields[i].name, "netdata") == 0) { 281 | assert(fields[i].type == PluginFieldType::kFLOAT32); 282 | int *tmp = (int*)(fields[i].data); 283 | class_count = tmp[0]; 284 | input_w = tmp[1]; 285 | input_h = tmp[2]; 286 | max_output_object_count = tmp[3]; 287 | } else if (strstr(fields[i].name, "yolodata") != NULL) { 288 | assert(fields[i].type == PluginFieldType::kFLOAT32); 289 | int *tmp = (int*)(fields[i].data); 290 | YoloKernel kernel; 291 | kernel.width = tmp[0]; 292 | kernel.height = tmp[1]; 293 | for (int j = 0; j < fields[i].length - 2; j++) { 294 | kernel.anchors[j] = tmp[j + 2]; 295 | } 296 | yolo_kernels[2 - (fields[i].name[8] - '1')] = kernel; 297 | } 298 | } 299 | YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, yolo_kernels); 300 | obj->setPluginNamespace(mNamespace.c_str()); 301 | return obj; 302 | } 303 | 304 | IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) 305 | { 306 | // This object will be deleted when the network is destroyed, which will 307 | // call YoloLayerPlugin::destroy() 308 | YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength); 309 | obj->setPluginNamespace(mNamespace.c_str()); 310 | return obj; 311 | } 312 | } 313 | -------------------------------------------------------------------------------- /yolov5/common.hpp: -------------------------------------------------------------------------------- 1 | #ifndef YOLOV5_COMMON_H_ 2 | #define YOLOV5_COMMON_H_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include "NvInfer.h" 11 | #include "yololayer.h" 12 | 13 | #define CHECK(status) \ 14 | do\ 15 | {\ 16 | auto ret = (status);\ 17 | if (ret != 0)\ 18 | {\ 19 | std::cerr << "Cuda failure: " << ret << std::endl;\ 20 | abort();\ 21 | }\ 22 | } while (0) 23 | 24 | using namespace nvinfer1; 25 | 26 | 27 | float iou(float lbox[4], float rbox[4]) { 28 | float interBox[] = { 29 | (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left 30 | (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right 31 | (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top 32 | (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom 33 | }; 34 | 35 | if (interBox[2] > interBox[3] || interBox[0] > interBox[1]) 36 | return 0.0f; 37 | 38 | float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]); 39 | return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS); 40 | } 41 | 42 | bool cmp(const Yolo::Detection& a, const Yolo::Detection& b) { 43 | return a.conf > b.conf; 44 | } 45 | 46 | void nms(std::vector& res, float *output, float conf_thresh, float nms_thresh = 0.5) { 47 | int det_size = sizeof(Yolo::Detection) / sizeof(float); 48 | std::map> m; 49 | for (int i = 0; i < output[0] && i < Yolo::MAX_OUTPUT_BBOX_COUNT; i++) { 50 | if (output[1 + det_size * i + 4] <= conf_thresh) continue; 51 | Yolo::Detection det; 52 | memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float)); 53 | if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); 54 | m[det.class_id].push_back(det); 55 | } 56 | for (auto it = m.begin(); it != m.end(); it++) { 57 | //std::cout << it->second[0].class_id << " --- " << std::endl; 58 | auto& dets = it->second; 59 | std::sort(dets.begin(), dets.end(), cmp); 60 | for (size_t m = 0; m < dets.size(); ++m) { 61 | auto& item = dets[m]; 62 | res.push_back(item); 63 | for (size_t n = m + 1; n < dets.size(); ++n) { 64 | if (iou(item.bbox, dets[n].bbox) > nms_thresh) { 65 | dets.erase(dets.begin() + n); 66 | --n; 67 | } 68 | } 69 | } 70 | } 71 | } 72 | 73 | // TensorRT weight files have a simple space delimited format: 74 | // [type] [size] 75 | std::map loadWeights(const std::string file) { 76 | std::cout << "Loading weights: " << file << std::endl; 77 | std::map weightMap; 78 | 79 | // Open weights file 80 | std::ifstream input(file); 81 | assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!"); 82 | 83 | // Read number of weight blobs 84 | int32_t count; 85 | input >> count; 86 | assert(count > 0 && "Invalid weight map file."); 87 | 88 | while (count--) 89 | { 90 | Weights wt{ DataType::kFLOAT, nullptr, 0 }; 91 | uint32_t size; 92 | 93 | // Read name and type of blob 94 | std::string name; 95 | input >> name >> std::dec >> size; 96 | wt.type = DataType::kFLOAT; 97 | 98 | // Load blob 99 | uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); 100 | for (uint32_t x = 0, y = size; x < y; ++x) 101 | { 102 | input >> std::hex >> val[x]; 103 | } 104 | wt.values = val; 105 | 106 | wt.count = size; 107 | weightMap[name] = wt; 108 | } 109 | 110 | return weightMap; 111 | } 112 | 113 | IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::string lname, float eps) { 114 | float *gamma = (float*)weightMap[lname + ".weight"].values; 115 | float *beta = (float*)weightMap[lname + ".bias"].values; 116 | float *mean = (float*)weightMap[lname + ".running_mean"].values; 117 | float *var = (float*)weightMap[lname + ".running_var"].values; 118 | int len = weightMap[lname + ".running_var"].count; 119 | 120 | float *scval = reinterpret_cast(malloc(sizeof(float) * len)); 121 | for (int i = 0; i < len; i++) { 122 | scval[i] = gamma[i] / sqrt(var[i] + eps); 123 | } 124 | Weights scale{ DataType::kFLOAT, scval, len }; 125 | 126 | float *shval = reinterpret_cast(malloc(sizeof(float) * len)); 127 | for (int i = 0; i < len; i++) { 128 | shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); 129 | } 130 | Weights shift{ DataType::kFLOAT, shval, len }; 131 | 132 | float *pval = reinterpret_cast(malloc(sizeof(float) * len)); 133 | for (int i = 0; i < len; i++) { 134 | pval[i] = 1.0; 135 | } 136 | Weights power{ DataType::kFLOAT, pval, len }; 137 | 138 | weightMap[lname + ".scale"] = scale; 139 | weightMap[lname + ".shift"] = shift; 140 | weightMap[lname + ".power"] = power; 141 | IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); 142 | assert(scale_1); 143 | return scale_1; 144 | } 145 | 146 | ILayer* convBlock(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int g, std::string lname) { 147 | Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; 148 | int p = ksize / 2; 149 | IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[lname + ".conv.weight"], emptywts); 150 | assert(conv1); 151 | conv1->setStrideNd(DimsHW{ s, s }); 152 | conv1->setPaddingNd(DimsHW{ p, p }); 153 | conv1->setNbGroups(g); 154 | IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn", 1e-3); 155 | 156 | // hard_swish = x * hard_sigmoid 157 | auto hsig = network->addActivation(*bn1->getOutput(0), ActivationType::kHARD_SIGMOID); 158 | assert(hsig); 159 | hsig->setAlpha(1.0 / 6.0); 160 | hsig->setBeta(0.5); 161 | auto ew = network->addElementWise(*bn1->getOutput(0), *hsig->getOutput(0), ElementWiseOperation::kPROD); 162 | assert(ew); 163 | return ew; 164 | } 165 | 166 | ILayer* focus(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int ksize, std::string lname) { 167 | ISliceLayer *s1 = network->addSlice(input, Dims3{ 0, 0, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 168 | ISliceLayer *s2 = network->addSlice(input, Dims3{ 0, 1, 0 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 169 | ISliceLayer *s3 = network->addSlice(input, Dims3{ 0, 0, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 170 | ISliceLayer *s4 = network->addSlice(input, Dims3{ 0, 1, 1 }, Dims3{ inch, Yolo::INPUT_H / 2, Yolo::INPUT_W / 2 }, Dims3{ 1, 2, 2 }); 171 | ITensor* inputTensors[] = { s1->getOutput(0), s2->getOutput(0), s3->getOutput(0), s4->getOutput(0) }; 172 | auto cat = network->addConcatenation(inputTensors, 4); 173 | auto conv = convBlock(network, weightMap, *cat->getOutput(0), outch, ksize, 1, 1, lname + ".conv"); 174 | return conv; 175 | } 176 | 177 | ILayer* bottleneck(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, bool shortcut, int g, float e, std::string lname) { 178 | auto cv1 = convBlock(network, weightMap, input, (int)((float)c2 * e), 1, 1, 1, lname + ".cv1"); 179 | auto cv2 = convBlock(network, weightMap, *cv1->getOutput(0), c2, 3, 1, g, lname + ".cv2"); 180 | if (shortcut && c1 == c2) { 181 | auto ew = network->addElementWise(input, *cv2->getOutput(0), ElementWiseOperation::kSUM); 182 | return ew; 183 | } 184 | return cv2; 185 | } 186 | 187 | ILayer* bottleneckCSP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { 188 | Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; 189 | int c_ = (int)((float)c2 * e); 190 | auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); 191 | auto cv2 = network->addConvolutionNd(input, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv2.weight"], emptywts); 192 | ITensor *y1 = cv1->getOutput(0); 193 | for (int i = 0; i < n; i++) { 194 | auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); 195 | y1 = b->getOutput(0); 196 | } 197 | auto cv3 = network->addConvolutionNd(*y1, c_, DimsHW{ 1, 1 }, weightMap[lname + ".cv3.weight"], emptywts); 198 | 199 | ITensor* inputTensors[] = { cv3->getOutput(0), cv2->getOutput(0) }; 200 | auto cat = network->addConcatenation(inputTensors, 2); 201 | 202 | IScaleLayer* bn = addBatchNorm2d(network, weightMap, *cat->getOutput(0), lname + ".bn", 1e-4); 203 | auto lr = network->addActivation(*bn->getOutput(0), ActivationType::kLEAKY_RELU); 204 | lr->setAlpha(0.1); 205 | 206 | auto cv4 = convBlock(network, weightMap, *lr->getOutput(0), c2, 1, 1, 1, lname + ".cv4"); 207 | return cv4; 208 | } 209 | 210 | ILayer* SPP(INetworkDefinition *network, std::map& weightMap, ITensor& input, int c1, int c2, int k1, int k2, int k3, std::string lname) { 211 | int c_ = c1 / 2; 212 | auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); 213 | 214 | auto pool1 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k1, k1 }); 215 | pool1->setPaddingNd(DimsHW{ k1 / 2, k1 / 2 }); 216 | pool1->setStrideNd(DimsHW{ 1, 1 }); 217 | auto pool2 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k2, k2 }); 218 | pool2->setPaddingNd(DimsHW{ k2 / 2, k2 / 2 }); 219 | pool2->setStrideNd(DimsHW{ 1, 1 }); 220 | auto pool3 = network->addPoolingNd(*cv1->getOutput(0), PoolingType::kMAX, DimsHW{ k3, k3 }); 221 | pool3->setPaddingNd(DimsHW{ k3 / 2, k3 / 2 }); 222 | pool3->setStrideNd(DimsHW{ 1, 1 }); 223 | 224 | ITensor* inputTensors[] = { cv1->getOutput(0), pool1->getOutput(0), pool2->getOutput(0), pool3->getOutput(0) }; 225 | auto cat = network->addConcatenation(inputTensors, 4); 226 | 227 | auto cv2 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv2"); 228 | return cv2; 229 | } 230 | 231 | ILayer* C3(INetworkDefinition* network, std::map& weightMap, ITensor& input, int c1, int c2, int n, bool shortcut, int g, float e, std::string lname) { 232 | int c_ = (int)((float)c2 * e); 233 | auto cv1 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv1"); 234 | auto cv2 = convBlock(network, weightMap, input, c_, 1, 1, 1, lname + ".cv2"); 235 | ITensor* y1 = cv1->getOutput(0); 236 | for (int i = 0; i < n; i++) { 237 | auto b = bottleneck(network, weightMap, *y1, c_, c_, shortcut, g, 1.0, lname + ".m." + std::to_string(i)); 238 | y1 = b->getOutput(0); 239 | } 240 | 241 | ITensor* inputTensors[] = { y1, cv2->getOutput(0) }; 242 | auto cat = network->addConcatenation(inputTensors, 2); 243 | 244 | auto cv3 = convBlock(network, weightMap, *cat->getOutput(0), c2, 1, 1, 1, lname + ".cv3"); 245 | return cv3; 246 | } 247 | 248 | 249 | int read_files_in_dir(const char *p_dir_name, std::vector &file_names) { 250 | DIR *p_dir = opendir(p_dir_name); 251 | if (p_dir == nullptr) { 252 | return -1; 253 | } 254 | 255 | struct dirent* p_file = nullptr; 256 | while ((p_file = readdir(p_dir)) != nullptr) { 257 | if (strcmp(p_file->d_name, ".") != 0 && 258 | strcmp(p_file->d_name, "..") != 0) { 259 | //std::string cur_file_name(p_dir_name); 260 | //cur_file_name += "/"; 261 | //cur_file_name += p_file->d_name; 262 | std::string cur_file_name(p_file->d_name); 263 | file_names.push_back(cur_file_name); 264 | } 265 | } 266 | 267 | closedir(p_dir); 268 | return 0; 269 | } 270 | 271 | std::vector getAnchors(std::map& weightMap) 272 | { 273 | std::vector anchors_yolo; 274 | Weights Yolo_Anchors = weightMap["model.24.anchor_grid"]; 275 | assert(Yolo_Anchors.count == 18); 276 | int each_yololayer_anchorsnum = Yolo_Anchors.count / 3; 277 | const float* tempAnchors = (const float*)(Yolo_Anchors.values); 278 | for (int i = 0; i < Yolo_Anchors.count; i++) 279 | { 280 | if (i < each_yololayer_anchorsnum) 281 | { 282 | anchors_yolo.push_back(const_cast(tempAnchors)[i]); 283 | } 284 | if ((i >= each_yololayer_anchorsnum) && (i < (2 * each_yololayer_anchorsnum))) 285 | { 286 | anchors_yolo.push_back(const_cast(tempAnchors)[i]); 287 | } 288 | if (i >= (2 * each_yololayer_anchorsnum)) 289 | { 290 | anchors_yolo.push_back(const_cast(tempAnchors)[i]); 291 | } 292 | } 293 | return anchors_yolo; 294 | } 295 | 296 | IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, IConvolutionLayer* det0, IConvolutionLayer* det1, IConvolutionLayer* det2) 297 | { 298 | auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); 299 | std::vector anchors_yolo = getAnchors(weightMap); 300 | PluginField pluginMultidata[4]; 301 | int NetData[4]; 302 | NetData[0] = Yolo::CLASS_NUM; 303 | NetData[1] = Yolo::INPUT_W; 304 | NetData[2] = Yolo::INPUT_H; 305 | NetData[3] = Yolo::MAX_OUTPUT_BBOX_COUNT; 306 | pluginMultidata[0].data = NetData; 307 | pluginMultidata[0].length = 3; 308 | pluginMultidata[0].name = "netdata"; 309 | pluginMultidata[0].type = PluginFieldType::kFLOAT32; 310 | int scale[3] = { 8, 16, 32 }; 311 | int plugindata[3][8]; 312 | std::string names[3]; 313 | for (int k = 1; k < 4; k++) 314 | { 315 | plugindata[k - 1][0] = Yolo::INPUT_W / scale[k - 1]; 316 | plugindata[k - 1][1] = Yolo::INPUT_H / scale[k - 1]; 317 | for (int i = 2; i < 8; i++) 318 | { 319 | plugindata[k - 1][i] = int(anchors_yolo[(k - 1) * 6 + i - 2]); 320 | } 321 | pluginMultidata[k].data = plugindata[k - 1]; 322 | pluginMultidata[k].length = 8; 323 | names[k - 1] = "yolodata" + std::to_string(k); 324 | pluginMultidata[k].name = names[k - 1].c_str(); 325 | pluginMultidata[k].type = PluginFieldType::kFLOAT32; 326 | } 327 | PluginFieldCollection pluginData; 328 | pluginData.nbFields = 4; 329 | pluginData.fields = pluginMultidata; 330 | IPluginV2 *pluginObj = creator->createPlugin("yololayer", &pluginData); 331 | ITensor* inputTensors_yolo[] = { det2->getOutput(0), det1->getOutput(0), det0->getOutput(0) }; 332 | auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); 333 | return yolo; 334 | } 335 | #endif 336 | 337 | -------------------------------------------------------------------------------- /sort/include/munkres.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2007 John Weaver 3 | * Copyright (c) 2015 Miroslav Krajicek 4 | * 5 | * This program is free software; you can redistribute it and/or modify 6 | * it under the terms of the GNU General Public License as published by 7 | * the Free Software Foundation; either version 2 of the License, or 8 | * (at your option) any later version. 9 | * 10 | * This program is distributed in the hope that it will be useful, 11 | * but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 | * GNU General Public License for more details. 14 | * 15 | * You should have received a copy of the GNU General Public License 16 | * along with this program; if not, write to the Free Software 17 | * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 18 | */ 19 | 20 | #if !defined(_MUNKRES_H_) 21 | #define _MUNKRES_H_ 22 | 23 | #include "matrix.h" 24 | 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | 31 | template class Munkres 32 | { 33 | static constexpr int NORMAL = 0; 34 | static constexpr int STAR = 1; 35 | static constexpr int PRIME = 2; 36 | public: 37 | 38 | /* 39 | * 40 | * Linear assignment problem solution 41 | * [modifies matrix in-place.] 42 | * matrix(row,col): row major format assumed. 43 | * 44 | * Assignments are remaining 0 values 45 | * (extra 0 values are replaced with -1) 46 | * 47 | */ 48 | void solve(Matrix &m) { 49 | const size_t rows = m.rows(), 50 | columns = m.columns(), 51 | size = std::max(rows, columns); 52 | 53 | #ifdef DEBUG 54 | std::cout << "Munkres input: " << m << std::endl; 55 | #endif 56 | 57 | // Copy input matrix 58 | this->matrix = m; 59 | 60 | if ( rows != columns ) { 61 | // If the input matrix isn't square, make it square 62 | // and fill the empty values with the largest value present 63 | // in the matrix. 64 | matrix.resize(size, size, matrix.max()); 65 | } 66 | 67 | 68 | // STAR == 1 == starred, PRIME == 2 == primed 69 | mask_matrix.resize(size, size); 70 | 71 | row_mask = new bool[size]; 72 | col_mask = new bool[size]; 73 | for ( size_t i = 0 ; i < size ; i++ ) { 74 | row_mask[i] = false; 75 | } 76 | 77 | for ( size_t i = 0 ; i < size ; i++ ) { 78 | col_mask[i] = false; 79 | } 80 | 81 | // Prepare the matrix values... 82 | 83 | // If there were any infinities, replace them with a value greater 84 | // than the maximum value in the matrix. 85 | replace_infinites(matrix); 86 | 87 | minimize_along_direction(matrix, rows >= columns); 88 | minimize_along_direction(matrix, rows < columns); 89 | 90 | // Follow the steps 91 | int step = 1; 92 | while ( step ) { 93 | switch ( step ) { 94 | case 1: 95 | step = step1(); 96 | // step is always 2 97 | break; 98 | case 2: 99 | step = step2(); 100 | // step is always either 0 or 3 101 | break; 102 | case 3: 103 | step = step3(); 104 | // step in [3, 4, 5] 105 | break; 106 | case 4: 107 | step = step4(); 108 | // step is always 2 109 | break; 110 | case 5: 111 | step = step5(); 112 | // step is always 3 113 | break; 114 | } 115 | } 116 | 117 | // Store results 118 | for ( size_t row = 0 ; row < size ; row++ ) { 119 | for ( size_t col = 0 ; col < size ; col++ ) { 120 | if ( mask_matrix(row, col) == STAR ) { 121 | matrix(row, col) = 0; 122 | } else { 123 | matrix(row, col) = -1; 124 | } 125 | } 126 | } 127 | 128 | #ifdef DEBUG 129 | std::cout << "Munkres output: " << matrix << std::endl; 130 | #endif 131 | // Remove the excess rows or columns that we added to fit the 132 | // input to a square matrix. 133 | matrix.resize(rows, columns); 134 | 135 | m = matrix; 136 | 137 | delete [] row_mask; 138 | delete [] col_mask; 139 | } 140 | 141 | static void replace_infinites(Matrix &matrix) { 142 | const size_t rows = matrix.rows(), 143 | columns = matrix.columns(); 144 | assert( rows > 0 && columns > 0 ); 145 | double max = matrix(0, 0); 146 | constexpr auto infinity = std::numeric_limits::infinity(); 147 | 148 | // Find the greatest value in the matrix that isn't infinity. 149 | for ( size_t row = 0 ; row < rows ; row++ ) { 150 | for ( size_t col = 0 ; col < columns ; col++ ) { 151 | if ( matrix(row, col) != infinity ) { 152 | if ( max == infinity ) { 153 | max = matrix(row, col); 154 | } else { 155 | max = std::max(max, matrix(row, col)); 156 | } 157 | } 158 | } 159 | } 160 | 161 | // a value higher than the maximum value present in the matrix. 162 | if ( max == infinity ) { 163 | // This case only occurs when all values are infinite. 164 | max = 0; 165 | } else { 166 | max++; 167 | } 168 | 169 | for ( size_t row = 0 ; row < rows ; row++ ) { 170 | for ( size_t col = 0 ; col < columns ; col++ ) { 171 | if ( matrix(row, col) == infinity ) { 172 | matrix(row, col) = max; 173 | } 174 | } 175 | } 176 | 177 | } 178 | 179 | static void minimize_along_direction(Matrix &matrix, const bool over_columns) { 180 | const size_t outer_size = over_columns ? matrix.columns() : matrix.rows(), 181 | inner_size = over_columns ? matrix.rows() : matrix.columns(); 182 | 183 | // Look for a minimum value to subtract from all values along 184 | // the "outer" direction. 185 | for ( size_t i = 0 ; i < outer_size ; i++ ) { 186 | double min = over_columns ? matrix(0, i) : matrix(i, 0); 187 | 188 | // As long as the current minimum is greater than zero, 189 | // keep looking for the minimum. 190 | // Start at one because we already have the 0th value in min. 191 | for ( size_t j = 1 ; j < inner_size && min > 0 ; j++ ) { 192 | min = std::min( 193 | min, 194 | over_columns ? matrix(j, i) : matrix(i, j)); 195 | } 196 | 197 | if ( min > 0 ) { 198 | for ( size_t j = 0 ; j < inner_size ; j++ ) { 199 | if ( over_columns ) { 200 | matrix(j, i) -= min; 201 | } else { 202 | matrix(i, j) -= min; 203 | } 204 | } 205 | } 206 | } 207 | } 208 | 209 | private: 210 | 211 | inline bool find_uncovered_in_matrix(const double item, size_t &row, size_t &col) const { 212 | const size_t rows = matrix.rows(), 213 | columns = matrix.columns(); 214 | 215 | for ( row = 0 ; row < rows ; row++ ) { 216 | if ( !row_mask[row] ) { 217 | for ( col = 0 ; col < columns ; col++ ) { 218 | if ( !col_mask[col] ) { 219 | if ( matrix(row,col) == item ) { 220 | return true; 221 | } 222 | } 223 | } 224 | } 225 | } 226 | 227 | return false; 228 | } 229 | 230 | bool pair_in_list(const std::pair &needle, const std::list > &haystack) { 231 | for ( std::list >::const_iterator i = haystack.begin() ; i != haystack.end() ; i++ ) { 232 | if ( needle == *i ) { 233 | return true; 234 | } 235 | } 236 | 237 | return false; 238 | } 239 | 240 | int step1() { 241 | const size_t rows = matrix.rows(), 242 | columns = matrix.columns(); 243 | 244 | for ( size_t row = 0 ; row < rows ; row++ ) { 245 | for ( size_t col = 0 ; col < columns ; col++ ) { 246 | if ( 0 == matrix(row, col) ) { 247 | for ( size_t nrow = 0 ; nrow < row ; nrow++ ) 248 | if ( STAR == mask_matrix(nrow,col) ) 249 | goto next_column; 250 | 251 | mask_matrix(row,col) = STAR; 252 | goto next_row; 253 | } 254 | next_column:; 255 | } 256 | next_row:; 257 | } 258 | 259 | return 2; 260 | } 261 | 262 | int step2() { 263 | const size_t rows = matrix.rows(), 264 | columns = matrix.columns(); 265 | size_t covercount = 0; 266 | 267 | for ( size_t row = 0 ; row < rows ; row++ ) 268 | for ( size_t col = 0 ; col < columns ; col++ ) 269 | if ( STAR == mask_matrix(row, col) ) { 270 | col_mask[col] = true; 271 | covercount++; 272 | } 273 | 274 | if ( covercount >= matrix.minsize() ) { 275 | #ifdef DEBUG 276 | std::cout << "Final cover count: " << covercount << std::endl; 277 | #endif 278 | return 0; 279 | } 280 | 281 | #ifdef DEBUG 282 | std::cout << "Munkres matrix has " << covercount << " of " << matrix.minsize() << " Columns covered:" << std::endl; 283 | std::cout << matrix << std::endl; 284 | #endif 285 | 286 | 287 | return 3; 288 | } 289 | 290 | int step3() { 291 | /* 292 | Main Zero Search 293 | 294 | 1. Find an uncovered Z in the distance matrix and prime it. If no such zero exists, go to Step 5 295 | 2. If No Z* exists in the row of the Z', go to Step 4. 296 | 3. If a Z* exists, cover this row and uncover the column of the Z*. Return to Step 3.1 to find a new Z 297 | */ 298 | if ( find_uncovered_in_matrix(0, saverow, savecol) ) { 299 | mask_matrix(saverow,savecol) = PRIME; // prime it. 300 | } else { 301 | return 5; 302 | } 303 | 304 | for ( size_t ncol = 0 ; ncol < matrix.columns() ; ncol++ ) { 305 | if ( mask_matrix(saverow,ncol) == STAR ) { 306 | row_mask[saverow] = true; //cover this row and 307 | col_mask[ncol] = false; // uncover the column containing the starred zero 308 | return 3; // repeat 309 | } 310 | } 311 | 312 | return 4; // no starred zero in the row containing this primed zero 313 | } 314 | 315 | int step4() { 316 | const size_t rows = matrix.rows(), 317 | columns = matrix.columns(); 318 | 319 | // seq contains pairs of row/column values where we have found 320 | // either a star or a prime that is part of the ``alternating sequence``. 321 | std::list > seq; 322 | // use saverow, savecol from step 3. 323 | std::pair z0(saverow, savecol); 324 | seq.insert(seq.end(), z0); 325 | 326 | // We have to find these two pairs: 327 | std::pair z1(-1, -1); 328 | std::pair z2n(-1, -1); 329 | 330 | size_t row, col = savecol; 331 | /* 332 | Increment Set of Starred Zeros 333 | 334 | 1. Construct the ``alternating sequence'' of primed and starred zeros: 335 | 336 | Z0 : Unpaired Z' from Step 4.2 337 | Z1 : The Z* in the column of Z0 338 | Z[2N] : The Z' in the row of Z[2N-1], if such a zero exists 339 | Z[2N+1] : The Z* in the column of Z[2N] 340 | 341 | The sequence eventually terminates with an unpaired Z' = Z[2N] for some N. 342 | */ 343 | bool madepair; 344 | do { 345 | madepair = false; 346 | for ( row = 0 ; row < rows ; row++ ) { 347 | if ( mask_matrix(row,col) == STAR ) { 348 | z1.first = row; 349 | z1.second = col; 350 | if ( pair_in_list(z1, seq) ) { 351 | continue; 352 | } 353 | 354 | madepair = true; 355 | seq.insert(seq.end(), z1); 356 | break; 357 | } 358 | } 359 | 360 | if ( !madepair ) 361 | break; 362 | 363 | madepair = false; 364 | 365 | for ( col = 0 ; col < columns ; col++ ) { 366 | if ( mask_matrix(row, col) == PRIME ) { 367 | z2n.first = row; 368 | z2n.second = col; 369 | if ( pair_in_list(z2n, seq) ) { 370 | continue; 371 | } 372 | madepair = true; 373 | seq.insert(seq.end(), z2n); 374 | break; 375 | } 376 | } 377 | } while ( madepair ); 378 | 379 | for ( std::list >::iterator i = seq.begin() ; 380 | i != seq.end() ; 381 | i++ ) { 382 | // 2. Unstar each starred zero of the sequence. 383 | if ( mask_matrix(i->first,i->second) == STAR ) 384 | mask_matrix(i->first,i->second) = NORMAL; 385 | 386 | // 3. Star each primed zero of the sequence, 387 | // thus increasing the number of starred zeros by one. 388 | if ( mask_matrix(i->first,i->second) == PRIME ) 389 | mask_matrix(i->first,i->second) = STAR; 390 | } 391 | 392 | // 4. Erase all primes, uncover all columns and rows, 393 | for ( size_t row = 0 ; row < mask_matrix.rows() ; row++ ) { 394 | for ( size_t col = 0 ; col < mask_matrix.columns() ; col++ ) { 395 | if ( mask_matrix(row,col) == PRIME ) { 396 | mask_matrix(row,col) = NORMAL; 397 | } 398 | } 399 | } 400 | 401 | for ( size_t i = 0 ; i < rows ; i++ ) { 402 | row_mask[i] = false; 403 | } 404 | 405 | for ( size_t i = 0 ; i < columns ; i++ ) { 406 | col_mask[i] = false; 407 | } 408 | 409 | // and return to Step 2. 410 | return 2; 411 | } 412 | 413 | int step5() { 414 | const size_t rows = matrix.rows(), 415 | columns = matrix.columns(); 416 | /* 417 | New Zero Manufactures 418 | 419 | 1. Let h be the smallest uncovered entry in the (modified) distance matrix. 420 | 2. Add h to all covered rows. 421 | 3. Subtract h from all uncovered columns 422 | 4. Return to Step 3, without altering stars, primes, or covers. 423 | */ 424 | double h = std::numeric_limits::max(); 425 | for ( size_t row = 0 ; row < rows ; row++ ) { 426 | if ( !row_mask[row] ) { 427 | for ( size_t col = 0 ; col < columns ; col++ ) { 428 | if ( !col_mask[col] ) { 429 | if ( h > matrix(row, col) && matrix(row, col) != 0 ) { 430 | h = matrix(row, col); 431 | } 432 | } 433 | } 434 | } 435 | } 436 | 437 | for ( size_t row = 0 ; row < rows ; row++ ) { 438 | if ( row_mask[row] ) { 439 | for ( size_t col = 0 ; col < columns ; col++ ) { 440 | matrix(row, col) += h; 441 | } 442 | } 443 | } 444 | 445 | for ( size_t col = 0 ; col < columns ; col++ ) { 446 | if ( !col_mask[col] ) { 447 | for ( size_t row = 0 ; row < rows ; row++ ) { 448 | matrix(row, col) -= h; 449 | } 450 | } 451 | } 452 | 453 | return 3; 454 | } 455 | 456 | Matrix mask_matrix; 457 | Matrix matrix; 458 | bool *row_mask; 459 | bool *col_mask; 460 | size_t saverow = 0, savecol = 0; 461 | }; 462 | 463 | 464 | #endif /* !defined(_MUNKRES_H_) */ 465 | -------------------------------------------------------------------------------- /yolov5/logging.h: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | #ifndef TENSORRT_LOGGING_H 18 | #define TENSORRT_LOGGING_H 19 | 20 | #include "NvInferRuntimeCommon.h" 21 | #include 22 | #include 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | 29 | using Severity = nvinfer1::ILogger::Severity; 30 | 31 | class LogStreamConsumerBuffer : public std::stringbuf 32 | { 33 | public: 34 | LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) 35 | : mOutput(stream) 36 | , mPrefix(prefix) 37 | , mShouldLog(shouldLog) 38 | { 39 | } 40 | 41 | LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) 42 | : mOutput(other.mOutput) 43 | { 44 | } 45 | 46 | ~LogStreamConsumerBuffer() 47 | { 48 | // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence 49 | // std::streambuf::pptr() gives a pointer to the current position of the output sequence 50 | // if the pointer to the beginning is not equal to the pointer to the current position, 51 | // call putOutput() to log the output to the stream 52 | if (pbase() != pptr()) 53 | { 54 | putOutput(); 55 | } 56 | } 57 | 58 | // synchronizes the stream buffer and returns 0 on success 59 | // synchronizing the stream buffer consists of inserting the buffer contents into the stream, 60 | // resetting the buffer and flushing the stream 61 | virtual int sync() 62 | { 63 | putOutput(); 64 | return 0; 65 | } 66 | 67 | void putOutput() 68 | { 69 | if (mShouldLog) 70 | { 71 | // prepend timestamp 72 | std::time_t timestamp = std::time(nullptr); 73 | tm* tm_local = std::localtime(×tamp); 74 | std::cout << "["; 75 | std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; 76 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; 77 | std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; 78 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; 79 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; 80 | std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; 81 | // std::stringbuf::str() gets the string contents of the buffer 82 | // insert the buffer contents pre-appended by the appropriate prefix into the stream 83 | mOutput << mPrefix << str(); 84 | // set the buffer to empty 85 | str(""); 86 | // flush the stream 87 | mOutput.flush(); 88 | } 89 | } 90 | 91 | void setShouldLog(bool shouldLog) 92 | { 93 | mShouldLog = shouldLog; 94 | } 95 | 96 | private: 97 | std::ostream& mOutput; 98 | std::string mPrefix; 99 | bool mShouldLog; 100 | }; 101 | 102 | //! 103 | //! \class LogStreamConsumerBase 104 | //! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer 105 | //! 106 | class LogStreamConsumerBase 107 | { 108 | public: 109 | LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) 110 | : mBuffer(stream, prefix, shouldLog) 111 | { 112 | } 113 | 114 | protected: 115 | LogStreamConsumerBuffer mBuffer; 116 | }; 117 | 118 | //! 119 | //! \class LogStreamConsumer 120 | //! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. 121 | //! Order of base classes is LogStreamConsumerBase and then std::ostream. 122 | //! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field 123 | //! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. 124 | //! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. 125 | //! Please do not change the order of the parent classes. 126 | //! 127 | class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream 128 | { 129 | public: 130 | //! \brief Creates a LogStreamConsumer which logs messages with level severity. 131 | //! Reportable severity determines if the messages are severe enough to be logged. 132 | LogStreamConsumer(Severity reportableSeverity, Severity severity) 133 | : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) 134 | , std::ostream(&mBuffer) // links the stream buffer with the stream 135 | , mShouldLog(severity <= reportableSeverity) 136 | , mSeverity(severity) 137 | { 138 | } 139 | 140 | LogStreamConsumer(LogStreamConsumer&& other) 141 | : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) 142 | , std::ostream(&mBuffer) // links the stream buffer with the stream 143 | , mShouldLog(other.mShouldLog) 144 | , mSeverity(other.mSeverity) 145 | { 146 | } 147 | 148 | void setReportableSeverity(Severity reportableSeverity) 149 | { 150 | mShouldLog = mSeverity <= reportableSeverity; 151 | mBuffer.setShouldLog(mShouldLog); 152 | } 153 | 154 | private: 155 | static std::ostream& severityOstream(Severity severity) 156 | { 157 | return severity >= Severity::kINFO ? std::cout : std::cerr; 158 | } 159 | 160 | static std::string severityPrefix(Severity severity) 161 | { 162 | switch (severity) 163 | { 164 | case Severity::kINTERNAL_ERROR: return "[F] "; 165 | case Severity::kERROR: return "[E] "; 166 | case Severity::kWARNING: return "[W] "; 167 | case Severity::kINFO: return "[I] "; 168 | case Severity::kVERBOSE: return "[V] "; 169 | default: assert(0); return ""; 170 | } 171 | } 172 | 173 | bool mShouldLog; 174 | Severity mSeverity; 175 | }; 176 | 177 | //! \class Logger 178 | //! 179 | //! \brief Class which manages logging of TensorRT tools and samples 180 | //! 181 | //! \details This class provides a common interface for TensorRT tools and samples to log information to the console, 182 | //! and supports logging two types of messages: 183 | //! 184 | //! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) 185 | //! - Test pass/fail messages 186 | //! 187 | //! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is 188 | //! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. 189 | //! 190 | //! In the future, this class could be extended to support dumping test results to a file in some standard format 191 | //! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). 192 | //! 193 | //! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger 194 | //! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT 195 | //! library and messages coming from the sample. 196 | //! 197 | //! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the 198 | //! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger 199 | //! object. 200 | 201 | class Logger : public nvinfer1::ILogger 202 | { 203 | public: 204 | Logger(Severity severity = Severity::kWARNING) 205 | : mReportableSeverity(severity) 206 | { 207 | } 208 | 209 | //! 210 | //! \enum TestResult 211 | //! \brief Represents the state of a given test 212 | //! 213 | enum class TestResult 214 | { 215 | kRUNNING, //!< The test is running 216 | kPASSED, //!< The test passed 217 | kFAILED, //!< The test failed 218 | kWAIVED //!< The test was waived 219 | }; 220 | 221 | //! 222 | //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger 223 | //! \return The nvinfer1::ILogger associated with this Logger 224 | //! 225 | //! TODO Once all samples are updated to use this method to register the logger with TensorRT, 226 | //! we can eliminate the inheritance of Logger from ILogger 227 | //! 228 | nvinfer1::ILogger& getTRTLogger() 229 | { 230 | return *this; 231 | } 232 | 233 | //! 234 | //! \brief Implementation of the nvinfer1::ILogger::log() virtual method 235 | //! 236 | //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the 237 | //! inheritance from nvinfer1::ILogger 238 | //! 239 | void log(Severity severity, const char* msg) override 240 | { 241 | LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; 242 | } 243 | 244 | //! 245 | //! \brief Method for controlling the verbosity of logging output 246 | //! 247 | //! \param severity The logger will only emit messages that have severity of this level or higher. 248 | //! 249 | void setReportableSeverity(Severity severity) 250 | { 251 | mReportableSeverity = severity; 252 | } 253 | 254 | //! 255 | //! \brief Opaque handle that holds logging information for a particular test 256 | //! 257 | //! This object is an opaque handle to information used by the Logger to print test results. 258 | //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used 259 | //! with Logger::reportTest{Start,End}(). 260 | //! 261 | class TestAtom 262 | { 263 | public: 264 | TestAtom(TestAtom&&) = default; 265 | 266 | private: 267 | friend class Logger; 268 | 269 | TestAtom(bool started, const std::string& name, const std::string& cmdline) 270 | : mStarted(started) 271 | , mName(name) 272 | , mCmdline(cmdline) 273 | { 274 | } 275 | 276 | bool mStarted; 277 | std::string mName; 278 | std::string mCmdline; 279 | }; 280 | 281 | //! 282 | //! \brief Define a test for logging 283 | //! 284 | //! \param[in] name The name of the test. This should be a string starting with 285 | //! "TensorRT" and containing dot-separated strings containing 286 | //! the characters [A-Za-z0-9_]. 287 | //! For example, "TensorRT.sample_googlenet" 288 | //! \param[in] cmdline The command line used to reproduce the test 289 | // 290 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 291 | //! 292 | static TestAtom defineTest(const std::string& name, const std::string& cmdline) 293 | { 294 | return TestAtom(false, name, cmdline); 295 | } 296 | 297 | //! 298 | //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments 299 | //! as input 300 | //! 301 | //! \param[in] name The name of the test 302 | //! \param[in] argc The number of command-line arguments 303 | //! \param[in] argv The array of command-line arguments (given as C strings) 304 | //! 305 | //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). 306 | static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) 307 | { 308 | auto cmdline = genCmdlineString(argc, argv); 309 | return defineTest(name, cmdline); 310 | } 311 | 312 | //! 313 | //! \brief Report that a test has started. 314 | //! 315 | //! \pre reportTestStart() has not been called yet for the given testAtom 316 | //! 317 | //! \param[in] testAtom The handle to the test that has started 318 | //! 319 | static void reportTestStart(TestAtom& testAtom) 320 | { 321 | reportTestResult(testAtom, TestResult::kRUNNING); 322 | assert(!testAtom.mStarted); 323 | testAtom.mStarted = true; 324 | } 325 | 326 | //! 327 | //! \brief Report that a test has ended. 328 | //! 329 | //! \pre reportTestStart() has been called for the given testAtom 330 | //! 331 | //! \param[in] testAtom The handle to the test that has ended 332 | //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, 333 | //! TestResult::kFAILED, TestResult::kWAIVED 334 | //! 335 | static void reportTestEnd(const TestAtom& testAtom, TestResult result) 336 | { 337 | assert(result != TestResult::kRUNNING); 338 | assert(testAtom.mStarted); 339 | reportTestResult(testAtom, result); 340 | } 341 | 342 | static int reportPass(const TestAtom& testAtom) 343 | { 344 | reportTestEnd(testAtom, TestResult::kPASSED); 345 | return EXIT_SUCCESS; 346 | } 347 | 348 | static int reportFail(const TestAtom& testAtom) 349 | { 350 | reportTestEnd(testAtom, TestResult::kFAILED); 351 | return EXIT_FAILURE; 352 | } 353 | 354 | static int reportWaive(const TestAtom& testAtom) 355 | { 356 | reportTestEnd(testAtom, TestResult::kWAIVED); 357 | return EXIT_SUCCESS; 358 | } 359 | 360 | static int reportTest(const TestAtom& testAtom, bool pass) 361 | { 362 | return pass ? reportPass(testAtom) : reportFail(testAtom); 363 | } 364 | 365 | Severity getReportableSeverity() const 366 | { 367 | return mReportableSeverity; 368 | } 369 | 370 | private: 371 | //! 372 | //! \brief returns an appropriate string for prefixing a log message with the given severity 373 | //! 374 | static const char* severityPrefix(Severity severity) 375 | { 376 | switch (severity) 377 | { 378 | case Severity::kINTERNAL_ERROR: return "[F] "; 379 | case Severity::kERROR: return "[E] "; 380 | case Severity::kWARNING: return "[W] "; 381 | case Severity::kINFO: return "[I] "; 382 | case Severity::kVERBOSE: return "[V] "; 383 | default: assert(0); return ""; 384 | } 385 | } 386 | 387 | //! 388 | //! \brief returns an appropriate string for prefixing a test result message with the given result 389 | //! 390 | static const char* testResultString(TestResult result) 391 | { 392 | switch (result) 393 | { 394 | case TestResult::kRUNNING: return "RUNNING"; 395 | case TestResult::kPASSED: return "PASSED"; 396 | case TestResult::kFAILED: return "FAILED"; 397 | case TestResult::kWAIVED: return "WAIVED"; 398 | default: assert(0); return ""; 399 | } 400 | } 401 | 402 | //! 403 | //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity 404 | //! 405 | static std::ostream& severityOstream(Severity severity) 406 | { 407 | return severity >= Severity::kINFO ? std::cout : std::cerr; 408 | } 409 | 410 | //! 411 | //! \brief method that implements logging test results 412 | //! 413 | static void reportTestResult(const TestAtom& testAtom, TestResult result) 414 | { 415 | severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " 416 | << testAtom.mCmdline << std::endl; 417 | } 418 | 419 | //! 420 | //! \brief generate a command line string from the given (argc, argv) values 421 | //! 422 | static std::string genCmdlineString(int argc, char const* const* argv) 423 | { 424 | std::stringstream ss; 425 | for (int i = 0; i < argc; i++) 426 | { 427 | if (i > 0) 428 | ss << " "; 429 | ss << argv[i]; 430 | } 431 | return ss.str(); 432 | } 433 | 434 | Severity mReportableSeverity; 435 | }; 436 | 437 | namespace 438 | { 439 | 440 | //! 441 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE 442 | //! 443 | //! Example usage: 444 | //! 445 | //! LOG_VERBOSE(logger) << "hello world" << std::endl; 446 | //! 447 | inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) 448 | { 449 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); 450 | } 451 | 452 | //! 453 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO 454 | //! 455 | //! Example usage: 456 | //! 457 | //! LOG_INFO(logger) << "hello world" << std::endl; 458 | //! 459 | inline LogStreamConsumer LOG_INFO(const Logger& logger) 460 | { 461 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); 462 | } 463 | 464 | //! 465 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING 466 | //! 467 | //! Example usage: 468 | //! 469 | //! LOG_WARN(logger) << "hello world" << std::endl; 470 | //! 471 | inline LogStreamConsumer LOG_WARN(const Logger& logger) 472 | { 473 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); 474 | } 475 | 476 | //! 477 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR 478 | //! 479 | //! Example usage: 480 | //! 481 | //! LOG_ERROR(logger) << "hello world" << std::endl; 482 | //! 483 | inline LogStreamConsumer LOG_ERROR(const Logger& logger) 484 | { 485 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); 486 | } 487 | 488 | //! 489 | //! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR 490 | // ("fatal" severity) 491 | //! 492 | //! Example usage: 493 | //! 494 | //! LOG_FATAL(logger) << "hello world" << std::endl; 495 | //! 496 | inline LogStreamConsumer LOG_FATAL(const Logger& logger) 497 | { 498 | return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); 499 | } 500 | 501 | } // anonymous namespace 502 | 503 | #endif // TENSORRT_LOGGING_H 504 | -------------------------------------------------------------------------------- /yolov5/utils/plots.py: -------------------------------------------------------------------------------- 1 | # Plotting utils 2 | 3 | import glob 4 | import math 5 | import os 6 | import random 7 | from copy import copy 8 | from pathlib import Path 9 | 10 | import cv2 11 | import matplotlib 12 | import matplotlib.pyplot as plt 13 | import numpy as np 14 | import pandas as pd 15 | import seaborn as sns 16 | import torch 17 | import yaml 18 | from PIL import Image, ImageDraw 19 | from scipy.signal import butter, filtfilt 20 | 21 | from utils.general import xywh2xyxy, xyxy2xywh 22 | from utils.metrics import fitness 23 | 24 | # Settings 25 | matplotlib.rc('font', **{'size': 11}) 26 | matplotlib.use('Agg') # for writing to files only 27 | 28 | 29 | def color_list(): 30 | # Return first 10 plt colors as (r,g,b) https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb 31 | def hex2rgb(h): 32 | return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) 33 | 34 | return [hex2rgb(h) for h in plt.rcParams['axes.prop_cycle'].by_key()['color']] 35 | 36 | 37 | def hist2d(x, y, n=100): 38 | # 2d histogram used in labels.png and evolve.png 39 | xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n) 40 | hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges)) 41 | xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1) 42 | yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1) 43 | return np.log(hist[xidx, yidx]) 44 | 45 | 46 | def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5): 47 | # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy 48 | def butter_lowpass(cutoff, fs, order): 49 | nyq = 0.5 * fs 50 | normal_cutoff = cutoff / nyq 51 | return butter(order, normal_cutoff, btype='low', analog=False) 52 | 53 | b, a = butter_lowpass(cutoff, fs, order=order) 54 | return filtfilt(b, a, data) # forward-backward filter 55 | 56 | 57 | def plot_one_box(x, img, color=None, label=None, line_thickness=None): 58 | # Plots one bounding box on image img 59 | tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness 60 | color = color or [random.randint(0, 255) for _ in range(3)] 61 | c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) 62 | cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) 63 | if label: 64 | tf = max(tl - 1, 1) # font thickness 65 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] 66 | c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 67 | cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled 68 | cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) 69 | 70 | 71 | def plot_wh_methods(): # from utils.plots import *; plot_wh_methods() 72 | # Compares the two methods for width-height anchor multiplication 73 | # https://github.com/ultralytics/yolov3/issues/168 74 | x = np.arange(-4.0, 4.0, .1) 75 | ya = np.exp(x) 76 | yb = torch.sigmoid(torch.from_numpy(x)).numpy() * 2 77 | 78 | fig = plt.figure(figsize=(6, 3), tight_layout=True) 79 | plt.plot(x, ya, '.-', label='YOLOv3') 80 | plt.plot(x, yb ** 2, '.-', label='YOLOv5 ^2') 81 | plt.plot(x, yb ** 1.6, '.-', label='YOLOv5 ^1.6') 82 | plt.xlim(left=-4, right=4) 83 | plt.ylim(bottom=0, top=6) 84 | plt.xlabel('input') 85 | plt.ylabel('output') 86 | plt.grid() 87 | plt.legend() 88 | fig.savefig('comparison.png', dpi=200) 89 | 90 | 91 | def output_to_target(output): 92 | # Convert model output to target format [batch_id, class_id, x, y, w, h, conf] 93 | targets = [] 94 | for i, o in enumerate(output): 95 | for *box, conf, cls in o.cpu().numpy(): 96 | targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf]) 97 | return np.array(targets) 98 | 99 | 100 | def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16): 101 | # Plot image grid with labels 102 | 103 | if isinstance(images, torch.Tensor): 104 | images = images.cpu().float().numpy() 105 | if isinstance(targets, torch.Tensor): 106 | targets = targets.cpu().numpy() 107 | 108 | # un-normalise 109 | if np.max(images[0]) <= 1: 110 | images *= 255 111 | 112 | tl = 3 # line thickness 113 | tf = max(tl - 1, 1) # font thickness 114 | bs, _, h, w = images.shape # batch size, _, height, width 115 | bs = min(bs, max_subplots) # limit plot images 116 | ns = np.ceil(bs ** 0.5) # number of subplots (square) 117 | 118 | # Check if we should resize 119 | scale_factor = max_size / max(h, w) 120 | if scale_factor < 1: 121 | h = math.ceil(scale_factor * h) 122 | w = math.ceil(scale_factor * w) 123 | 124 | colors = color_list() # list of colors 125 | mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init 126 | for i, img in enumerate(images): 127 | if i == max_subplots: # if last batch has fewer images than we expect 128 | break 129 | 130 | block_x = int(w * (i // ns)) 131 | block_y = int(h * (i % ns)) 132 | 133 | img = img.transpose(1, 2, 0) 134 | if scale_factor < 1: 135 | img = cv2.resize(img, (w, h)) 136 | 137 | mosaic[block_y:block_y + h, block_x:block_x + w, :] = img 138 | if len(targets) > 0: 139 | image_targets = targets[targets[:, 0] == i] 140 | boxes = xywh2xyxy(image_targets[:, 2:6]).T 141 | classes = image_targets[:, 1].astype('int') 142 | labels = image_targets.shape[1] == 6 # labels if no conf column 143 | conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred) 144 | 145 | if boxes.shape[1]: 146 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01 147 | boxes[[0, 2]] *= w # scale to pixels 148 | boxes[[1, 3]] *= h 149 | elif scale_factor < 1: # absolute coords need scale if image scales 150 | boxes *= scale_factor 151 | boxes[[0, 2]] += block_x 152 | boxes[[1, 3]] += block_y 153 | for j, box in enumerate(boxes.T): 154 | cls = int(classes[j]) 155 | color = colors[cls % len(colors)] 156 | cls = names[cls] if names else cls 157 | if labels or conf[j] > 0.25: # 0.25 conf thresh 158 | label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j]) 159 | plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl) 160 | 161 | # Draw image filename labels 162 | if paths: 163 | label = Path(paths[i]).name[:40] # trim to 40 char 164 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] 165 | cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf, 166 | lineType=cv2.LINE_AA) 167 | 168 | # Image border 169 | cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3) 170 | 171 | if fname: 172 | r = min(1280. / max(h, w) / ns, 1.0) # ratio to limit image size 173 | mosaic = cv2.resize(mosaic, (int(ns * w * r), int(ns * h * r)), interpolation=cv2.INTER_AREA) 174 | # cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB)) # cv2 save 175 | Image.fromarray(mosaic).save(fname) # PIL save 176 | return mosaic 177 | 178 | 179 | def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''): 180 | # Plot LR simulating training for full epochs 181 | optimizer, scheduler = copy(optimizer), copy(scheduler) # do not modify originals 182 | y = [] 183 | for _ in range(epochs): 184 | scheduler.step() 185 | y.append(optimizer.param_groups[0]['lr']) 186 | plt.plot(y, '.-', label='LR') 187 | plt.xlabel('epoch') 188 | plt.ylabel('LR') 189 | plt.grid() 190 | plt.xlim(0, epochs) 191 | plt.ylim(0) 192 | plt.savefig(Path(save_dir) / 'LR.png', dpi=200) 193 | 194 | 195 | def plot_test_txt(): # from utils.plots import *; plot_test() 196 | # Plot test.txt histograms 197 | x = np.loadtxt('test.txt', dtype=np.float32) 198 | box = xyxy2xywh(x[:, :4]) 199 | cx, cy = box[:, 0], box[:, 1] 200 | 201 | fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True) 202 | ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0) 203 | ax.set_aspect('equal') 204 | plt.savefig('hist2d.png', dpi=300) 205 | 206 | fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True) 207 | ax[0].hist(cx, bins=600) 208 | ax[1].hist(cy, bins=600) 209 | plt.savefig('hist1d.png', dpi=200) 210 | 211 | 212 | def plot_targets_txt(): # from utils.plots import *; plot_targets_txt() 213 | # Plot targets.txt histograms 214 | x = np.loadtxt('targets.txt', dtype=np.float32).T 215 | s = ['x targets', 'y targets', 'width targets', 'height targets'] 216 | fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True) 217 | ax = ax.ravel() 218 | for i in range(4): 219 | ax[i].hist(x[i], bins=100, label='%.3g +/- %.3g' % (x[i].mean(), x[i].std())) 220 | ax[i].legend() 221 | ax[i].set_title(s[i]) 222 | plt.savefig('targets.jpg', dpi=200) 223 | 224 | 225 | def plot_study_txt(path='', x=None): # from utils.plots import *; plot_study_txt() 226 | # Plot study.txt generated by test.py 227 | fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True) 228 | ax = ax.ravel() 229 | 230 | fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True) 231 | for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]: 232 | y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T 233 | x = np.arange(y.shape[1]) if x is None else np.array(x) 234 | s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)'] 235 | for i in range(7): 236 | ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8) 237 | ax[i].set_title(s[i]) 238 | 239 | j = y[3].argmax() + 1 240 | ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8, 241 | label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO')) 242 | 243 | ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5], 244 | 'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet') 245 | 246 | ax2.grid() 247 | ax2.set_xlim(0, 30) 248 | ax2.set_ylim(28, 50) 249 | ax2.set_yticks(np.arange(30, 55, 5)) 250 | ax2.set_xlabel('GPU Speed (ms/img)') 251 | ax2.set_ylabel('COCO AP val') 252 | ax2.legend(loc='lower right') 253 | plt.savefig('test_study.png', dpi=300) 254 | 255 | 256 | def plot_labels(labels, save_dir=Path(''), loggers=None): 257 | # plot dataset labels 258 | print('Plotting labels... ') 259 | c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes 260 | nc = int(c.max() + 1) # number of classes 261 | colors = color_list() 262 | x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height']) 263 | 264 | # seaborn correlogram 265 | sns.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9)) 266 | plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200) 267 | plt.close() 268 | 269 | # matplotlib labels 270 | matplotlib.use('svg') # faster 271 | ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel() 272 | ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8) 273 | ax[0].set_xlabel('classes') 274 | sns.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9) 275 | sns.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9) 276 | 277 | # rectangles 278 | labels[:, 1:3] = 0.5 # center 279 | labels[:, 1:] = xywh2xyxy(labels[:, 1:]) * 2000 280 | img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255) 281 | for cls, *box in labels[:1000]: 282 | ImageDraw.Draw(img).rectangle(box, width=1, outline=colors[int(cls) % 10]) # plot 283 | ax[1].imshow(img) 284 | ax[1].axis('off') 285 | 286 | for a in [0, 1, 2, 3]: 287 | for s in ['top', 'right', 'left', 'bottom']: 288 | ax[a].spines[s].set_visible(False) 289 | 290 | plt.savefig(save_dir / 'labels.jpg', dpi=200) 291 | matplotlib.use('Agg') 292 | plt.close() 293 | 294 | # loggers 295 | for k, v in loggers.items() or {}: 296 | if k == 'wandb' and v: 297 | v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]}) 298 | 299 | 300 | def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution() 301 | # Plot hyperparameter evolution results in evolve.txt 302 | with open(yaml_file) as f: 303 | hyp = yaml.load(f, Loader=yaml.FullLoader) 304 | x = np.loadtxt('evolve.txt', ndmin=2) 305 | f = fitness(x) 306 | # weights = (f - f.min()) ** 2 # for weighted results 307 | plt.figure(figsize=(10, 12), tight_layout=True) 308 | matplotlib.rc('font', **{'size': 8}) 309 | for i, (k, v) in enumerate(hyp.items()): 310 | y = x[:, i + 7] 311 | # mu = (y * weights).sum() / weights.sum() # best weighted result 312 | mu = y[f.argmax()] # best single result 313 | plt.subplot(6, 5, i + 1) 314 | plt.scatter(y, f, c=hist2d(y, f, 20), cmap='viridis', alpha=.8, edgecolors='none') 315 | plt.plot(mu, f.max(), 'k+', markersize=15) 316 | plt.title('%s = %.3g' % (k, mu), fontdict={'size': 9}) # limit to 40 characters 317 | if i % 5 != 0: 318 | plt.yticks([]) 319 | print('%15s: %.3g' % (k, mu)) 320 | plt.savefig('evolve.png', dpi=200) 321 | print('\nPlot saved as evolve.png') 322 | 323 | 324 | def profile_idetection(start=0, stop=0, labels=(), save_dir=''): 325 | # Plot iDetection '*.txt' per-image logs. from utils.plots import *; profile_idetection() 326 | ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel() 327 | s = ['Images', 'Free Storage (GB)', 'RAM Usage (GB)', 'Battery', 'dt_raw (ms)', 'dt_smooth (ms)', 'real-world FPS'] 328 | files = list(Path(save_dir).glob('frames*.txt')) 329 | for fi, f in enumerate(files): 330 | try: 331 | results = np.loadtxt(f, ndmin=2).T[:, 90:-30] # clip first and last rows 332 | n = results.shape[1] # number of rows 333 | x = np.arange(start, min(stop, n) if stop else n) 334 | results = results[:, x] 335 | t = (results[0] - results[0].min()) # set t0=0s 336 | results[0] = x 337 | for i, a in enumerate(ax): 338 | if i < len(results): 339 | label = labels[fi] if len(labels) else f.stem.replace('frames_', '') 340 | a.plot(t, results[i], marker='.', label=label, linewidth=1, markersize=5) 341 | a.set_title(s[i]) 342 | a.set_xlabel('time (s)') 343 | # if fi == len(files) - 1: 344 | # a.set_ylim(bottom=0) 345 | for side in ['top', 'right']: 346 | a.spines[side].set_visible(False) 347 | else: 348 | a.remove() 349 | except Exception as e: 350 | print('Warning: Plotting error for %s; %s' % (f, e)) 351 | 352 | ax[1].legend() 353 | plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200) 354 | 355 | 356 | def plot_results_overlay(start=0, stop=0): # from utils.plots import *; plot_results_overlay() 357 | # Plot training 'results*.txt', overlaying train and val losses 358 | s = ['train', 'train', 'train', 'Precision', 'mAP@0.5', 'val', 'val', 'val', 'Recall', 'mAP@0.5:0.95'] # legends 359 | t = ['Box', 'Objectness', 'Classification', 'P-R', 'mAP-F1'] # titles 360 | for f in sorted(glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt')): 361 | results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T 362 | n = results.shape[1] # number of rows 363 | x = range(start, min(stop, n) if stop else n) 364 | fig, ax = plt.subplots(1, 5, figsize=(14, 3.5), tight_layout=True) 365 | ax = ax.ravel() 366 | for i in range(5): 367 | for j in [i, i + 5]: 368 | y = results[j, x] 369 | ax[i].plot(x, y, marker='.', label=s[j]) 370 | # y_smooth = butter_lowpass_filtfilt(y) 371 | # ax[i].plot(x, np.gradient(y_smooth), marker='.', label=s[j]) 372 | 373 | ax[i].set_title(t[i]) 374 | ax[i].legend() 375 | ax[i].set_ylabel(f) if i == 0 else None # add filename 376 | fig.savefig(f.replace('.txt', '.png'), dpi=200) 377 | 378 | 379 | def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''): 380 | # Plot training 'results*.txt'. from utils.plots import *; plot_results(save_dir='runs/train/exp') 381 | fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True) 382 | ax = ax.ravel() 383 | s = ['Box', 'Objectness', 'Classification', 'Precision', 'Recall', 384 | 'val Box', 'val Objectness', 'val Classification', 'mAP@0.5', 'mAP@0.5:0.95'] 385 | if bucket: 386 | # files = ['https://storage.googleapis.com/%s/results%g.txt' % (bucket, x) for x in id] 387 | files = ['results%g.txt' % x for x in id] 388 | c = ('gsutil cp ' + '%s ' * len(files) + '.') % tuple('gs://%s/results%g.txt' % (bucket, x) for x in id) 389 | os.system(c) 390 | else: 391 | files = list(Path(save_dir).glob('results*.txt')) 392 | assert len(files), 'No results.txt files found in %s, nothing to plot.' % os.path.abspath(save_dir) 393 | for fi, f in enumerate(files): 394 | try: 395 | results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T 396 | n = results.shape[1] # number of rows 397 | x = range(start, min(stop, n) if stop else n) 398 | for i in range(10): 399 | y = results[i, x] 400 | if i in [0, 1, 2, 5, 6, 7]: 401 | y[y == 0] = np.nan # don't show zero loss values 402 | # y /= y[0] # normalize 403 | label = labels[fi] if len(labels) else f.stem 404 | ax[i].plot(x, y, marker='.', label=label, linewidth=2, markersize=8) 405 | ax[i].set_title(s[i]) 406 | # if i in [5, 6, 7]: # share train and val loss y axes 407 | # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) 408 | except Exception as e: 409 | print('Warning: Plotting error for %s; %s' % (f, e)) 410 | 411 | ax[1].legend() 412 | fig.savefig(Path(save_dir) / 'results.png', dpi=200) 413 | -------------------------------------------------------------------------------- /yolov5/utils/general.py: -------------------------------------------------------------------------------- 1 | # General utils 2 | 3 | import glob 4 | import logging 5 | import math 6 | import os 7 | import platform 8 | import random 9 | import re 10 | import subprocess 11 | import time 12 | from pathlib import Path 13 | 14 | import cv2 15 | import numpy as np 16 | import torch 17 | import torchvision 18 | import yaml 19 | 20 | from utils.google_utils import gsutil_getsize 21 | from utils.metrics import fitness 22 | from utils.torch_utils import init_torch_seeds 23 | 24 | # Settings 25 | torch.set_printoptions(linewidth=320, precision=5, profile='long') 26 | np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 27 | cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) 28 | 29 | 30 | def set_logging(rank=-1): 31 | logging.basicConfig( 32 | format="%(message)s", 33 | level=logging.INFO if rank in [-1, 0] else logging.WARN) 34 | 35 | 36 | def init_seeds(seed=0): 37 | random.seed(seed) 38 | np.random.seed(seed) 39 | init_torch_seeds(seed) 40 | 41 | 42 | def get_latest_run(search_dir='.'): 43 | # Return path to most recent 'last.pt' in /runs (i.e. to --resume from) 44 | last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True) 45 | return max(last_list, key=os.path.getctime) if last_list else '' 46 | 47 | 48 | def check_git_status(): 49 | # Suggest 'git pull' if repo is out of date 50 | if platform.system() in ['Linux', 'Darwin'] and not os.path.isfile('/.dockerenv'): 51 | s = subprocess.check_output('if [ -d .git ]; then git fetch && git status -uno; fi', shell=True).decode('utf-8') 52 | if 'Your branch is behind' in s: 53 | print(s[s.find('Your branch is behind'):s.find('\n\n')] + '\n') 54 | 55 | 56 | def check_img_size(img_size, s=32): 57 | # Verify img_size is a multiple of stride s 58 | new_size = make_divisible(img_size, int(s)) # ceil gs-multiple 59 | if new_size != img_size: 60 | print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size)) 61 | return new_size 62 | 63 | 64 | def check_file(file): 65 | # Search for file if not found 66 | if os.path.isfile(file) or file == '': 67 | return file 68 | else: 69 | files = glob.glob('./**/' + file, recursive=True) # find file 70 | assert len(files), 'File Not Found: %s' % file # assert file was found 71 | assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files) # assert unique 72 | return files[0] # return file 73 | 74 | 75 | def check_dataset(dict): 76 | # Download dataset if not found locally 77 | val, s = dict.get('val'), dict.get('download') 78 | if val and len(val): 79 | val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path 80 | if not all(x.exists() for x in val): 81 | print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()]) 82 | if s and len(s): # download script 83 | print('Downloading %s ...' % s) 84 | if s.startswith('http') and s.endswith('.zip'): # URL 85 | f = Path(s).name # filename 86 | torch.hub.download_url_to_file(s, f) 87 | r = os.system('unzip -q %s -d ../ && rm %s' % (f, f)) # unzip 88 | else: # bash script 89 | r = os.system(s) 90 | print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure')) # analyze return value 91 | else: 92 | raise Exception('Dataset not found.') 93 | 94 | 95 | def make_divisible(x, divisor): 96 | # Returns x evenly divisible by divisor 97 | return math.ceil(x / divisor) * divisor 98 | 99 | 100 | def clean_str(s): 101 | # Cleans a string by replacing special characters with underscore _ 102 | return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s) 103 | 104 | 105 | def labels_to_class_weights(labels, nc=80): 106 | # Get class weights (inverse frequency) from training labels 107 | if labels[0] is None: # no labels loaded 108 | return torch.Tensor() 109 | 110 | labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO 111 | classes = labels[:, 0].astype(np.int) # labels = [class xywh] 112 | weights = np.bincount(classes, minlength=nc) # occurrences per class 113 | 114 | # Prepend gridpoint count (for uCE training) 115 | # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image 116 | # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start 117 | 118 | weights[weights == 0] = 1 # replace empty bins with 1 119 | weights = 1 / weights # number of targets per class 120 | weights /= weights.sum() # normalize 121 | return torch.from_numpy(weights) 122 | 123 | 124 | def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)): 125 | # Produces image weights based on class_weights and image contents 126 | class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels]) 127 | image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1) 128 | # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample 129 | return image_weights 130 | 131 | 132 | def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper) 133 | # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/ 134 | # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') 135 | # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') 136 | # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco 137 | # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet 138 | x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 139 | 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 140 | 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] 141 | return x 142 | 143 | 144 | def xyxy2xywh(x): 145 | # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right 146 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 147 | y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center 148 | y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center 149 | y[:, 2] = x[:, 2] - x[:, 0] # width 150 | y[:, 3] = x[:, 3] - x[:, 1] # height 151 | return y 152 | 153 | 154 | def xywh2xyxy(x): 155 | # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right 156 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x) 157 | y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x 158 | y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y 159 | y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x 160 | y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y 161 | return y 162 | 163 | 164 | def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): 165 | # Rescale coords (xyxy) from img1_shape to img0_shape 166 | if ratio_pad is None: # calculate from img0_shape 167 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new 168 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding 169 | else: 170 | gain = ratio_pad[0][0] 171 | pad = ratio_pad[1] 172 | 173 | coords[:, [0, 2]] -= pad[0] # x padding 174 | coords[:, [1, 3]] -= pad[1] # y padding 175 | coords[:, :4] /= gain 176 | clip_coords(coords, img0_shape) 177 | return coords 178 | 179 | 180 | def clip_coords(boxes, img_shape): 181 | # Clip bounding xyxy bounding boxes to image shape (height, width) 182 | boxes[:, 0].clamp_(0, img_shape[1]) # x1 183 | boxes[:, 1].clamp_(0, img_shape[0]) # y1 184 | boxes[:, 2].clamp_(0, img_shape[1]) # x2 185 | boxes[:, 3].clamp_(0, img_shape[0]) # y2 186 | 187 | 188 | def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9): 189 | # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4 190 | box2 = box2.T 191 | 192 | # Get the coordinates of bounding boxes 193 | if x1y1x2y2: # x1, y1, x2, y2 = box1 194 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3] 195 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] 196 | else: # transform from xywh to xyxy 197 | b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2 198 | b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2 199 | b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2 200 | b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 201 | 202 | # Intersection area 203 | inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ 204 | (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) 205 | 206 | # Union Area 207 | w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps 208 | w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps 209 | union = w1 * h1 + w2 * h2 - inter + eps 210 | 211 | iou = inter / union 212 | if GIoU or DIoU or CIoU: 213 | cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width 214 | ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height 215 | if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 216 | c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared 217 | rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + 218 | (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared 219 | if DIoU: 220 | return iou - rho2 / c2 # DIoU 221 | elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 222 | v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2) 223 | with torch.no_grad(): 224 | alpha = v / ((1 + eps) - iou + v) 225 | return iou - (rho2 / c2 + v * alpha) # CIoU 226 | else: # GIoU https://arxiv.org/pdf/1902.09630.pdf 227 | c_area = cw * ch + eps # convex area 228 | return iou - (c_area - union) / c_area # GIoU 229 | else: 230 | return iou # IoU 231 | 232 | 233 | def box_iou(box1, box2): 234 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py 235 | """ 236 | Return intersection-over-union (Jaccard index) of boxes. 237 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format. 238 | Arguments: 239 | box1 (Tensor[N, 4]) 240 | box2 (Tensor[M, 4]) 241 | Returns: 242 | iou (Tensor[N, M]): the NxM matrix containing the pairwise 243 | IoU values for every element in boxes1 and boxes2 244 | """ 245 | 246 | def box_area(box): 247 | # box = 4xn 248 | return (box[2] - box[0]) * (box[3] - box[1]) 249 | 250 | area1 = box_area(box1.T) 251 | area2 = box_area(box2.T) 252 | 253 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2) 254 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2) 255 | return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) 256 | 257 | 258 | def wh_iou(wh1, wh2): 259 | # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2 260 | wh1 = wh1[:, None] # [N,1,2] 261 | wh2 = wh2[None] # [1,M,2] 262 | inter = torch.min(wh1, wh2).prod(2) # [N,M] 263 | return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter) 264 | 265 | 266 | def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()): 267 | """Performs Non-Maximum Suppression (NMS) on inference results 268 | 269 | Returns: 270 | detections with shape: nx6 (x1, y1, x2, y2, conf, cls) 271 | """ 272 | 273 | nc = prediction.shape[2] - 5 # number of classes 274 | xc = prediction[..., 4] > conf_thres # candidates 275 | 276 | # Settings 277 | min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height 278 | max_det = 300 # maximum number of detections per image 279 | time_limit = 10.0 # seconds to quit after 280 | redundant = True # require redundant detections 281 | multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) 282 | merge = False # use merge-NMS 283 | 284 | t = time.time() 285 | output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0] 286 | for xi, x in enumerate(prediction): # image index, image inference 287 | # Apply constraints 288 | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height 289 | x = x[xc[xi]] # confidence 290 | 291 | # Cat apriori labels if autolabelling 292 | if labels and len(labels[xi]): 293 | l = labels[xi] 294 | v = torch.zeros((len(l), nc + 5), device=x.device) 295 | v[:, :4] = l[:, 1:5] # box 296 | v[:, 4] = 1.0 # conf 297 | v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls 298 | x = torch.cat((x, v), 0) 299 | 300 | # If none remain process next image 301 | if not x.shape[0]: 302 | continue 303 | 304 | # Compute conf 305 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 306 | 307 | # Box (center x, center y, width, height) to (x1, y1, x2, y2) 308 | box = xywh2xyxy(x[:, :4]) 309 | 310 | # Detections matrix nx6 (xyxy, conf, cls) 311 | if multi_label: 312 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 313 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 314 | else: # best class only 315 | conf, j = x[:, 5:].max(1, keepdim=True) 316 | x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] 317 | 318 | # Filter by class 319 | if classes is not None: 320 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 321 | 322 | # Apply finite constraint 323 | # if not torch.isfinite(x).all(): 324 | # x = x[torch.isfinite(x).all(1)] 325 | 326 | # If none remain process next image 327 | n = x.shape[0] # number of boxes 328 | if not n: 329 | continue 330 | 331 | # Sort by confidence 332 | # x = x[x[:, 4].argsort(descending=True)] 333 | 334 | # Batched NMS 335 | c = x[:, 5:6] * (0 if agnostic else max_wh) # classes 336 | boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores 337 | i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS 338 | if i.shape[0] > max_det: # limit detections 339 | i = i[:max_det] 340 | if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) 341 | # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) 342 | iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix 343 | weights = iou * scores[None] # box weights 344 | x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes 345 | if redundant: 346 | i = i[iou.sum(1) > 1] # require redundancy 347 | 348 | output[xi] = x[i] 349 | if (time.time() - t) > time_limit: 350 | break # time limit exceeded 351 | 352 | return output 353 | 354 | 355 | def strip_optimizer(f='weights/best.pt', s=''): # from utils.general import *; strip_optimizer() 356 | # Strip optimizer from 'f' to finalize training, optionally save as 's' 357 | x = torch.load(f, map_location=torch.device('cpu')) 358 | x['optimizer'] = None 359 | x['training_results'] = None 360 | x['epoch'] = -1 361 | x['model'].half() # to FP16 362 | for p in x['model'].parameters(): 363 | p.requires_grad = False 364 | torch.save(x, s or f) 365 | mb = os.path.getsize(s or f) / 1E6 # filesize 366 | print('Optimizer stripped from %s,%s %.1fMB' % (f, (' saved as %s,' % s) if s else '', mb)) 367 | 368 | 369 | def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''): 370 | # Print mutation results to evolve.txt (for use with train.py --evolve) 371 | a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys 372 | b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values 373 | c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3) 374 | print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c)) 375 | 376 | if bucket: 377 | url = 'gs://%s/evolve.txt' % bucket 378 | if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0): 379 | os.system('gsutil cp %s .' % url) # download evolve.txt if larger than local 380 | 381 | with open('evolve.txt', 'a') as f: # append result 382 | f.write(c + b + '\n') 383 | x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows 384 | x = x[np.argsort(-fitness(x))] # sort 385 | np.savetxt('evolve.txt', x, '%10.3g') # save sort by fitness 386 | 387 | # Save yaml 388 | for i, k in enumerate(hyp.keys()): 389 | hyp[k] = float(x[0, i + 7]) 390 | with open(yaml_file, 'w') as f: 391 | results = tuple(x[0, :7]) 392 | c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3) 393 | f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n') 394 | yaml.dump(hyp, f, sort_keys=False) 395 | 396 | if bucket: 397 | os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket)) # upload 398 | 399 | 400 | def apply_classifier(x, model, img, im0): 401 | # applies a second stage classifier to yolo outputs 402 | im0 = [im0] if isinstance(im0, np.ndarray) else im0 403 | for i, d in enumerate(x): # per image 404 | if d is not None and len(d): 405 | d = d.clone() 406 | 407 | # Reshape and pad cutouts 408 | b = xyxy2xywh(d[:, :4]) # boxes 409 | b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square 410 | b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad 411 | d[:, :4] = xywh2xyxy(b).long() 412 | 413 | # Rescale boxes from img_size to im0 size 414 | scale_coords(img.shape[2:], d[:, :4], im0[i].shape) 415 | 416 | # Classes 417 | pred_cls1 = d[:, 5].long() 418 | ims = [] 419 | for j, a in enumerate(d): # per item 420 | cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])] 421 | im = cv2.resize(cutout, (224, 224)) # BGR 422 | # cv2.imwrite('test%i.jpg' % j, cutout) 423 | 424 | im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 425 | im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32 426 | im /= 255.0 # 0 - 255 to 0.0 - 1.0 427 | ims.append(im) 428 | 429 | pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction 430 | x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections 431 | 432 | return x 433 | 434 | 435 | def increment_path(path, exist_ok=True, sep=''): 436 | # Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc. 437 | path = Path(path) # os-agnostic 438 | if (path.exists() and exist_ok) or (not path.exists()): 439 | return str(path) 440 | else: 441 | dirs = glob.glob(f"{path}{sep}*") # similar paths 442 | matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs] 443 | i = [int(m.groups()[0]) for m in matches if m] # indices 444 | n = max(i) + 1 if i else 2 # increment number 445 | return f"{path}{sep}{n}" # update path 446 | --------------------------------------------------------------------------------