├── README.md ├── SDK ├── Dockerfile ├── convert_openvino.sh └── ji.py └── yolov5 ├── .dockerignore ├── .gitattributes ├── .github ├── ISSUE_TEMPLATE │ ├── --bug-report.md │ ├── --feature-request.md │ └── -question.md └── workflows │ ├── ci-testing.yml │ ├── greetings.yml │ ├── rebase.yml │ └── stale.yml ├── .ipynb_checkpoints ├── detect-checkpoint.py ├── requirements-checkpoint.txt └── train-checkpoint.py ├── LICENSE ├── README.md ├── __pycache__ └── test.cpython-36.pyc ├── data ├── .ipynb_checkpoints │ ├── coco128-checkpoint.yaml │ ├── hyp.finetune-checkpoint.yaml │ └── hyp.scratch-checkpoint.yaml ├── coco.yaml ├── coco128.yaml ├── hyp.finetune.yaml ├── hyp.scratch.yaml ├── scripts │ ├── get_coco.sh │ └── get_voc.sh └── voc.yaml ├── detect.py ├── hubconf.py ├── inference ├── images │ ├── bus.jpg │ └── zidane.jpg └── output │ ├── .ipynb_checkpoints │ └── 2209a117-checkpoint.jpg │ └── 2209a117.jpg ├── models ├── .ipynb_checkpoints │ ├── common-checkpoint.py │ ├── experimental-checkpoint.py │ ├── export-checkpoint.py │ ├── yolo-checkpoint.py │ ├── yolo_1-checkpoint.py │ └── yolov5s-checkpoint.yaml ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── common.cpython-36.pyc │ ├── experimental.cpython-36.pyc │ └── yolo.cpython-36.pyc ├── common.py ├── experimental.py ├── export.py ├── hub │ ├── yolov3-spp.yaml │ ├── yolov5-fpn.yaml │ └── yolov5-panet.yaml ├── yolo.py ├── yolo_1.py ├── yolov5l.yaml ├── yolov5m.yaml ├── yolov5s.yaml └── yolov5x.yaml ├── requirements.txt ├── test.py ├── train.py ├── tutorial.ipynb ├── utils ├── .ipynb_checkpoints │ └── general-checkpoint.py ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── datasets.cpython-36.pyc │ ├── general.cpython-36.pyc │ ├── google_utils.cpython-36.pyc │ └── torch_utils.cpython-36.pyc ├── activations.py ├── datasets.py ├── general.py ├── google_utils.py └── torch_utils.py └── weights └── download_weights.sh /README.md: -------------------------------------------------------------------------------- 1 | # yolov5_openvino_sdk 2 | an SDK about how to use openvino model transformed from yolov5 3 | 4 | ## Before Start 5 | train your yolov5 model on your own dataset following [Train Custom Data](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data) 6 | 7 | ## export model to onnx，using export.py. 8 | ### before run export.py: 9 | - **changeto torch==1.5.1,torchvision==0.6.1** 10 | - **delete [yolo.py](https://github.com/linhaoqi027/yolov5_openvino_sdk/blob/master/yolov5/models/yolo.py), and rename [yolo_1.py](https://github.com/linhaoqi027/yolov5_openvino_sdk/blob/master/yolov5/models/yolo_1.py) to yolo.py** 11 | 12 | then you can run export.py to export onnx model. 13 | 14 | ## install OPENVINO2020R4 15 | you can install following https://bbs.cvmart.net/topics/3117 16 | 17 | 18 | ## Convert onnx to openvino 19 | AS for how to use openvino to inference ,please refer to [SDK](https://github.com/linhaoqi027/yolov5_openvino_sdk/tree/master/SDK) 20 | using [convert_openvino.sh](https://github.com/linhaoqi027/yolov5_openvino_sdk/blob/master/SDK/convert_openvino.sh) to convert onnx to openvino.Then you can get openvino model .bin and .xml. 21 | 22 | ## use openvino model to inference picture. 23 | This profile provide an SDK. 24 | using [ji.py](https://github.com/linhaoqi027/yolov5_openvino_sdk/blob/master/SDK/ji.py) to using transformed openvino model. 25 | (WARNING:NMS is not included in openvino model.) 26 | -------------------------------------------------------------------------------- /SDK/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM uhub.service.ucloud.cn/eagle_nest/cuda10.1-cudnn7.5.1-dev-ubuntu16.04-opencv4.1.1-torch1.4.0-openvino2020r1 2 | 3 | RUN pip install lxml 4 | RUN pip install Cython 5 | RUN pip install matplotlib>=3.2.2 6 | RUN pip install numpy>=1.18.5 7 | RUN pip install opencv-python 8 | RUN pip install pillow 9 | RUN pip install PyYAML==5.3.1 10 | RUN pip install scipy==1.5.2 11 | RUN pip install tensorboard>=2.2 12 | RUN pip install torch==1.5.1 13 | RUN pip install torchvision==0.6.1 14 | RUN pip install tqdm>=4.41.0 15 | RUN pip install coremltools==4.0b2 16 | RUN pip install onnx==1.6.0 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | RUN rm -rf /usr/local/ev_sdk && mkdir -p /usr/local/ev_sdk 25 | COPY ./ /usr/local/ev_sdk 26 | 27 | RUN \ 28 | cd /usr/local/ev_sdk && mkdir -p build && rm -rf build/* \ 29 | && cd build && cmake -DCMAKE_BUILD_TYPE=Release .. && make -j4 install && rm -rf ../build/* 30 | 31 | # 如果使用Python接口编写代码需要添加一个环境变量 32 | ENV AUTO_TEST_USE_JI_PYTHON_API=1 33 | 34 | 35 | 36 | 37 | # Remove old version 38 | RUN rm -rf /opt/intel/ 39 | # Download and install 2020r4 40 | RUN apt-get update \ 41 | && apt-get install -y --no-install-recommends cpio sudo lsb-release \ 42 | && cd /usr/local/src \ 43 | && wget http://10.9.0.146:8888/group1/M00/01/0B/CgkAkl8yRomEOlssAAAAAHMwDr4446.tgz -O l_openvino_toolkit_p_2020.4.287.tgz \ 44 | && tar xf l_openvino_toolkit_p_2020.4.287.tgz \ 45 | && cd l_openvino_toolkit_p_2020.4.287 \ 46 | && sed -i 's/decline/accept/g' silent.cfg \ 47 | && ./install.sh -s silent.cfg \ 48 | && clean-layer.sh 49 | 50 | # Install prerequisite 51 | RUN bash -c "source /opt/intel/openvino/bin/setupvars.sh" && cd /opt/intel/openvino/install_dependencies/ && sed -i 's/sudo -E //' ./install_openvino_dependencies.sh && ./install_openvino_dependencies.sh \ 52 | && cd /opt/intel/openvino/deployment_tools/model_optimizer/install_prerequisites && sed -i 's/sudo -E //' ./install_prerequisites.sh && sed -i 's/python3-pip python3-venv//' ./install_prerequisites.sh && sed -i 's|pip install|pip install -i https://pypi.tuna.tsinghua.edu.cn/simple/|' ./install_prerequisites.sh && ./install_prerequisites_onnx.sh \ 53 | && sed -i 's/source "$INSTALLDIR\/opencv\/setupvars.sh"/echo "$INSTALLDIR\/opencv\/setupvars.sh"/' /opt/intel/openvino/bin/setupvars.sh \ 54 | && echo "source /opt/intel/openvino/bin/setupvars.sh" >> /etc/profile.d/openvino.sh \ 55 | && clean-layer.sh 56 | 57 | # Setup env 58 | ENV HDDL_INSTALL_DIR=/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/external/hddl 59 | ENV LD_LIBRARY_PATH=/opt/intel/openvino_2020.4.287/deployment_tools/ngraph/lib:/opt/intel/opencl:/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/external/hddl/lib:/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/external/gna/lib:/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/external/mkltiny_lnx/lib:/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/external/tbb/lib:/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/lib/intel64:${LD_LIBRARY_PATH} 60 | ENV PATH=/opt/intel/openvino_2020.4.287/deployment_tools/model_optimizer:${PATH} 61 | ENV InferenceEngine_DIR=/opt/intel/openvino_2020.4.287/deployment_tools/inference_engine/share 62 | ENV PYTHONPATH=/opt/intel/openvino_2020.4.287/python/python3.6:/opt/intel/openvino_2020.4.287/python/python3:/opt/intel/openvino_2020.4.287/deployment_tools/open_model_zoo/tools/accuracy_checker:/opt/intel/openvino_2020.4.287/deployment_tools/model_optimizer:${PYTHONPATH} 63 | ENV INTEL_OPENVINO_DIR=/opt/intel/openvino_2020.4.287 64 | ENV INTEL_CVSDK_DIR=/opt/intel/openvino_2020.4.287 65 | ENV PYTHONPATH=/usr/local/lib/python3.6/dist-packages/:${PYTHONPATH} -------------------------------------------------------------------------------- /SDK/convert_openvino.sh: -------------------------------------------------------------------------------- 1 | # #!/bin/bash 2 | 3 | # # Root directory of the script 4 | # SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" 5 | 6 | # /opt/intel/openvino/deployment_tools/model_optimizer/mo_tf.py \ 7 | # --input_model ${SCRIPT_DIR}/ssd_inception_v2.pb \ 8 | # --transformations_config /opt/intel/openvino/deployment_tools/model_optimizer/extensions/front/tf/ssd_v2_support.json \ 9 | # --tensorflow_object_detection_api_pipeline_config ${SCRIPT_DIR}/ssd_inception_v2_coco.config \ 10 | # --output_dir ${SCRIPT_DIR}/openvino \ 11 | # --model_name ssd_inception_v2 \ 12 | # --input image_tensor 13 | 14 | 15 | # Root directory of the script 16 | 17 | 18 | 19 | 20 | #source /opt/intel/openvino/deployment_tools/model_optimizer/venv/bin/activate 21 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" 22 | 23 | cd /usr/local/ev_sdk/yolov5 24 | python models/export.py --weights ${SCRIPT_DIR}/yolov5x_10_best.pt --img 480 --batch 1 25 | 26 | python3 /opt/intel/openvino_2020.4.287/deployment_tools/model_optimizer/mo.py \ 27 | --input_model ${SCRIPT_DIR}/yolov5x_10_best.onnx \ 28 | --output_dir ${SCRIPT_DIR}/openvino \ 29 | --input_shape [1,3,480,480] 30 | 31 | -------------------------------------------------------------------------------- /SDK/ji.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import logging as log 4 | import os 5 | import pathlib 6 | import json 7 | import cv2 8 | import numpy as np 9 | from openvino.inference_engine import IENetwork, IECore 10 | import torch 11 | import torchvision 12 | import time 13 | 14 | 15 | 16 | def xywh2xyxy(x): 17 | # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right 18 | y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x) 19 | y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x 20 | y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y 21 | y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x 22 | y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y 23 | return y 24 | 25 | 26 | def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, merge=False, classes=None, agnostic=False): 27 | """Performs Non-Maximum Suppression (NMS) on inference results 28 | 29 | Returns: 30 | detections with shape: nx6 (x1, y1, x2, y2, conf, cls) 31 | """ 32 | prediction=torch.from_numpy(prediction) 33 | if prediction.dtype is torch.float16: 34 | prediction = prediction.float() # to FP32 35 | 36 | nc = prediction[0].shape[1] - 5 # number of classes 37 | xc = prediction[..., 4] > conf_thres # candidates 38 | 39 | # Settings 40 | min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height 41 | max_det = 300 # maximum number of detections per image 42 | time_limit = 10.0 # seconds to quit after 43 | redundant = True # require redundant detections 44 | multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img) 45 | 46 | t = time.time() 47 | output = [None] * prediction.shape[0] 48 | for xi, x in enumerate(prediction): # image index, image inference 49 | # Apply constraints 50 | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height 51 | x = x[xc[xi]] # confidence 52 | 53 | # If none remain process next image 54 | if not x.shape[0]: 55 | continue 56 | 57 | # Compute conf 58 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf 59 | 60 | # Box (center x, center y, width, height) to (x1, y1, x2, y2) 61 | box = xywh2xyxy(x[:, :4]) 62 | 63 | # Detections matrix nx6 (xyxy, conf, cls) 64 | if multi_label: 65 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T 66 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1) 67 | else: # best class only 68 | conf, j = x[:, 5:].max(1, keepdim=True) 69 | x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres] 70 | 71 | # Filter by class 72 | if classes: 73 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)] 74 | 75 | # Apply finite constraint 76 | # if not torch.isfinite(x).all(): 77 | # x = x[torch.isfinite(x).all(1)] 78 | 79 | # If none remain process next image 80 | n = x.shape[0] # number of boxes 81 | if not n: 82 | continue 83 | 84 | # Sort by confidence 85 | # x = x[x[:, 4].argsort(descending=True)] 86 | 87 | # Batched NMS 88 | c = x[:, 5:6] * (0 if agnostic else max_wh) # classes 89 | boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores 90 | i = torchvision.ops.boxes.nms(boxes, scores, iou_thres) 91 | if i.shape[0] > max_det: # limit detections 92 | i = i[:max_det] 93 | if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) 94 | try: # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) 95 | iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix 96 | weights = iou * scores[None] # box weights 97 | x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes 98 | if redundant: 99 | i = i[iou.sum(1) > 1] # require redundancy 100 | except: # possible CUDA error https://github.com/ultralytics/yolov3/issues/1139 101 | print(x, i, x.shape, i.shape) 102 | pass 103 | 104 | output[xi] = x[i] 105 | if (time.time() - t) > time_limit: 106 | break # time limit exceeded 107 | 108 | return output 109 | 110 | 111 | 112 | device = 'CPU' 113 | input_h, input_w, input_c, input_n = (480, 480, 3, 1) 114 | log.basicConfig(level=log.DEBUG) 115 | 116 | # For objection detection task, replace your target labels here. 117 | label_id_map = { 118 | 0: "fire", 119 | } 120 | exec_net = None 121 | 122 | 123 | def init(): 124 | """Initialize model 125 | 126 | Returns: model 127 | 128 | """ 129 | #model_xml = "/project/train/src_repo/yolov5/runs/exp0/weights/best.xml" 130 | model_xml = "/usr/local/ev_sdk/model/openvino/yolov5x_10_best.xml" 131 | if not os.path.isfile(model_xml): 132 | log.error(f'{model_xml} does not exist') 133 | return None 134 | model_bin = pathlib.Path(model_xml).with_suffix('.bin').as_posix() 135 | # log.info("Loading network files:\n\t{}\n\t{}".format(model_xml, model_bin)) 136 | net = IENetwork(model=model_xml, weights=model_bin) 137 | 138 | # Load Inference Engine 139 | # log.info('Loading Inference Engine') 140 | ie = IECore() 141 | global exec_net 142 | exec_net = ie.load_network(network=net, device_name=device) 143 | # log.info('Device info:') 144 | # versions = ie.get_versions(device) 145 | # print("{}".format(device)) 146 | # print("MKLDNNPlugin version ......... {}.{}".format(versions[device].major, versions[device].minor)) 147 | # print("Build ........... {}".format(versions[device].build_number)) 148 | 149 | input_blob = next(iter(net.inputs)) 150 | n, c, h, w = net.inputs[input_blob].shape 151 | global input_h, input_w, input_c, input_n 152 | input_h, input_w, input_c, input_n = h, w, c, n 153 | 154 | return net 155 | 156 | 157 | def process_image(net, input_image): 158 | """Do inference to analysis input_image and get output 159 | 160 | Attributes: 161 | net: model handle 162 | input_image (numpy.ndarray): image to be process, format: (h, w, c), BGR 163 | thresh: thresh value 164 | 165 | Returns: process result 166 | 167 | """ 168 | if not net or input_image is None: 169 | log.error('Invalid input args') 170 | return None 171 | # log.info(f'process_image, ({input_image.shape}') 172 | ih, iw, _ = input_image.shape 173 | 174 | # --------------------------- Prepare input blobs ----------------------------------------------------- 175 | if ih != input_h or iw != input_w: 176 | input_image = cv2.resize(input_image, (input_w, input_h)) 177 | input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB) 178 | input_image = input_image/255 179 | input_image = input_image.transpose((2, 0, 1)) 180 | images = np.ndarray(shape=(input_n, input_c, input_h, input_w)) 181 | images[0] = input_image 182 | 183 | input_blob = next(iter(net.inputs)) 184 | out_blob = next(iter(net.outputs)) 185 | 186 | # --------------------------- Prepare output blobs ---------------------------------------------------- 187 | # log.info('Preparing output blobs') 188 | # log.info(f"The output_name{net.outputs}") 189 | #print(net.outputs) 190 | # output_name = "Transpose_305" 191 | # try: 192 | # output_info = net.outputs[output_name] 193 | # except KeyError: 194 | # log.error(f"Can't find a {output_name} layer in the topology") 195 | # return None 196 | 197 | # output_dims = output_info.shape 198 | # log.info(f"The output_dims{output_dims}") 199 | # if len(output_dims) != 4: 200 | # log.error("Incorrect output dimensions for yolo model") 201 | # max_proposal_count, object_size = output_dims[2], output_dims[3] 202 | 203 | # if object_size != 7: 204 | # log.error("Output item should have 7 as a last dimension") 205 | 206 | #output_info.precision = "FP32" 207 | 208 | # --------------------------- Performing inference ---------------------------------------------------- 209 | # log.info("Creating infer request and starting inference") 210 | res = exec_net.infer(inputs={input_blob: images}) 211 | 212 | # --------------------------- Read and postprocess output --------------------------------------------- 213 | # log.info("Processing output blobs") 214 | 215 | # res = res[out_blob] 216 | data = res[out_blob] 217 | 218 | 219 | data=non_max_suppression(data, 0.4, 0.5) 220 | detect_objs = [] 221 | if data[0]==None: 222 | return json.dumps({"objects": detect_objs}) 223 | else: 224 | data=data[0].numpy() 225 | for proposal in data: 226 | if proposal[4] > 0 : 227 | confidence = proposal[4] 228 | xmin = np.int(iw * (proposal[0]/480)) 229 | ymin = np.int(ih * (proposal[1]/480)) 230 | xmax = np.int(iw * (proposal[2]/480)) 231 | ymax = np.int(ih * (proposal[3]/480)) 232 | # if label not in label_id_map: 233 | # log.warning(f'{label} does not in {label_id_map}') 234 | # continue 235 | detect_objs.append({ 236 | 'name': label_id_map[0], 237 | 'xmin': int(xmin), 238 | 'ymin': int(ymin), 239 | 'xmax': int(xmax), 240 | 'ymax': int(ymax), 241 | 'confidence': float(confidence) 242 | }) 243 | return json.dumps({"objects": detect_objs}) 244 | 245 | 246 | if __name__ == '__main__': 247 | # Test API 248 | img = cv2.imread('/home/data/19/2209a117.jpg') 249 | predictor = init() 250 | result = process_image(predictor, img) 251 | log.info(result) 252 | 253 | # device = 'CPU' 254 | # input_h, input_w, input_c, input_n = (640, 640, 3, 1) 255 | # log.basicConfig(level=log.DEBUG) 256 | 257 | # # For objection detection task, replace your target labels here. 258 | # label_id_map = { 259 | # 1: "fire", 260 | # } 261 | # #ie = None 262 | # exec_net = None 263 | 264 | # def init(): 265 | # """Initialize model 266 | 267 | # Returns: model 268 | 269 | # """ 270 | # model_xml = "/project/train/src_repo/yolov5/runs/exp0/weights/best.xml" 271 | # if not os.path.isfile(model_xml): 272 | # log.error(f'{model_xml} does not exist') 273 | # return None 274 | # model_bin = pathlib.Path(model_xml).with_suffix('.bin').as_posix() 275 | # log.info("Loading network files:\n\t{}\n\t{}".format(model_xml, model_bin)) 276 | # net = IENetwork(model=model_xml, weights=model_bin) 277 | 278 | # # Load Inference Engine 279 | # log.info('Loading Inference Engine') 280 | # global ie 281 | # ie = IECore() 282 | # global exec_net 283 | # exec_net = ie.load_network(network=net, device_name=device) 284 | # log.info('Device info:') 285 | # versions = ie.get_versions(device) 286 | # print("{}".format(device)) 287 | # print("MKLDNNPlugin version ......... {}.{}".format(versions[device].major, versions[device].minor)) 288 | # print("Build ........... {}".format(versions[device].build_number)) 289 | 290 | # input_blob = next(iter(net.inputs)) 291 | # n, c, h, w = net.inputs[input_blob].shape 292 | # global input_h, input_w, input_c, input_n 293 | # input_h, input_w, input_c, input_n = h, w, c, n 294 | 295 | # return net 296 | 297 | 298 | # def process_image(net, input_image, thresh): 299 | # """Do inference to analysis input_image and get output 300 | 301 | # Attributes: 302 | # net: model handle 303 | # input_image (numpy.ndarray): image to be process, format: (h, w, c), BGR 304 | # thresh: thresh value 305 | 306 | # Returns: process result 307 | 308 | # """ 309 | # if not net or input_image is None: 310 | # log.error('Invalid input args') 311 | # return None 312 | # log.info(f'process_image, ({input_image.shape}') 313 | # ih, iw, _ = input_image.shape 314 | 315 | # # --------------------------- Prepare input blobs ----------------------------------------------------- 316 | # if ih != input_h or iw != input_w: 317 | # input_image = cv2.resize(input_image, (input_w, input_h)) 318 | # input_image = input_image.transpose((2, 0, 1)) 319 | # input_image = input_image/255 320 | # images = np.ndarray(shape=(input_n, input_c, input_h, input_w)) 321 | # images[0] = input_image 322 | 323 | # input_blob = next(iter(net.inputs)) 324 | # out_blob = next(iter(net.outputs)) 325 | 326 | # # --------------------------- Prepare output blobs ---------------------------------------------------- 327 | # log.info('Preparing output blobs') 328 | # #print(net.outputs) 329 | 330 | # output_name = "Transpose_305" 331 | # try: 332 | # output_info = net.outputs[output_name] 333 | # except KeyError: 334 | # log.error(f"Can't find a {output_name} layer in the topology") 335 | # return None 336 | 337 | # output_dims = output_info.shape 338 | # if len(output_dims) != 4: 339 | # log.error("Incorrect output dimensions for SSD model") 340 | # max_proposal_count, object_size = output_dims[2], output_dims[3] 341 | 342 | # if object_size != 7: 343 | # log.error("Output item should have 7 as a last dimension") 344 | 345 | # output_info.precision = "FP32" 346 | 347 | # # --------------------------- Performing inference ---------------------------------------------------- 348 | # log.info("Creating infer request and starting inference") 349 | # res = exec_net.infer(inputs={input_blob: images}) 350 | 351 | # # --------------------------- Read and postprocess output --------------------------------------------- 352 | # log.info("Processing output blobs") 353 | # res = res[out_blob] 354 | # data = res 355 | # data=data.reshape(data.shape[0],-1,data.shape[4]) 356 | # data = non_max_suppression(data) 357 | 358 | # detect_objs = [] 359 | # for proposal in enumerate(data): 360 | # if proposal[2] > 0: 361 | # #batch_id = np.int(proposal[0]) 362 | # label = np.int(proposal[5]) 363 | # confidence = proposal[4] 364 | # xmin = np.int(iw * proposal[0]) 365 | # ymin = np.int(ih * proposal[1]) 366 | # xmax = np.int(iw * proposal[2]) 367 | # ymax = np.int(ih * proposal[3]) 368 | # if label not in label_id_map: 369 | # log.warning(f'{label} does not in {label_id_map}') 370 | # continue 371 | # if proposal[4] > thresh: 372 | # detect_objs.append({ 373 | # 'name': label_id_map[label], 374 | # 'xmin': int(xmin), 375 | # 'ymin': int(ymin), 376 | # 'xmax': int(xmax), 377 | # 'ymax': int(ymax) 378 | # }) 379 | # return json.dumps(detect_objs) 380 | 381 | 382 | # if __name__ == '__main__': 383 | # # Test API 384 | # img = cv2.imread('/project/train/images/train2017/1014a144.jpg') 385 | # predictor = init() 386 | # result = process_image(predictor, img, 0.5) 387 | # log.info(result) 388 | 389 | 390 | -------------------------------------------------------------------------------- /yolov5/.dockerignore: -------------------------------------------------------------------------------- 1 | # Repo-specific DockerIgnore ------------------------------------------------------------------------------------------- 2 | # .git 3 | .cache 4 | .idea 5 | runs 6 | output 7 | coco 8 | storage.googleapis.com 9 | 10 | data/samples/* 11 | **/results*.txt 12 | *.jpg 13 | 14 | # Neural Network weights ----------------------------------------------------------------------------------------------- 15 | **/*.weights 16 | **/*.pt 17 | **/*.pth 18 | **/*.onnx 19 | **/*.mlmodel 20 | **/*.torchscript 21 | 22 | 23 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 24 | # Below Copied From .gitignore ----------------------------------------------------------------------------------------- 25 | 26 | 27 | # GitHub Python GitIgnore ---------------------------------------------------------------------------------------------- 28 | # Byte-compiled / optimized / DLL files 29 | __pycache__/ 30 | *.py[cod] 31 | *$py.class 32 | 33 | # C extensions 34 | *.so 35 | 36 | # Distribution / packaging 37 | .Python 38 | env/ 39 | build/ 40 | develop-eggs/ 41 | dist/ 42 | downloads/ 43 | eggs/ 44 | .eggs/ 45 | lib/ 46 | lib64/ 47 | parts/ 48 | sdist/ 49 | var/ 50 | wheels/ 51 | *.egg-info/ 52 | .installed.cfg 53 | *.egg 54 | 55 | # PyInstaller 56 | # Usually these files are written by a python script from a template 57 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 58 | *.manifest 59 | *.spec 60 | 61 | # Installer logs 62 | pip-log.txt 63 | pip-delete-this-directory.txt 64 | 65 | # Unit test / coverage reports 66 | htmlcov/ 67 | .tox/ 68 | .coverage 69 | .coverage.* 70 | .cache 71 | nosetests.xml 72 | coverage.xml 73 | *.cover 74 | .hypothesis/ 75 | 76 | # Translations 77 | *.mo 78 | *.pot 79 | 80 | # Django stuff: 81 | *.log 82 | local_settings.py 83 | 84 | # Flask stuff: 85 | instance/ 86 | .webassets-cache 87 | 88 | # Scrapy stuff: 89 | .scrapy 90 | 91 | # Sphinx documentation 92 | docs/_build/ 93 | 94 | # PyBuilder 95 | target/ 96 | 97 | # Jupyter Notebook 98 | .ipynb_checkpoints 99 | 100 | # pyenv 101 | .python-version 102 | 103 | # celery beat schedule file 104 | celerybeat-schedule 105 | 106 | # SageMath parsed files 107 | *.sage.py 108 | 109 | # dotenv 110 | .env 111 | 112 | # virtualenv 113 | .venv 114 | venv*/ 115 | ENV/ 116 | 117 | # Spyder project settings 118 | .spyderproject 119 | .spyproject 120 | 121 | # Rope project settings 122 | .ropeproject 123 | 124 | # mkdocs documentation 125 | /site 126 | 127 | # mypy 128 | .mypy_cache/ 129 | 130 | 131 | # https://github.com/github/gitignore/blob/master/Global/macOS.gitignore ----------------------------------------------- 132 | 133 | # General 134 | .DS_Store 135 | .AppleDouble 136 | .LSOverride 137 | 138 | # Icon must end with two \r 139 | Icon 140 | Icon? 141 | 142 | # Thumbnails 143 | ._* 144 | 145 | # Files that might appear in the root of a volume 146 | .DocumentRevisions-V100 147 | .fseventsd 148 | .Spotlight-V100 149 | .TemporaryItems 150 | .Trashes 151 | .VolumeIcon.icns 152 | .com.apple.timemachine.donotpresent 153 | 154 | # Directories potentially created on remote AFP share 155 | .AppleDB 156 | .AppleDesktop 157 | Network Trash Folder 158 | Temporary Items 159 | .apdisk 160 | 161 | 162 | # https://github.com/github/gitignore/blob/master/Global/JetBrains.gitignore 163 | # Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm 164 | # Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 165 | 166 | # User-specific stuff: 167 | .idea/* 168 | .idea/**/workspace.xml 169 | .idea/**/tasks.xml 170 | .idea/dictionaries 171 | .html # Bokeh Plots 172 | .pg # TensorFlow Frozen Graphs 173 | .avi # videos 174 | 175 | # Sensitive or high-churn files: 176 | .idea/**/dataSources/ 177 | .idea/**/dataSources.ids 178 | .idea/**/dataSources.local.xml 179 | .idea/**/sqlDataSources.xml 180 | .idea/**/dynamic.xml 181 | .idea/**/uiDesigner.xml 182 | 183 | # Gradle: 184 | .idea/**/gradle.xml 185 | .idea/**/libraries 186 | 187 | # CMake 188 | cmake-build-debug/ 189 | cmake-build-release/ 190 | 191 | # Mongo Explorer plugin: 192 | .idea/**/mongoSettings.xml 193 | 194 | ## File-based project format: 195 | *.iws 196 | 197 | ## Plugin-specific files: 198 | 199 | # IntelliJ 200 | out/ 201 | 202 | # mpeltonen/sbt-idea plugin 203 | .idea_modules/ 204 | 205 | # JIRA plugin 206 | atlassian-ide-plugin.xml 207 | 208 | # Cursive Clojure plugin 209 | .idea/replstate.xml 210 | 211 | # Crashlytics plugin (for Android Studio and IntelliJ) 212 | com_crashlytics_export_strings.xml 213 | crashlytics.properties 214 | crashlytics-build.properties 215 | fabric.properties 216 | -------------------------------------------------------------------------------- /yolov5/.gitattributes: -------------------------------------------------------------------------------- 1 | # this drop notebooks from GitHub language stats 2 | *.ipynb linguist-vendored 3 | -------------------------------------------------------------------------------- /yolov5/.github/ISSUE_TEMPLATE/--bug-report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "\U0001F41BBug report" 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: bug 6 | assignees: '' 7 | 8 | --- 9 | 10 | Before submitting a bug report, please be aware that your issue **must be reproducible** with all of the following, otherwise it is non-actionable, and we can not help you: 11 | - **Current repo**: run `git fetch && git status -uno` to check and `git pull` to update repo 12 | - **Common dataset**: coco.yaml or coco128.yaml 13 | - **Common environment**: Colab, Google Cloud, or Docker image. See https://github.com/ultralytics/yolov5#environments 14 | 15 | If this is a custom dataset/training question you **must include** your `train*.jpg`, `test*.jpg` and `results.png` figures, or we can not help you. You can generate these with `utils.plot_results()`. 16 | 17 | 18 | ## 🐛 Bug 19 | A clear and concise description of what the bug is. 20 | 21 | 22 | ## To Reproduce (REQUIRED) 23 | 24 | Input: 25 | ``` 26 | import torch 27 | 28 | a = torch.tensor([5]) 29 | c = a / 0 30 | ``` 31 | 32 | Output: 33 | ``` 34 | Traceback (most recent call last): 35 | File "/Users/glennjocher/opt/anaconda3/envs/env1/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3331, in run_code 36 | exec(code_obj, self.user_global_ns, self.user_ns) 37 | File "", line 5, in 38 | c = a / 0 39 | RuntimeError: ZeroDivisionError 40 | ``` 41 | 42 | 43 | ## Expected behavior 44 | A clear and concise description of what you expected to happen. 45 | 46 | 47 | ## Environment 48 | If applicable, add screenshots to help explain your problem. 49 | 50 | - OS: [e.g. Ubuntu] 51 | - GPU [e.g. 2080 Ti] 52 | 53 | 54 | ## Additional context 55 | Add any other context about the problem here. 56 | -------------------------------------------------------------------------------- /yolov5/.github/ISSUE_TEMPLATE/--feature-request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "\U0001F680Feature request" 3 | about: Suggest an idea for this project 4 | title: '' 5 | labels: enhancement 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## 🚀 Feature 11 | 12 | 13 | ## Motivation 14 | 15 | 16 | 17 | ## Pitch 18 | 19 | 20 | 21 | ## Alternatives 22 | 23 | 24 | 25 | ## Additional context 26 | 27 | 28 | -------------------------------------------------------------------------------- /yolov5/.github/ISSUE_TEMPLATE/-question.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: "❓Question" 3 | about: Ask a general question 4 | title: '' 5 | labels: question 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## ❔Question 11 | 12 | 13 | ## Additional context 14 | -------------------------------------------------------------------------------- /yolov5/.github/workflows/ci-testing.yml: -------------------------------------------------------------------------------- 1 | name: CI CPU testing 2 | 3 | on: # https://help.github.com/en/actions/reference/events-that-trigger-workflows 4 | push: 5 | pull_request: 6 | schedule: 7 | - cron: "0 0 * * *" 8 | 9 | jobs: 10 | cpu-tests: 11 | 12 | runs-on: ${{ matrix.os }} 13 | strategy: 14 | fail-fast: false 15 | matrix: 16 | os: [ubuntu-latest, macos-latest, windows-latest] 17 | python-version: [3.8] 18 | model: ['yolov5s'] # models to test 19 | 20 | # Timeout: https://stackoverflow.com/a/59076067/4521646 21 | timeout-minutes: 50 22 | steps: 23 | - uses: actions/checkout@v2 24 | - name: Set up Python ${{ matrix.python-version }} 25 | uses: actions/setup-python@v2 26 | with: 27 | python-version: ${{ matrix.python-version }} 28 | 29 | # Note: This uses an internal pip API and may not always work 30 | # https://github.com/actions/cache/blob/master/examples.md#multiple-oss-in-a-workflow 31 | - name: Get pip cache 32 | id: pip-cache 33 | run: | 34 | python -c "from pip._internal.locations import USER_CACHE_DIR; print('::set-output name=dir::' + USER_CACHE_DIR)" 35 | 36 | - name: Cache pip 37 | uses: actions/cache@v1 38 | with: 39 | path: ${{ steps.pip-cache.outputs.dir }} 40 | key: ${{ runner.os }}-${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt') }} 41 | restore-keys: | 42 | ${{ runner.os }}-${{ matrix.python-version }}-pip- 43 | 44 | - name: Install dependencies 45 | run: | 46 | python -m pip install --upgrade pip 47 | pip install -qr requirements.txt -f https://download.pytorch.org/whl/cpu/torch_stable.html 48 | pip install -q onnx 49 | python --version 50 | pip --version 51 | pip list 52 | shell: bash 53 | 54 | - name: Download data 55 | run: | 56 | curl -L -o temp.zip https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 57 | unzip -q temp.zip -d ../ 58 | rm temp.zip 59 | 60 | - name: Tests workflow 61 | run: | 62 | export PYTHONPATH="$PWD" # to run *.py. files in subdirectories 63 | di=cpu # inference devices # define device 64 | 65 | # train 66 | python train.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --cfg models/${{ matrix.model }}.yaml --epochs 1 --device $di 67 | # detect 68 | python detect.py --weights weights/${{ matrix.model }}.pt --device $di 69 | python detect.py --weights runs/exp0/weights/last.pt --device $di 70 | # test 71 | python test.py --img 256 --batch 8 --weights weights/${{ matrix.model }}.pt --device $di 72 | python test.py --img 256 --batch 8 --weights runs/exp0/weights/last.pt --device $di 73 | 74 | python models/yolo.py --cfg models/${{ matrix.model }}.yaml # inspect 75 | python models/export.py --img 256 --batch 1 --weights weights/${{ matrix.model }}.pt # export 76 | shell: bash 77 | -------------------------------------------------------------------------------- /yolov5/.github/workflows/greetings.yml: -------------------------------------------------------------------------------- 1 | name: Greetings 2 | 3 | on: [pull_request_target, issues] 4 | 5 | jobs: 6 | greeting: 7 | runs-on: ubuntu-latest 8 | steps: 9 | - uses: actions/first-interaction@v1 10 | with: 11 | repo-token: ${{ secrets.GITHUB_TOKEN }} 12 | pr-message: | 13 | Hello @${{ github.actor }}, thank you for submitting a PR! To allow your work to be integrated as seamlessly as possible, we advise you to: 14 | - Verify your PR is **up-to-date with origin/master.** If your PR is behind origin/master update by running the following, replacing 'feature' with the name of your local branch: 15 | ```bash 16 | git remote add upstream https://github.com/ultralytics/yolov5.git 17 | git fetch upstream 18 | git checkout feature # <----- replace 'feature' with local branch name 19 | git rebase upstream/master 20 | git push -u origin -f 21 | ``` 22 | - Verify all Continuous Integration (CI) **checks are passing**. 23 | - Reduce changes to the absolute **minimum** required for your bug fix or feature addition. _"It is not daily increase but daily decrease, hack away the unessential. The closer to the source, the less wastage there is."_ -Bruce Lee 24 | 25 | issue-message: | 26 | Hello @${{ github.actor }}, thank you for your interest in our work! Please visit our [Custom Training Tutorial](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data) to get started, and see our [Jupyter Notebook](https://github.com/ultralytics/yolov5/blob/master/tutorial.ipynb) Open In Colab, [Docker Image](https://hub.docker.com/r/ultralytics/yolov5), and [Google Cloud Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) for example environments. 27 | 28 | If this is a bug report, please provide screenshots and **minimum viable code to reproduce your issue**, otherwise we can not help you. 29 | 30 | If this is a custom model or data training question, please note Ultralytics does **not** provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as: 31 | - **Cloud-based AI** systems operating on **hundreds of HD video streams in realtime.** 32 | - **Edge AI** integrated into custom iOS and Android apps for realtime **30 FPS video inference.** 33 | - **Custom data training**, hyperparameter evolution, and model exportation to any destination. 34 | 35 | For more information please visit https://www.ultralytics.com. 36 | -------------------------------------------------------------------------------- /yolov5/.github/workflows/rebase.yml: -------------------------------------------------------------------------------- 1 | name: Automatic Rebase 2 | # https://github.com/marketplace/actions/automatic-rebase 3 | 4 | on: 5 | issue_comment: 6 | types: [created] 7 | 8 | jobs: 9 | rebase: 10 | name: Rebase 11 | if: github.event.issue.pull_request != '' && contains(github.event.comment.body, '/rebase') 12 | runs-on: ubuntu-latest 13 | steps: 14 | - name: Checkout the latest code 15 | uses: actions/checkout@v2 16 | with: 17 | fetch-depth: 0 18 | - name: Automatic Rebase 19 | uses: cirrus-actions/rebase@1.3.1 20 | env: 21 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 22 | -------------------------------------------------------------------------------- /yolov5/.github/workflows/stale.yml: -------------------------------------------------------------------------------- 1 | name: Close stale issues 2 | on: 3 | schedule: 4 | - cron: "0 0 * * *" 5 | 6 | jobs: 7 | stale: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/stale@v1 11 | with: 12 | repo-token: ${{ secrets.GITHUB_TOKEN }} 13 | stale-issue-message: 'This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.' 14 | stale-pr-message: 'This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.' 15 | days-before-stale: 30 16 | days-before-close: 5 17 | exempt-issue-label: 'documentation,tutorial' 18 | operations-per-run: 100 # The maximum number of operations per run, used to control rate limiting. 19 | -------------------------------------------------------------------------------- /yolov5/.ipynb_checkpoints/detect-checkpoint.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import platform 4 | import shutil 5 | import time 6 | from pathlib import Path 7 | 8 | import cv2 9 | import torch 10 | import torch.backends.cudnn as cudnn 11 | from numpy import random 12 | 13 | from models.experimental import attempt_load 14 | from utils.datasets import LoadStreams, LoadImages 15 | from utils.general import ( 16 | check_img_size, non_max_suppression, apply_classifier, scale_coords, 17 | xyxy2xywh, plot_one_box, strip_optimizer, set_logging) 18 | from utils.torch_utils import select_device, load_classifier, time_synchronized 19 | 20 | 21 | def detect(save_img=False): 22 | out, source, weights, view_img, save_txt, imgsz = \ 23 | opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size 24 | webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt') 25 | 26 | # Initialize 27 | set_logging() 28 | device = select_device(opt.device) 29 | if os.path.exists(out): 30 | shutil.rmtree(out) # delete output folder 31 | os.makedirs(out) # make new output folder 32 | half = device.type != 'cpu' # half precision only supported on CUDA 33 | 34 | # Load model 35 | model = attempt_load(weights, map_location=device) # load FP32 model 36 | imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size 37 | if half: 38 | model.half() # to FP16 39 | 40 | # Second-stage classifier 41 | classify = False 42 | if classify: 43 | modelc = load_classifier(name='resnet101', n=2) # initialize 44 | modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights 45 | modelc.to(device).eval() 46 | 47 | # Set Dataloader 48 | vid_path, vid_writer = None, None 49 | if webcam: 50 | view_img = True 51 | cudnn.benchmark = True # set True to speed up constant image size inference 52 | dataset = LoadStreams(source, img_size=imgsz) 53 | else: 54 | save_img = True 55 | dataset = LoadImages(source, img_size=imgsz) 56 | 57 | # Get names and colors 58 | names = model.module.names if hasattr(model, 'module') else model.names 59 | colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] 60 | 61 | # Run inference 62 | t0 = time.time() 63 | img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img 64 | _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once 65 | for path, img, im0s, vid_cap in dataset: 66 | img = torch.from_numpy(img).to(device) 67 | img = img.half() if half else img.float() # uint8 to fp16/32 68 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 69 | if img.ndimension() == 3: 70 | img = img.unsqueeze(0) 71 | 72 | # Inference 73 | t1 = time_synchronized() 74 | pred = model(img, augment=opt.augment)[0] 75 | 76 | # Apply NMS 77 | print(pred) 78 | pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms) 79 | 80 | t2 = time_synchronized() 81 | 82 | # Apply Classifier 83 | if classify: 84 | pred = apply_classifier(pred, modelc, img, im0s) 85 | 86 | # Process detections 87 | for i, det in enumerate(pred): # detections per image 88 | if webcam: # batch_size >= 1 89 | p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() 90 | else: 91 | p, s, im0 = path, '', im0s 92 | 93 | save_path = str(Path(out) / Path(p).name) 94 | txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '') 95 | s += '%gx%g ' % img.shape[2:] # print string 96 | gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh 97 | if det is not None and len(det): 98 | # Rescale boxes from img_size to im0 size 99 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() 100 | 101 | # Print results 102 | for c in det[:, -1].unique(): 103 | n = (det[:, -1] == c).sum() # detections per class 104 | s += '%g %ss, ' % (n, names[int(c)]) # add to string 105 | 106 | # Write results 107 | for *xyxy, conf, cls in reversed(det): 108 | if save_txt: # Write to file 109 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 110 | with open(txt_path + '.txt', 'a') as f: 111 | f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format 112 | 113 | if save_img or view_img: # Add bbox to image 114 | label = '%s %.2f' % (names[int(cls)], conf) 115 | plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) 116 | 117 | # Print time (inference + NMS) 118 | print('%sDone. (%.3fs)' % (s, t2 - t1)) 119 | 120 | # Stream results 121 | if view_img: 122 | cv2.imshow(p, im0) 123 | if cv2.waitKey(1) == ord('q'): # q to quit 124 | raise StopIteration 125 | 126 | # Save results (image with detections) 127 | if save_img: 128 | if dataset.mode == 'images': 129 | cv2.imwrite(save_path, im0) 130 | else: 131 | if vid_path != save_path: # new video 132 | vid_path = save_path 133 | if isinstance(vid_writer, cv2.VideoWriter): 134 | vid_writer.release() # release previous video writer 135 | 136 | fourcc = 'mp4v' # output video codec 137 | fps = vid_cap.get(cv2.CAP_PROP_FPS) 138 | w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 139 | h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 140 | vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) 141 | vid_writer.write(im0) 142 | 143 | if save_txt or save_img: 144 | print('Results saved to %s' % Path(out)) 145 | if platform.system() == 'Darwin' and not opt.update: # MacOS 146 | os.system('open ' + save_path) 147 | 148 | print('Done. (%.3fs)' % (time.time() - t0)) 149 | 150 | 151 | if __name__ == '__main__': 152 | parser = argparse.ArgumentParser() 153 | parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 154 | parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam 155 | parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder 156 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 157 | parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') 158 | parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') 159 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 160 | parser.add_argument('--view-img', action='store_true', help='display results') 161 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 162 | parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') 163 | parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') 164 | parser.add_argument('--augment', action='store_true', help='augmented inference') 165 | parser.add_argument('--update', action='store_true', help='update all models') 166 | opt = parser.parse_args() 167 | print(opt) 168 | 169 | with torch.no_grad(): 170 | if opt.update: # update all models (to fix SourceChangeWarning) 171 | for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 172 | detect() 173 | strip_optimizer(opt.weights) 174 | else: 175 | detect() 176 | -------------------------------------------------------------------------------- /yolov5/.ipynb_checkpoints/requirements-checkpoint.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # base ---------------------------------------- 4 | Cython 5 | matplotlib>=3.2.2 6 | numpy>=1.18.5 7 | opencv-python>=4.1.2 8 | pillow 9 | PyYAML>=5.3 10 | scipy>=1.4.1 11 | tensorboard>=2.2 12 | torch>=1.6.0 13 | torchvision>=0.7.0 14 | tqdm>=4.41.0 15 | 16 | # coco ---------------------------------------- 17 | # pycocotools>=2.0 18 | 19 | # export -------------------------------------- 20 | # coremltools==4.0b2 21 | # onnx>=1.7.0 22 | 23 | # extras -------------------------------------- 24 | # thop 25 | -------------------------------------------------------------------------------- /yolov5/README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 |   4 | 5 | ![CI CPU testing](https://github.com/ultralytics/yolov5/workflows/CI%20CPU%20testing/badge.svg) 6 | 7 | This repository represents Ultralytics open-source research into future object detection methods, and incorporates our lessons learned and best practices evolved over training thousands of models on custom client datasets with our previous YOLO repository https://github.com/ultralytics/yolov3. **All code and models are under active development, and are subject to modification or deletion without notice.** Use at your own risk. 8 | 9 | ** GPU Speed measures end-to-end time per image averaged over 5000 COCO val2017 images using a V100 GPU with batch size 32, and includes image preprocessing, PyTorch FP16 inference, postprocessing and NMS. EfficientDet data from [google/automl](https://github.com/google/automl) at batch size 8. 10 | 11 | - **August 13, 2020**: [v3.0 release](https://github.com/ultralytics/yolov5/releases/tag/v3.0): nn.Hardswish() activations, data autodownload, native AMP. 12 | - **July 23, 2020**: [v2.0 release](https://github.com/ultralytics/yolov5/releases/tag/v2.0): improved model definition, training and mAP. 13 | - **June 22, 2020**: [PANet](https://arxiv.org/abs/1803.01534) updates: new heads, reduced parameters, improved speed and mAP [364fcfd](https://github.com/ultralytics/yolov5/commit/364fcfd7dba53f46edd4f04c037a039c0a287972). 14 | - **June 19, 2020**: [FP16](https://pytorch.org/docs/stable/nn.html#torch.nn.Module.half) as new default for smaller checkpoints and faster inference [d4c6674](https://github.com/ultralytics/yolov5/commit/d4c6674c98e19df4c40e33a777610a18d1961145). 15 | - **June 9, 2020**: [CSP](https://github.com/WongKinYiu/CrossStagePartialNetworks) updates: improved speed, size, and accuracy (credit to @WongKinYiu for CSP). 16 | - **May 27, 2020**: Public release. YOLOv5 models are SOTA among all known YOLO implementations. 17 | - **April 1, 2020**: Start development of future compound-scaled [YOLOv3](https://github.com/ultralytics/yolov3)/[YOLOv4](https://github.com/AlexeyAB/darknet)-based PyTorch models. 18 | 19 | 20 | ## Pretrained Checkpoints 21 | 22 | | Model | APval | APtest | AP50 | SpeedGPU | FPSGPU || params | FLOPS | 23 | |---------- |------ |------ |------ | -------- | ------| ------ |------ | :------: | 24 | | [YOLOv5s](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 37.0 | 37.0 | 56.2 | **2.4ms** | **416** || 7.5M | 13.2B 25 | | [YOLOv5m](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 44.3 | 44.3 | 63.2 | 3.4ms | 294 || 21.8M | 39.4B 26 | | [YOLOv5l](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 47.7 | 47.7 | 66.5 | 4.4ms | 227 || 47.8M | 88.1B 27 | | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | **49.2** | **49.2** | **67.7** | 6.9ms | 145 || 89.0M | 166.4B 28 | | | | | | | || | 29 | | [YOLOv5x](https://github.com/ultralytics/yolov5/releases/tag/v3.0) + TTA|**50.8**| **50.8** | **68.9** | 25.5ms | 39 || 89.0M | 354.3B 30 | | | | | | | || | 31 | | [YOLOv3-SPP](https://github.com/ultralytics/yolov5/releases/tag/v3.0) | 45.6 | 45.5 | 65.2 | 4.5ms | 222 || 63.0M | 118.0B 32 | 33 | ** APtest denotes COCO [test-dev2017](http://cocodataset.org/#upload) server results, all other AP results in the table denote val2017 accuracy. 34 | ** All AP numbers are for single-model single-scale without ensemble or test-time augmentation. **Reproduce** by `python test.py --data coco.yaml --img 640 --conf 0.001` 35 | ** SpeedGPU measures end-to-end time per image averaged over 5000 COCO val2017 images using a GCP [n1-standard-16](https://cloud.google.com/compute/docs/machine-types#n1_standard_machine_types) instance with one V100 GPU, and includes image preprocessing, PyTorch FP16 image inference at --batch-size 32 --img-size 640, postprocessing and NMS. Average NMS time included in this chart is 1-2ms/img. **Reproduce** by `python test.py --data coco.yaml --img 640 --conf 0.1` 36 | ** All checkpoints are trained to 300 epochs with default settings and hyperparameters (no autoaugmentation). 37 | ** Test Time Augmentation ([TTA](https://github.com/ultralytics/yolov5/issues/303)) runs at 3 image sizes. **Reproduce** by `python test.py --data coco.yaml --img 832 --augment` 38 | 39 | ## Requirements 40 | 41 | Python 3.8 or later with all [requirements.txt](https://github.com/ultralytics/yolov5/blob/master/requirements.txt) dependencies installed, including `torch>=1.6`. To install run: 42 | ```bash 43 | $ pip install -r requirements.txt 44 | ``` 45 | 46 | 47 | ## Tutorials 48 | 49 | * [Train Custom Data](https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data) 50 | * [Multi-GPU Training](https://github.com/ultralytics/yolov5/issues/475) 51 | * [PyTorch Hub](https://github.com/ultralytics/yolov5/issues/36) 52 | * [ONNX and TorchScript Export](https://github.com/ultralytics/yolov5/issues/251) 53 | * [Test-Time Augmentation (TTA)](https://github.com/ultralytics/yolov5/issues/303) 54 | * [Model Ensembling](https://github.com/ultralytics/yolov5/issues/318) 55 | * [Model Pruning/Sparsity](https://github.com/ultralytics/yolov5/issues/304) 56 | * [Hyperparameter Evolution](https://github.com/ultralytics/yolov5/issues/607) 57 | * [TensorRT Deployment](https://github.com/wang-xinyu/tensorrtx) 58 | 59 | 60 | ## Environments 61 | 62 | YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled): 63 | 64 | - **Google Colab Notebook** with free GPU: Open In Colab 65 | - **Kaggle Notebook** with free GPU: [https://www.kaggle.com/ultralytics/yolov5](https://www.kaggle.com/ultralytics/yolov5) 66 | - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart) 67 | - **Docker Image** https://hub.docker.com/r/ultralytics/yolov5. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) ![Docker Pulls](https://img.shields.io/docker/pulls/ultralytics/yolov5?logo=docker) 68 | 69 | 70 | ## Inference 71 | 72 | Inference can be run on most common media formats. Model [checkpoints](https://drive.google.com/open?id=1Drs_Aiu7xx6S-ix95f9kNsA6ueKRpN2J) are downloaded automatically if available. Results are saved to `./inference/output`. 73 | ```bash 74 | $ python detect.py --source 0 # webcam 75 | file.jpg # image 76 | file.mp4 # video 77 | path/ # directory 78 | path/*.jpg # glob 79 | rtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa # rtsp stream 80 | http://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8 # http stream 81 | ``` 82 | 83 | To run inference on examples in the `./inference/images` folder: 84 | 85 | ```bash 86 | $ python detect.py --source ./inference/images/ --weights yolov5s.pt --conf 0.4 87 | 88 | Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.4, device='', fourcc='mp4v', half=False, img_size=640, iou_thres=0.5, output='inference/output', save_txt=False, source='./inference/images/', view_img=False, weights='yolov5s.pt') 89 | Using CUDA device0 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', total_memory=16280MB) 90 | 91 | Downloading https://drive.google.com/uc?export=download&id=1R5T6rIyy3lLwgFXNms8whc-387H0tMQO as yolov5s.pt... Done (2.6s) 92 | 93 | image 1/2 inference/images/bus.jpg: 640x512 3 persons, 1 buss, Done. (0.009s) 94 | image 2/2 inference/images/zidane.jpg: 384x640 2 persons, 2 ties, Done. (0.009s) 95 | Results saved to /content/yolov5/inference/output 96 | ``` 97 | 98 | 99 | 100 | 101 | ## Training 102 | 103 | Download [COCO](https://github.com/ultralytics/yolov5/blob/master/data/scripts/get_coco.sh) and run command below. Training times for YOLOv5s/m/l/x are 2/4/6/8 days on a single V100 (multi-GPU times faster). Use the largest `--batch-size` your GPU allows (batch sizes shown for 16 GB devices). 104 | ```bash 105 | $ python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64 106 | yolov5m 40 107 | yolov5l 24 108 | yolov5x 16 109 | ``` 110 | 111 | 112 | 113 | ## Citation 114 | 115 | [![DOI](https://zenodo.org/badge/264818686.svg)](https://zenodo.org/badge/latestdoi/264818686) 116 | 117 | 118 | ## About Us 119 | 120 | Ultralytics is a U.S.-based particle physics and AI startup with over 6 years of expertise supporting government, academic and business clients. We offer a wide range of vision AI services, spanning from simple expert advice up to delivery of fully customized, end-to-end production solutions, including: 121 | - **Cloud-based AI** systems operating on **hundreds of HD video streams in realtime.** 122 | - **Edge AI** integrated into custom iOS and Android apps for realtime **30 FPS video inference.** 123 | - **Custom data training**, hyperparameter evolution, and model exportation to any destination. 124 | 125 | For business inquiries and professional support requests please visit us at https://www.ultralytics.com. 126 | 127 | 128 | ## Contact 129 | 130 | **Issues should be raised directly in the repository.** For business inquiries or professional support requests please visit https://www.ultralytics.com or email Glenn Jocher at glenn.jocher@ultralytics.com. 131 | -------------------------------------------------------------------------------- /yolov5/__pycache__/test.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/__pycache__/test.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/data/.ipynb_checkpoints/coco128-checkpoint.yaml: -------------------------------------------------------------------------------- 1 | # COCO 2017 dataset http://cocodataset.org - first 128 training images 2 | # Train command: python train.py --data coco128.yaml 3 | # Default dataset location is next to /yolov5: 4 | # /parent_folder 5 | # /coco128 6 | # /yolov5 7 | 8 | 9 | # # download command/URL (optional) 10 | # download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 11 | 12 | # # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] 13 | # train: ../coco128/images/train2017/ # 128 images 14 | # val: ../coco128/images/train2017/ # 128 images 15 | 16 | # # number of classes 17 | # nc: 80 18 | 19 | # # class names 20 | # names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 21 | # 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 22 | # 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 23 | # 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 24 | # 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 25 | # 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 26 | # 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 27 | # 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 28 | # 'hair drier', 'toothbrush'] 29 | train: /project/train/images/train2017/ # 128 images 30 | val: /project/train/images/train2017/ # 128 images 31 | 32 | # number of classes 33 | nc: 1 34 | 35 | # class names 36 | names: ['fire'] 37 | -------------------------------------------------------------------------------- /yolov5/data/.ipynb_checkpoints/hyp.finetune-checkpoint.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for VOC fine-tuning 2 | # python train.py --batch 64 --cfg '' --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50 3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 4 | 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | momentum: 0.937 # SGD momentum/Adam beta1 8 | weight_decay: 0.0005 # optimizer weight decay 5e-4 9 | giou: 0.05 # GIoU loss gain 10 | cls: 0.5 # cls loss gain 11 | cls_pw: 1.0 # cls BCELoss positive_weight 12 | obj: 1.0 # obj loss gain (scale with pixels) 13 | obj_pw: 1.0 # obj BCELoss positive_weight 14 | iou_t: 0.20 # IoU training threshold 15 | anchor_t: 4.0 # anchor-multiple threshold 16 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 17 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 18 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 19 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 20 | degrees: 0.0 # image rotation (+/- deg) 21 | translate: 0.1 # image translation (+/- fraction) 22 | scale: 0.5 # image scale (+/- gain) 23 | shear: 0.0 # image shear (+/- deg) 24 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 25 | flipud: 0.0 # image flip up-down (probability) 26 | fliplr: 0.5 # image flip left-right (probability) 27 | mixup: 0.0 # image mixup (probability) 28 | -------------------------------------------------------------------------------- /yolov5/data/.ipynb_checkpoints/hyp.scratch-checkpoint.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for COCO training from scratch 2 | # python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300 3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 4 | 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | momentum: 0.937 # SGD momentum/Adam beta1 8 | weight_decay: 0.0005 # optimizer weight decay 5e-4 9 | giou: 0.05 # GIoU loss gain 10 | cls: 0.5 # cls loss gain 11 | cls_pw: 1.0 # cls BCELoss positive_weight 12 | obj: 1.0 # obj loss gain (scale with pixels) 13 | obj_pw: 1.0 # obj BCELoss positive_weight 14 | iou_t: 0.20 # IoU training threshold 15 | anchor_t: 4.0 # anchor-multiple threshold 16 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 17 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 18 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 19 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 20 | degrees: 0.0 # image rotation (+/- deg) 21 | translate: 0.1 # image translation (+/- fraction) 22 | scale: 0.5 # image scale (+/- gain) 23 | shear: 0.0 # image shear (+/- deg) 24 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 25 | flipud: 0.0 # image flip up-down (probability) 26 | fliplr: 0.5 # image flip left-right (probability) 27 | mixup: 0.0 # image mixup (probability) 28 | -------------------------------------------------------------------------------- /yolov5/data/coco.yaml: -------------------------------------------------------------------------------- 1 | # COCO 2017 dataset http://cocodataset.org 2 | # Train command: python train.py --data coco.yaml 3 | # Default dataset location is next to /yolov5: 4 | # /parent_folder 5 | # /coco 6 | # /yolov5 7 | 8 | 9 | # download command/URL (optional) 10 | download: bash data/scripts/get_coco.sh 11 | 12 | # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] 13 | train: ../coco/train2017.txt # 118287 images 14 | val: ../coco/val2017.txt # 5000 images 15 | test: ../coco/test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794 16 | 17 | # number of classes 18 | nc: 80 19 | 20 | # class names 21 | names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 22 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 23 | 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 24 | 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 25 | 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 26 | 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 27 | 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 28 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 29 | 'hair drier', 'toothbrush'] 30 | 31 | # Print classes 32 | # with open('data/coco.yaml') as f: 33 | # d = yaml.load(f, Loader=yaml.FullLoader) # dict 34 | # for i, x in enumerate(d['names']): 35 | # print(i, x) 36 | -------------------------------------------------------------------------------- /yolov5/data/coco128.yaml: -------------------------------------------------------------------------------- 1 | # COCO 2017 dataset http://cocodataset.org - first 128 training images 2 | # Train command: python train.py --data coco128.yaml 3 | # Default dataset location is next to /yolov5: 4 | # /parent_folder 5 | # /coco128 6 | # /yolov5 7 | 8 | 9 | # # download command/URL (optional) 10 | # download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip 11 | 12 | # # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] 13 | # train: ../coco128/images/train2017/ # 128 images 14 | # val: ../coco128/images/train2017/ # 128 images 15 | 16 | # # number of classes 17 | # nc: 80 18 | 19 | # # class names 20 | # names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 21 | # 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 22 | # 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 23 | # 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 24 | # 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 25 | # 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 26 | # 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 27 | # 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 28 | # 'hair drier', 'toothbrush'] 29 | train: /project/train/images/train2017/ # 128 images 30 | val: /project/train/images/train2017/ # 128 images 31 | 32 | # number of classes 33 | nc: 1 34 | 35 | # class names 36 | names: ['fire'] 37 | -------------------------------------------------------------------------------- /yolov5/data/hyp.finetune.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for VOC fine-tuning 2 | # python train.py --batch 64 --cfg '' --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50 3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 4 | 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | momentum: 0.937 # SGD momentum/Adam beta1 8 | weight_decay: 0.0005 # optimizer weight decay 5e-4 9 | giou: 0.05 # GIoU loss gain 10 | cls: 0.5 # cls loss gain 11 | cls_pw: 1.0 # cls BCELoss positive_weight 12 | obj: 1.0 # obj loss gain (scale with pixels) 13 | obj_pw: 1.0 # obj BCELoss positive_weight 14 | iou_t: 0.20 # IoU training threshold 15 | anchor_t: 4.0 # anchor-multiple threshold 16 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 17 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 18 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 19 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 20 | degrees: 0.0 # image rotation (+/- deg) 21 | translate: 0.1 # image translation (+/- fraction) 22 | scale: 0.5 # image scale (+/- gain) 23 | shear: 0.0 # image shear (+/- deg) 24 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 25 | flipud: 0.0 # image flip up-down (probability) 26 | fliplr: 0.5 # image flip left-right (probability) 27 | mixup: 0.0 # image mixup (probability) 28 | -------------------------------------------------------------------------------- /yolov5/data/hyp.scratch.yaml: -------------------------------------------------------------------------------- 1 | # Hyperparameters for COCO training from scratch 2 | # python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300 3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials 4 | 5 | 6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3) 7 | momentum: 0.937 # SGD momentum/Adam beta1 8 | weight_decay: 0.0005 # optimizer weight decay 5e-4 9 | giou: 0.05 # GIoU loss gain 10 | cls: 0.5 # cls loss gain 11 | cls_pw: 1.0 # cls BCELoss positive_weight 12 | obj: 1.0 # obj loss gain (scale with pixels) 13 | obj_pw: 1.0 # obj BCELoss positive_weight 14 | iou_t: 0.20 # IoU training threshold 15 | anchor_t: 4.0 # anchor-multiple threshold 16 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5) 17 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction) 18 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction) 19 | hsv_v: 0.4 # image HSV-Value augmentation (fraction) 20 | degrees: 0.0 # image rotation (+/- deg) 21 | translate: 0.1 # image translation (+/- fraction) 22 | scale: 0.5 # image scale (+/- gain) 23 | shear: 0.0 # image shear (+/- deg) 24 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001 25 | flipud: 0.0 # image flip up-down (probability) 26 | fliplr: 0.5 # image flip left-right (probability) 27 | mixup: 0.0 # image mixup (probability) 28 | -------------------------------------------------------------------------------- /yolov5/data/scripts/get_coco.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # COCO 2017 dataset http://cocodataset.org 3 | # Download command: bash data/scripts/get_coco.sh 4 | # Train command: python train.py --data coco.yaml 5 | # Default dataset location is next to /yolov5: 6 | # /parent_folder 7 | # /coco 8 | # /yolov5 9 | 10 | # Download/unzip labels 11 | echo 'Downloading COCO 2017 labels ...' 12 | d='../' # unzip directory 13 | f='coco2017labels.zip' && curl -L https://github.com/ultralytics/yolov5/releases/download/v1.0/$f -o $f 14 | unzip -q $f -d $d && rm $f 15 | 16 | # Download/unzip images 17 | echo 'Downloading COCO 2017 images ...' 18 | d='../coco/images' # unzip directory 19 | f='train2017.zip' && curl http://images.cocodataset.org/zips/$f -o $f && unzip -q $f -d $d && rm $f # 19G, 118k images 20 | f='val2017.zip' && curl http://images.cocodataset.org/zips/$f -o $f && unzip -q $f -d $d && rm $f # 1G, 5k images 21 | # f='test2017.zip' && curl http://images.cocodataset.org/zips/$f -o $f && unzip -q $f -d $d && rm $f # 7G, 41k images 22 | -------------------------------------------------------------------------------- /yolov5/data/scripts/get_voc.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC/ 3 | # Download command: bash data/scripts/get_voc.sh 4 | # Train command: python train.py --data voc.yaml 5 | # Default dataset location is next to /yolov5: 6 | # /parent_folder 7 | # /VOC 8 | # /yolov5 9 | 10 | start=$(date +%s) 11 | 12 | # handle optional download dir 13 | if [ -z "$1" ]; then 14 | # navigate to ~/tmp 15 | echo "navigating to ../tmp/ ..." 16 | mkdir -p ../tmp 17 | cd ../tmp/ 18 | else 19 | # check if is valid directory 20 | if [ ! -d $1 ]; then 21 | echo $1 "is not a valid directory" 22 | exit 0 23 | fi 24 | echo "navigating to" $1 "..." 25 | cd $1 26 | fi 27 | 28 | echo "Downloading VOC2007 trainval ..." 29 | # Download data 30 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar 31 | echo "Downloading VOC2007 test data ..." 32 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar 33 | echo "Done downloading." 34 | 35 | # Extract data 36 | echo "Extracting trainval ..." 37 | tar -xf VOCtrainval_06-Nov-2007.tar 38 | echo "Extracting test ..." 39 | tar -xf VOCtest_06-Nov-2007.tar 40 | echo "removing tars ..." 41 | rm VOCtrainval_06-Nov-2007.tar 42 | rm VOCtest_06-Nov-2007.tar 43 | 44 | end=$(date +%s) 45 | runtime=$((end - start)) 46 | 47 | echo "Completed in" $runtime "seconds" 48 | 49 | start=$(date +%s) 50 | 51 | # handle optional download dir 52 | if [ -z "$1" ]; then 53 | # navigate to ~/tmp 54 | echo "navigating to ../tmp/ ..." 55 | mkdir -p ../tmp 56 | cd ../tmp/ 57 | else 58 | # check if is valid directory 59 | if [ ! -d $1 ]; then 60 | echo $1 "is not a valid directory" 61 | exit 0 62 | fi 63 | echo "navigating to" $1 "..." 64 | cd $1 65 | fi 66 | 67 | echo "Downloading VOC2012 trainval ..." 68 | # Download data 69 | curl -LO http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar 70 | echo "Done downloading." 71 | 72 | # Extract data 73 | echo "Extracting trainval ..." 74 | tar -xf VOCtrainval_11-May-2012.tar 75 | echo "removing tar ..." 76 | rm VOCtrainval_11-May-2012.tar 77 | 78 | end=$(date +%s) 79 | runtime=$((end - start)) 80 | 81 | echo "Completed in" $runtime "seconds" 82 | 83 | cd ../tmp 84 | echo "Spliting dataset..." 85 | python3 - "$@" <train.txt 145 | cat 2007_train.txt 2007_val.txt 2007_test.txt 2012_train.txt 2012_val.txt >train.all.txt 146 | 147 | python3 - "$@" <= 1 89 | p, s, im0 = path[i], '%g: ' % i, im0s[i].copy() 90 | else: 91 | p, s, im0 = path, '', im0s 92 | 93 | save_path = str(Path(out) / Path(p).name) 94 | txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '') 95 | s += '%gx%g ' % img.shape[2:] # print string 96 | gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh 97 | if det is not None and len(det): 98 | # Rescale boxes from img_size to im0 size 99 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() 100 | 101 | # Print results 102 | for c in det[:, -1].unique(): 103 | n = (det[:, -1] == c).sum() # detections per class 104 | s += '%g %ss, ' % (n, names[int(c)]) # add to string 105 | 106 | # Write results 107 | for *xyxy, conf, cls in reversed(det): 108 | if save_txt: # Write to file 109 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 110 | with open(txt_path + '.txt', 'a') as f: 111 | f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format 112 | 113 | if save_img or view_img: # Add bbox to image 114 | label = '%s %.2f' % (names[int(cls)], conf) 115 | plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) 116 | 117 | # Print time (inference + NMS) 118 | print('%sDone. (%.3fs)' % (s, t2 - t1)) 119 | 120 | # Stream results 121 | if view_img: 122 | cv2.imshow(p, im0) 123 | if cv2.waitKey(1) == ord('q'): # q to quit 124 | raise StopIteration 125 | 126 | # Save results (image with detections) 127 | if save_img: 128 | if dataset.mode == 'images': 129 | cv2.imwrite(save_path, im0) 130 | else: 131 | if vid_path != save_path: # new video 132 | vid_path = save_path 133 | if isinstance(vid_writer, cv2.VideoWriter): 134 | vid_writer.release() # release previous video writer 135 | 136 | fourcc = 'mp4v' # output video codec 137 | fps = vid_cap.get(cv2.CAP_PROP_FPS) 138 | w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) 139 | h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) 140 | vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) 141 | vid_writer.write(im0) 142 | 143 | if save_txt or save_img: 144 | print('Results saved to %s' % Path(out)) 145 | if platform.system() == 'Darwin' and not opt.update: # MacOS 146 | os.system('open ' + save_path) 147 | 148 | print('Done. (%.3fs)' % (time.time() - t0)) 149 | 150 | 151 | if __name__ == '__main__': 152 | parser = argparse.ArgumentParser() 153 | parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 154 | parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam 155 | parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder 156 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 157 | parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold') 158 | parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS') 159 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 160 | parser.add_argument('--view-img', action='store_true', help='display results') 161 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 162 | parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') 163 | parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS') 164 | parser.add_argument('--augment', action='store_true', help='augmented inference') 165 | parser.add_argument('--update', action='store_true', help='update all models') 166 | opt = parser.parse_args() 167 | print(opt) 168 | 169 | with torch.no_grad(): 170 | if opt.update: # update all models (to fix SourceChangeWarning) 171 | for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 172 | detect() 173 | strip_optimizer(opt.weights) 174 | else: 175 | detect() 176 | -------------------------------------------------------------------------------- /yolov5/hubconf.py: -------------------------------------------------------------------------------- 1 | """File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/ 2 | 3 | Usage: 4 | import torch 5 | model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80) 6 | """ 7 | 8 | dependencies = ['torch', 'yaml'] 9 | import os 10 | 11 | import torch 12 | 13 | from models.yolo import Model 14 | from utils.google_utils import attempt_download 15 | 16 | 17 | def create(name, pretrained, channels, classes): 18 | """Creates a specified YOLOv5 model 19 | 20 | Arguments: 21 | name (str): name of model, i.e. 'yolov5s' 22 | pretrained (bool): load pretrained weights into the model 23 | channels (int): number of input channels 24 | classes (int): number of model classes 25 | 26 | Returns: 27 | pytorch model 28 | """ 29 | config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) # model.yaml path 30 | try: 31 | model = Model(config, channels, classes) 32 | if pretrained: 33 | ckpt = '%s.pt' % name # checkpoint filename 34 | attempt_download(ckpt) # download if not found locally 35 | state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32 36 | state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter 37 | model.load_state_dict(state_dict, strict=False) # load 38 | return model 39 | 40 | except Exception as e: 41 | help_url = 'https://github.com/ultralytics/yolov5/issues/36' 42 | s = 'Cache maybe be out of date, deleting cache and retrying may solve this. See %s for help.' % help_url 43 | raise Exception(s) from e 44 | 45 | 46 | def yolov5s(pretrained=False, channels=3, classes=80): 47 | """YOLOv5-small model from https://github.com/ultralytics/yolov5 48 | 49 | Arguments: 50 | pretrained (bool): load pretrained weights into the model, default=False 51 | channels (int): number of input channels, default=3 52 | classes (int): number of model classes, default=80 53 | 54 | Returns: 55 | pytorch model 56 | """ 57 | return create('yolov5s', pretrained, channels, classes) 58 | 59 | 60 | def yolov5m(pretrained=False, channels=3, classes=80): 61 | """YOLOv5-medium model from https://github.com/ultralytics/yolov5 62 | 63 | Arguments: 64 | pretrained (bool): load pretrained weights into the model, default=False 65 | channels (int): number of input channels, default=3 66 | classes (int): number of model classes, default=80 67 | 68 | Returns: 69 | pytorch model 70 | """ 71 | return create('yolov5m', pretrained, channels, classes) 72 | 73 | 74 | def yolov5l(pretrained=False, channels=3, classes=80): 75 | """YOLOv5-large model from https://github.com/ultralytics/yolov5 76 | 77 | Arguments: 78 | pretrained (bool): load pretrained weights into the model, default=False 79 | channels (int): number of input channels, default=3 80 | classes (int): number of model classes, default=80 81 | 82 | Returns: 83 | pytorch model 84 | """ 85 | return create('yolov5l', pretrained, channels, classes) 86 | 87 | 88 | def yolov5x(pretrained=False, channels=3, classes=80): 89 | """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5 90 | 91 | Arguments: 92 | pretrained (bool): load pretrained weights into the model, default=False 93 | channels (int): number of input channels, default=3 94 | classes (int): number of model classes, default=80 95 | 96 | Returns: 97 | pytorch model 98 | """ 99 | return create('yolov5x', pretrained, channels, classes) 100 | -------------------------------------------------------------------------------- /yolov5/inference/images/bus.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/inference/images/bus.jpg -------------------------------------------------------------------------------- /yolov5/inference/images/zidane.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/inference/images/zidane.jpg -------------------------------------------------------------------------------- /yolov5/inference/output/.ipynb_checkpoints/2209a117-checkpoint.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/inference/output/.ipynb_checkpoints/2209a117-checkpoint.jpg -------------------------------------------------------------------------------- /yolov5/inference/output/2209a117.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/inference/output/2209a117.jpg -------------------------------------------------------------------------------- /yolov5/models/.ipynb_checkpoints/common-checkpoint.py: -------------------------------------------------------------------------------- 1 | # This file contains modules common to various models 2 | import math 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | def autopad(k, p=None): # kernel, padding 9 | # Pad to 'same' 10 | if p is None: 11 | p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad 12 | return p 13 | 14 | 15 | def DWConv(c1, c2, k=1, s=1, act=True): 16 | # Depthwise convolution 17 | return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) 18 | 19 | 20 | class Conv(nn.Module): 21 | # Standard convolution 22 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 23 | super(Conv, self).__init__() 24 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) 25 | self.bn = nn.BatchNorm2d(c2) 26 | self.act = nn.LeakyReLU(0.1) if act else nn.Identity() 27 | 28 | def forward(self, x): 29 | return self.act(self.bn(self.conv(x))) 30 | 31 | def fuseforward(self, x): 32 | return self.act(self.conv(x)) 33 | 34 | 35 | class Bottleneck(nn.Module): 36 | # Standard bottleneck 37 | def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion 38 | super(Bottleneck, self).__init__() 39 | c_ = int(c2 * e) # hidden channels 40 | self.cv1 = Conv(c1, c_, 1, 1) 41 | self.cv2 = Conv(c_, c2, 3, 1, g=g) 42 | self.add = shortcut and c1 == c2 43 | 44 | def forward(self, x): 45 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 46 | 47 | 48 | class BottleneckCSP(nn.Module): 49 | # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks 50 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 51 | super(BottleneckCSP, self).__init__() 52 | c_ = int(c2 * e) # hidden channels 53 | self.cv1 = Conv(c1, c_, 1, 1) 54 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 55 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 56 | self.cv4 = Conv(2 * c_, c2, 1, 1) 57 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 58 | self.act = nn.LeakyReLU(0.1, inplace=True) 59 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) 60 | 61 | def forward(self, x): 62 | y1 = self.cv3(self.m(self.cv1(x))) 63 | y2 = self.cv2(x) 64 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 65 | 66 | 67 | class SPP(nn.Module): 68 | # Spatial pyramid pooling layer used in YOLOv3-SPP 69 | def __init__(self, c1, c2, k=(5, 9, 13)): 70 | super(SPP, self).__init__() 71 | c_ = c1 // 2 # hidden channels 72 | self.cv1 = Conv(c1, c_, 1, 1) 73 | self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) 74 | self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) 75 | 76 | def forward(self, x): 77 | x = self.cv1(x) 78 | return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) 79 | 80 | 81 | class Focus(nn.Module): 82 | # Focus wh information into c-space 83 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 84 | super(Focus, self).__init__() 85 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 86 | 87 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) 88 | return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) 89 | 90 | 91 | class Concat(nn.Module): 92 | # Concatenate a list of tensors along dimension 93 | def __init__(self, dimension=1): 94 | super(Concat, self).__init__() 95 | self.d = dimension 96 | 97 | def forward(self, x): 98 | return torch.cat(x, self.d) 99 | 100 | 101 | class Flatten(nn.Module): 102 | # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions 103 | @staticmethod 104 | def forward(x): 105 | return x.view(x.size(0), -1) 106 | 107 | 108 | class Classify(nn.Module): 109 | # Classification head, i.e. x(b,c1,20,20) to x(b,c2) 110 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups 111 | super(Classify, self).__init__() 112 | self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1) 113 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) # to x(b,c2,1,1) 114 | self.flat = Flatten() 115 | 116 | def forward(self, x): 117 | z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list 118 | return self.flat(self.conv(z)) # flatten to x(b,c2) 119 | -------------------------------------------------------------------------------- /yolov5/models/.ipynb_checkpoints/experimental-checkpoint.py: -------------------------------------------------------------------------------- 1 | # This file contains experimental modules 2 | 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | 7 | from models.common import Conv, DWConv 8 | from utils.google_utils import attempt_download 9 | 10 | 11 | class CrossConv(nn.Module): 12 | # Cross Convolution Downsample 13 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 14 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 15 | super(CrossConv, self).__init__() 16 | c_ = int(c2 * e) # hidden channels 17 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 18 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 19 | self.add = shortcut and c1 == c2 20 | 21 | def forward(self, x): 22 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 23 | 24 | 25 | class C3(nn.Module): 26 | # Cross Convolution CSP 27 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 28 | super(C3, self).__init__() 29 | c_ = int(c2 * e) # hidden channels 30 | self.cv1 = Conv(c1, c_, 1, 1) 31 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 32 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 33 | self.cv4 = Conv(2 * c_, c2, 1, 1) 34 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 35 | self.act = nn.LeakyReLU(0.1, inplace=True) 36 | self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)]) 37 | 38 | def forward(self, x): 39 | y1 = self.cv3(self.m(self.cv1(x))) 40 | y2 = self.cv2(x) 41 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 42 | 43 | 44 | class Sum(nn.Module): 45 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 46 | def __init__(self, n, weight=False): # n: number of inputs 47 | super(Sum, self).__init__() 48 | self.weight = weight # apply weights boolean 49 | self.iter = range(n - 1) # iter object 50 | if weight: 51 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights 52 | 53 | def forward(self, x): 54 | y = x[0] # no weight 55 | if self.weight: 56 | w = torch.sigmoid(self.w) * 2 57 | for i in self.iter: 58 | y = y + x[i + 1] * w[i] 59 | else: 60 | for i in self.iter: 61 | y = y + x[i + 1] 62 | return y 63 | 64 | 65 | class GhostConv(nn.Module): 66 | # Ghost Convolution https://github.com/huawei-noah/ghostnet 67 | def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups 68 | super(GhostConv, self).__init__() 69 | c_ = c2 // 2 # hidden channels 70 | self.cv1 = Conv(c1, c_, k, s, g, act) 71 | self.cv2 = Conv(c_, c_, 5, 1, c_, act) 72 | 73 | def forward(self, x): 74 | y = self.cv1(x) 75 | return torch.cat([y, self.cv2(y)], 1) 76 | 77 | 78 | class GhostBottleneck(nn.Module): 79 | # Ghost Bottleneck https://github.com/huawei-noah/ghostnet 80 | def __init__(self, c1, c2, k, s): 81 | super(GhostBottleneck, self).__init__() 82 | c_ = c2 // 2 83 | self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw 84 | DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw 85 | GhostConv(c_, c2, 1, 1, act=False)) # pw-linear 86 | self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), 87 | Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() 88 | 89 | def forward(self, x): 90 | return self.conv(x) + self.shortcut(x) 91 | 92 | 93 | class MixConv2d(nn.Module): 94 | # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595 95 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): 96 | super(MixConv2d, self).__init__() 97 | groups = len(k) 98 | if equal_ch: # equal c_ per group 99 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices 100 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels 101 | else: # equal weight.numel() per group 102 | b = [c2] + [0] * groups 103 | a = np.eye(groups + 1, groups, k=-1) 104 | a -= np.roll(a, 1, axis=1) 105 | a *= np.array(k) ** 2 106 | a[0] = 1 107 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 108 | 109 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) 110 | self.bn = nn.BatchNorm2d(c2) 111 | self.act = nn.LeakyReLU(0.1, inplace=True) 112 | 113 | def forward(self, x): 114 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 115 | 116 | 117 | class Ensemble(nn.ModuleList): 118 | # Ensemble of models 119 | def __init__(self): 120 | super(Ensemble, self).__init__() 121 | 122 | def forward(self, x, augment=False): 123 | y = [] 124 | for module in self: 125 | y.append(module(x, augment)[0]) 126 | # y = torch.stack(y).max(0)[0] # max ensemble 127 | # y = torch.cat(y, 1) # nms ensemble 128 | y = torch.stack(y).mean(0) # mean ensemble 129 | return y, None # inference, train output 130 | 131 | 132 | def attempt_load(weights, map_location=None): 133 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 134 | model = Ensemble() 135 | for w in weights if isinstance(weights, list) else [weights]: 136 | attempt_download(w) 137 | model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model 138 | 139 | if len(model) == 1: 140 | return model[-1] # return model 141 | else: 142 | print('Ensemble created with %s\n' % weights) 143 | for k in ['names', 'stride']: 144 | setattr(model, k, getattr(model[-1], k)) 145 | return model # return ensemble 146 | -------------------------------------------------------------------------------- /yolov5/models/.ipynb_checkpoints/export-checkpoint.py: -------------------------------------------------------------------------------- 1 | """Exports a YOLOv5 *.pt model to ONNX and TorchScript formats 2 | 3 | Usage: 4 | $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 5 | """ 6 | 7 | import argparse 8 | 9 | import torch 10 | 11 | from utils.google_utils import attempt_download 12 | from utils.general import set_logging 13 | 14 | if __name__ == '__main__': 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') 17 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') 18 | parser.add_argument('--batch-size', type=int, default=1, help='batch size') 19 | opt = parser.parse_args() 20 | opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand 21 | print(opt) 22 | set_logging() 23 | 24 | # Input 25 | img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection 26 | 27 | # Load PyTorch model 28 | attempt_download(opt.weights) 29 | model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float() 30 | model.eval() 31 | model.model[-1].export = False # set Detect() layer export=True 32 | y = model(img) # dry run 33 | 34 | # TorchScript export 35 | try: 36 | print('\nStarting TorchScript export with torch %s...' % torch.__version__) 37 | f = opt.weights.replace('.pt', '.torchscript.pt') # filename 38 | ts = torch.jit.trace(model, img) 39 | ts.save(f) 40 | print('TorchScript export success, saved as %s' % f) 41 | except Exception as e: 42 | print('TorchScript export failure: %s' % e) 43 | 44 | # ONNX export 45 | try: 46 | import onnx 47 | 48 | print('\nStarting ONNX export with onnx %s...' % onnx.__version__) 49 | f = opt.weights.replace('.pt', '.onnx') # filename 50 | model.fuse() # only for ONNX 51 | # torch.onnx.export(model, img, f, verbose=False, opset_version=10, input_names=['images'], 52 | # output_names=['classes', 'boxes'] if y is None else ['output']) 53 | torch.onnx.export(model, img, f, verbose=False, opset_version=10, input_names=['data'], 54 | output_names=['prob']if y is None else ['output']) 55 | 56 | # Checks 57 | onnx_model = onnx.load(f) # load onnx model 58 | onnx.checker.check_model(onnx_model) # check onnx model 59 | print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model 60 | print('ONNX export success, saved as %s' % f) 61 | except Exception as e: 62 | print('ONNX export failure: %s' % e) 63 | 64 | # CoreML export 65 | try: 66 | import coremltools as ct 67 | 68 | print('\nStarting CoreML export with coremltools %s...' % ct.__version__) 69 | # convert model from torchscript and apply pixel scaling as per detect.py 70 | model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) 71 | f = opt.weights.replace('.pt', '.mlmodel') # filename 72 | model.save(f) 73 | print('CoreML export success, saved as %s' % f) 74 | except Exception as e: 75 | print('CoreML export failure: %s' % e) 76 | 77 | # Finish 78 | print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.') 79 | -------------------------------------------------------------------------------- /yolov5/models/.ipynb_checkpoints/yolo-checkpoint.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import logging 4 | from copy import deepcopy 5 | from pathlib import Path 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat 11 | from models.experimental import MixConv2d, CrossConv, C3 12 | from utils.general import check_anchor_order, make_divisible, check_file, set_logging 13 | from utils.torch_utils import ( 14 | time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, select_device) 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | class Detect(nn.Module): 20 | stride = None # strides computed during build 21 | export = False # onnx export 22 | 23 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 24 | super(Detect, self).__init__() 25 | self.nc = nc # number of classes 26 | self.no = nc + 5 # number of outputs per anchor 27 | self.nl = len(anchors) # number of detection layers 28 | self.na = len(anchors[0]) // 2 # number of anchors 29 | self.grid = [torch.zeros(1)] * self.nl # init grid 30 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 31 | self.register_buffer('anchors', a) # shape(nl,na,2) 32 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 33 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 34 | 35 | def forward(self, x): 36 | # x = x.copy() # for profiling 37 | z = [] # inference output 38 | self.training |= self.export 39 | for i in range(self.nl): 40 | x[i] = self.m[i](x[i]) # conv 41 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 42 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 43 | 44 | if not self.training: # inference 45 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 46 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 47 | 48 | y = x[i].sigmoid() 49 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 50 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 51 | z.append(y.view(bs, -1, self.no)) 52 | 53 | return x if self.training else (torch.cat(z, 1), x) 54 | 55 | @staticmethod 56 | def _make_grid(nx=20, ny=20): 57 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 58 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 59 | 60 | 61 | class Model(nn.Module): 62 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes 63 | super(Model, self).__init__() 64 | if isinstance(cfg, dict): 65 | self.yaml = cfg # model dict 66 | else: # is *.yaml 67 | import yaml # for torch hub 68 | self.yaml_file = Path(cfg).name 69 | with open(cfg) as f: 70 | self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict 71 | 72 | # Define model 73 | if nc and nc != self.yaml['nc']: 74 | print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)) 75 | self.yaml['nc'] = nc # override yaml value 76 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out 77 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 78 | 79 | # Build strides, anchors 80 | m = self.model[-1] # Detect() 81 | if isinstance(m, Detect): 82 | s = 128 # 2x min stride 83 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 84 | m.anchors /= m.stride.view(-1, 1, 1) 85 | check_anchor_order(m) 86 | self.stride = m.stride 87 | self._initialize_biases() # only run once 88 | # print('Strides: %s' % m.stride.tolist()) 89 | 90 | # Init weights, biases 91 | initialize_weights(self) 92 | self.info() 93 | print('') 94 | 95 | def forward(self, x, augment=False, profile=False): 96 | if augment: 97 | img_size = x.shape[-2:] # height, width 98 | s = [1, 0.83, 0.67] # scales 99 | f = [None, 3, None] # flips (2-ud, 3-lr) 100 | y = [] # outputs 101 | for si, fi in zip(s, f): 102 | xi = scale_img(x.flip(fi) if fi else x, si) 103 | yi = self.forward_once(xi)[0] # forward 104 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 105 | yi[..., :4] /= si # de-scale 106 | if fi == 2: 107 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 108 | elif fi == 3: 109 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 110 | y.append(yi) 111 | return torch.cat(y, 1), None # augmented inference, train 112 | else: 113 | return self.forward_once(x, profile) # single-scale inference, train 114 | 115 | def forward_once(self, x, profile=False): 116 | y, dt = [], [] # outputs 117 | for m in self.model: 118 | if m.f != -1: # if not from previous layer 119 | x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers 120 | 121 | if profile: 122 | try: 123 | import thop 124 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS 125 | except: 126 | o = 0 127 | t = time_synchronized() 128 | for _ in range(10): 129 | _ = m(x) 130 | dt.append((time_synchronized() - t) * 100) 131 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 132 | 133 | x = m(x) # run 134 | y.append(x if m.i in self.save else None) # save output 135 | 136 | if profile: 137 | print('%.1fms total' % sum(dt)) 138 | return x 139 | 140 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 141 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 142 | m = self.model[-1] # Detect() module 143 | for mi, s in zip(m.m, m.stride): # from 144 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 145 | b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 146 | b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 147 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 148 | 149 | def _print_biases(self): 150 | m = self.model[-1] # Detect() module 151 | for mi in m.m: # from 152 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 153 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 154 | 155 | # def _print_weights(self): 156 | # for m in self.model.modules(): 157 | # if type(m) is Bottleneck: 158 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 159 | 160 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 161 | print('Fusing layers... ') 162 | for m in self.model.modules(): 163 | if type(m) is Conv: 164 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability 165 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 166 | m.bn = None # remove batchnorm 167 | m.forward = m.fuseforward # update forward 168 | self.info() 169 | return self 170 | 171 | def info(self): # print model information 172 | model_info(self) 173 | 174 | 175 | def parse_model(d, ch): # model_dict, input_channels(3) 176 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 177 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 178 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 179 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 180 | 181 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 182 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 183 | m = eval(m) if isinstance(m, str) else m # eval strings 184 | for j, a in enumerate(args): 185 | try: 186 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 187 | except: 188 | pass 189 | 190 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 191 | if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: 192 | c1, c2 = ch[f], args[0] 193 | 194 | # Normal 195 | # if i > 0 and args[0] != no: # channel expansion factor 196 | # ex = 1.75 # exponential (default 2.0) 197 | # e = math.log(c2 / ch[1]) / math.log(2) 198 | # c2 = int(ch[1] * ex ** e) 199 | # if m != Focus: 200 | 201 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 202 | 203 | # Experimental 204 | # if i > 0 and args[0] != no: # channel expansion factor 205 | # ex = 1 + gw # exponential (default 2.0) 206 | # ch1 = 32 # ch[1] 207 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n 208 | # c2 = int(ch1 * ex ** e) 209 | # if m != Focus: 210 | # c2 = make_divisible(c2, 8) if c2 != no else c2 211 | 212 | args = [c1, c2, *args[1:]] 213 | if m in [BottleneckCSP, C3]: 214 | args.insert(2, n) 215 | n = 1 216 | elif m is nn.BatchNorm2d: 217 | args = [ch[f]] 218 | elif m is Concat: 219 | c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) 220 | elif m is Detect: 221 | args.append([ch[x + 1] for x in f]) 222 | if isinstance(args[1], int): # number of anchors 223 | args[1] = [list(range(args[1] * 2))] * len(f) 224 | else: 225 | c2 = ch[f] 226 | 227 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 228 | t = str(m)[8:-2].replace('__main__.', '') # module type 229 | np = sum([x.numel() for x in m_.parameters()]) # number params 230 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 231 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print 232 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 233 | layers.append(m_) 234 | ch.append(c2) 235 | return nn.Sequential(*layers), sorted(save) 236 | 237 | 238 | if __name__ == '__main__': 239 | parser = argparse.ArgumentParser() 240 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 241 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 242 | opt = parser.parse_args() 243 | opt.cfg = check_file(opt.cfg) # check file 244 | set_logging() 245 | device = select_device(opt.device) 246 | 247 | # Create model 248 | model = Model(opt.cfg).to(device) 249 | model.train() 250 | 251 | # Profile 252 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 253 | # y = model(img, profile=True) 254 | 255 | # ONNX export 256 | # model.model[-1].export = True 257 | # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) 258 | 259 | # Tensorboard 260 | # from torch.utils.tensorboard import SummaryWriter 261 | # tb_writer = SummaryWriter() 262 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") 263 | # tb_writer.add_graph(model.model, img) # add model to tensorboard 264 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard 265 | -------------------------------------------------------------------------------- /yolov5/models/.ipynb_checkpoints/yolo_1-checkpoint.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import logging 4 | from copy import deepcopy 5 | from pathlib import Path 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat 11 | from models.experimental import MixConv2d, CrossConv, C3 12 | from utils.general import check_anchor_order, make_divisible, check_file, set_logging 13 | from utils.torch_utils import ( 14 | time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, select_device) 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | class Detect(nn.Module): 20 | stride = None # strides computed during build 21 | export = False # onnx export 22 | 23 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 24 | super(Detect, self).__init__() 25 | self.nc = nc # number of classes 26 | self.no = nc + 5 # number of outputs per anchor 27 | self.nl = len(anchors) # number of detection layers 28 | self.na = len(anchors[0]) // 2 # number of anchors 29 | self.grid = [torch.zeros(1)] * self.nl # init grid 30 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 31 | self.register_buffer('anchors', a) # shape(nl,na,2) 32 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 33 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 34 | 35 | def forward(self, x): 36 | # x = x.copy() # for profiling 37 | z = [] # inference output 38 | self.training |= self.export 39 | for i in range(self.nl): 40 | x[i] = self.m[i](x[i]) # conv 41 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 42 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 43 | 44 | if not self.training: # inference 45 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 46 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 47 | 48 | y = x[i].sigmoid() 49 | # y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 50 | # y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 51 | c=(y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 52 | d=(y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 53 | e=y[..., 4:6] 54 | f=torch.cat((c,d,e),4) 55 | z.append(f.view(bs, -1, self.no)) 56 | 57 | return x if self.training else torch.cat(z, 1) 58 | 59 | @staticmethod 60 | def _make_grid(nx=20, ny=20): 61 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 62 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 63 | 64 | 65 | class Model(nn.Module): 66 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes 67 | super(Model, self).__init__() 68 | if isinstance(cfg, dict): 69 | self.yaml = cfg # model dict 70 | else: # is *.yaml 71 | import yaml # for torch hub 72 | self.yaml_file = Path(cfg).name 73 | with open(cfg) as f: 74 | self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict 75 | 76 | # Define model 77 | if nc and nc != self.yaml['nc']: 78 | print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)) 79 | self.yaml['nc'] = nc # override yaml value 80 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out 81 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 82 | 83 | # Build strides, anchors 84 | m = self.model[-1] # Detect() 85 | if isinstance(m, Detect): 86 | s = 128 # 2x min stride 87 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 88 | m.anchors /= m.stride.view(-1, 1, 1) 89 | check_anchor_order(m) 90 | self.stride = m.stride 91 | self._initialize_biases() # only run once 92 | # print('Strides: %s' % m.stride.tolist()) 93 | 94 | # Init weights, biases 95 | initialize_weights(self) 96 | self.info() 97 | print('') 98 | 99 | def forward(self, x, augment=False, profile=False): 100 | if augment: 101 | img_size = x.shape[-2:] # height, width 102 | s = [1, 0.83, 0.67] # scales 103 | f = [None, 3, None] # flips (2-ud, 3-lr) 104 | y = [] # outputs 105 | for si, fi in zip(s, f): 106 | xi = scale_img(x.flip(fi) if fi else x, si) 107 | yi = self.forward_once(xi)[0] # forward 108 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 109 | yi[..., :4] /= si # de-scale 110 | if fi == 2: 111 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 112 | elif fi == 3: 113 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 114 | y.append(yi) 115 | return torch.cat(y, 1), None # augmented inference, train 116 | else: 117 | return self.forward_once(x, profile) # single-scale inference, train 118 | 119 | def forward_once(self, x, profile=False): 120 | y, dt = [], [] # outputs 121 | for m in self.model: 122 | if m.f != -1: # if not from previous layer 123 | x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers 124 | 125 | if profile: 126 | try: 127 | import thop 128 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS 129 | except: 130 | o = 0 131 | t = time_synchronized() 132 | for _ in range(10): 133 | _ = m(x) 134 | dt.append((time_synchronized() - t) * 100) 135 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 136 | 137 | x = m(x) # run 138 | y.append(x if m.i in self.save else None) # save output 139 | 140 | if profile: 141 | print('%.1fms total' % sum(dt)) 142 | return x 143 | 144 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 145 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 146 | m = self.model[-1] # Detect() module 147 | for mi, s in zip(m.m, m.stride): # from 148 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 149 | b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 150 | b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 151 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 152 | 153 | def _print_biases(self): 154 | m = self.model[-1] # Detect() module 155 | for mi in m.m: # from 156 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 157 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 158 | 159 | # def _print_weights(self): 160 | # for m in self.model.modules(): 161 | # if type(m) is Bottleneck: 162 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 163 | 164 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 165 | print('Fusing layers... ') 166 | for m in self.model.modules(): 167 | if type(m) is Conv: 168 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability 169 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 170 | m.bn = None # remove batchnorm 171 | m.forward = m.fuseforward # update forward 172 | self.info() 173 | return self 174 | 175 | def info(self): # print model information 176 | model_info(self) 177 | 178 | 179 | def parse_model(d, ch): # model_dict, input_channels(3) 180 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 181 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 182 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 183 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 184 | 185 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 186 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 187 | m = eval(m) if isinstance(m, str) else m # eval strings 188 | for j, a in enumerate(args): 189 | try: 190 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 191 | except: 192 | pass 193 | 194 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 195 | if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: 196 | c1, c2 = ch[f], args[0] 197 | 198 | # Normal 199 | # if i > 0 and args[0] != no: # channel expansion factor 200 | # ex = 1.75 # exponential (default 2.0) 201 | # e = math.log(c2 / ch[1]) / math.log(2) 202 | # c2 = int(ch[1] * ex ** e) 203 | # if m != Focus: 204 | 205 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 206 | 207 | # Experimental 208 | # if i > 0 and args[0] != no: # channel expansion factor 209 | # ex = 1 + gw # exponential (default 2.0) 210 | # ch1 = 32 # ch[1] 211 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n 212 | # c2 = int(ch1 * ex ** e) 213 | # if m != Focus: 214 | # c2 = make_divisible(c2, 8) if c2 != no else c2 215 | 216 | args = [c1, c2, *args[1:]] 217 | if m in [BottleneckCSP, C3]: 218 | args.insert(2, n) 219 | n = 1 220 | elif m is nn.BatchNorm2d: 221 | args = [ch[f]] 222 | elif m is Concat: 223 | c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) 224 | elif m is Detect: 225 | args.append([ch[x + 1] for x in f]) 226 | if isinstance(args[1], int): # number of anchors 227 | args[1] = [list(range(args[1] * 2))] * len(f) 228 | else: 229 | c2 = ch[f] 230 | 231 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 232 | t = str(m)[8:-2].replace('__main__.', '') # module type 233 | np = sum([x.numel() for x in m_.parameters()]) # number params 234 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 235 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print 236 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 237 | layers.append(m_) 238 | ch.append(c2) 239 | return nn.Sequential(*layers), sorted(save) 240 | 241 | 242 | if __name__ == '__main__': 243 | parser = argparse.ArgumentParser() 244 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 245 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 246 | opt = parser.parse_args() 247 | opt.cfg = check_file(opt.cfg) # check file 248 | set_logging() 249 | device = select_device(opt.device) 250 | 251 | # Create model 252 | model = Model(opt.cfg).to(device) 253 | model.train() 254 | 255 | # Profile 256 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 257 | # y = model(img, profile=True) 258 | 259 | # ONNX export 260 | # model.model[-1].export = True 261 | # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) 262 | 263 | # Tensorboard 264 | # from torch.utils.tensorboard import SummaryWriter 265 | # tb_writer = SummaryWriter() 266 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") 267 | # tb_writer.add_graph(model.model, img) # add model to tensorboard 268 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard 269 | -------------------------------------------------------------------------------- /yolov5/models/.ipynb_checkpoints/yolov5s-checkpoint.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 2 # number of classes 3 | depth_multiple: 0.33 # model depth multiple 4 | width_multiple: 0.50 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/models/__init__.py -------------------------------------------------------------------------------- /yolov5/models/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/models/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/models/__pycache__/common.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/models/__pycache__/common.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/models/__pycache__/experimental.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/models/__pycache__/experimental.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/models/__pycache__/yolo.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/models/__pycache__/yolo.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/models/common.py: -------------------------------------------------------------------------------- 1 | # This file contains modules common to various models 2 | import math 3 | 4 | import torch 5 | import torch.nn as nn 6 | 7 | 8 | def autopad(k, p=None): # kernel, padding 9 | # Pad to 'same' 10 | if p is None: 11 | p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad 12 | return p 13 | 14 | 15 | def DWConv(c1, c2, k=1, s=1, act=True): 16 | # Depthwise convolution 17 | return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) 18 | 19 | 20 | class Conv(nn.Module): 21 | # Standard convolution 22 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 23 | super(Conv, self).__init__() 24 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) 25 | self.bn = nn.BatchNorm2d(c2) 26 | self.act = nn.LeakyReLU(0.1) if act else nn.Identity() 27 | 28 | def forward(self, x): 29 | return self.act(self.bn(self.conv(x))) 30 | 31 | def fuseforward(self, x): 32 | return self.act(self.conv(x)) 33 | 34 | 35 | class Bottleneck(nn.Module): 36 | # Standard bottleneck 37 | def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion 38 | super(Bottleneck, self).__init__() 39 | c_ = int(c2 * e) # hidden channels 40 | self.cv1 = Conv(c1, c_, 1, 1) 41 | self.cv2 = Conv(c_, c2, 3, 1, g=g) 42 | self.add = shortcut and c1 == c2 43 | 44 | def forward(self, x): 45 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 46 | 47 | 48 | class BottleneckCSP(nn.Module): 49 | # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks 50 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 51 | super(BottleneckCSP, self).__init__() 52 | c_ = int(c2 * e) # hidden channels 53 | self.cv1 = Conv(c1, c_, 1, 1) 54 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 55 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 56 | self.cv4 = Conv(2 * c_, c2, 1, 1) 57 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 58 | self.act = nn.LeakyReLU(0.1, inplace=True) 59 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) 60 | 61 | def forward(self, x): 62 | y1 = self.cv3(self.m(self.cv1(x))) 63 | y2 = self.cv2(x) 64 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 65 | 66 | 67 | class SPP(nn.Module): 68 | # Spatial pyramid pooling layer used in YOLOv3-SPP 69 | def __init__(self, c1, c2, k=(5, 9, 13)): 70 | super(SPP, self).__init__() 71 | c_ = c1 // 2 # hidden channels 72 | self.cv1 = Conv(c1, c_, 1, 1) 73 | self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) 74 | self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) 75 | 76 | def forward(self, x): 77 | x = self.cv1(x) 78 | return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) 79 | 80 | 81 | class Focus(nn.Module): 82 | # Focus wh information into c-space 83 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups 84 | super(Focus, self).__init__() 85 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act) 86 | 87 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) 88 | return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) 89 | 90 | 91 | class Concat(nn.Module): 92 | # Concatenate a list of tensors along dimension 93 | def __init__(self, dimension=1): 94 | super(Concat, self).__init__() 95 | self.d = dimension 96 | 97 | def forward(self, x): 98 | return torch.cat(x, self.d) 99 | 100 | 101 | class Flatten(nn.Module): 102 | # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions 103 | @staticmethod 104 | def forward(x): 105 | return x.view(x.size(0), -1) 106 | 107 | 108 | class Classify(nn.Module): 109 | # Classification head, i.e. x(b,c1,20,20) to x(b,c2) 110 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups 111 | super(Classify, self).__init__() 112 | self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1) 113 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) # to x(b,c2,1,1) 114 | self.flat = Flatten() 115 | 116 | def forward(self, x): 117 | z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list 118 | return self.flat(self.conv(z)) # flatten to x(b,c2) 119 | -------------------------------------------------------------------------------- /yolov5/models/experimental.py: -------------------------------------------------------------------------------- 1 | # This file contains experimental modules 2 | 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | 7 | from models.common import Conv, DWConv 8 | from utils.google_utils import attempt_download 9 | 10 | 11 | class CrossConv(nn.Module): 12 | # Cross Convolution Downsample 13 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False): 14 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut 15 | super(CrossConv, self).__init__() 16 | c_ = int(c2 * e) # hidden channels 17 | self.cv1 = Conv(c1, c_, (1, k), (1, s)) 18 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g) 19 | self.add = shortcut and c1 == c2 20 | 21 | def forward(self, x): 22 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) 23 | 24 | 25 | class C3(nn.Module): 26 | # Cross Convolution CSP 27 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion 28 | super(C3, self).__init__() 29 | c_ = int(c2 * e) # hidden channels 30 | self.cv1 = Conv(c1, c_, 1, 1) 31 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) 32 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) 33 | self.cv4 = Conv(2 * c_, c2, 1, 1) 34 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) 35 | self.act = nn.LeakyReLU(0.1, inplace=True) 36 | self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)]) 37 | 38 | def forward(self, x): 39 | y1 = self.cv3(self.m(self.cv1(x))) 40 | y2 = self.cv2(x) 41 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) 42 | 43 | 44 | class Sum(nn.Module): 45 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 46 | def __init__(self, n, weight=False): # n: number of inputs 47 | super(Sum, self).__init__() 48 | self.weight = weight # apply weights boolean 49 | self.iter = range(n - 1) # iter object 50 | if weight: 51 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights 52 | 53 | def forward(self, x): 54 | y = x[0] # no weight 55 | if self.weight: 56 | w = torch.sigmoid(self.w) * 2 57 | for i in self.iter: 58 | y = y + x[i + 1] * w[i] 59 | else: 60 | for i in self.iter: 61 | y = y + x[i + 1] 62 | return y 63 | 64 | 65 | class GhostConv(nn.Module): 66 | # Ghost Convolution https://github.com/huawei-noah/ghostnet 67 | def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups 68 | super(GhostConv, self).__init__() 69 | c_ = c2 // 2 # hidden channels 70 | self.cv1 = Conv(c1, c_, k, s, g, act) 71 | self.cv2 = Conv(c_, c_, 5, 1, c_, act) 72 | 73 | def forward(self, x): 74 | y = self.cv1(x) 75 | return torch.cat([y, self.cv2(y)], 1) 76 | 77 | 78 | class GhostBottleneck(nn.Module): 79 | # Ghost Bottleneck https://github.com/huawei-noah/ghostnet 80 | def __init__(self, c1, c2, k, s): 81 | super(GhostBottleneck, self).__init__() 82 | c_ = c2 // 2 83 | self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw 84 | DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw 85 | GhostConv(c_, c2, 1, 1, act=False)) # pw-linear 86 | self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), 87 | Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() 88 | 89 | def forward(self, x): 90 | return self.conv(x) + self.shortcut(x) 91 | 92 | 93 | class MixConv2d(nn.Module): 94 | # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595 95 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): 96 | super(MixConv2d, self).__init__() 97 | groups = len(k) 98 | if equal_ch: # equal c_ per group 99 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices 100 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels 101 | else: # equal weight.numel() per group 102 | b = [c2] + [0] * groups 103 | a = np.eye(groups + 1, groups, k=-1) 104 | a -= np.roll(a, 1, axis=1) 105 | a *= np.array(k) ** 2 106 | a[0] = 1 107 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b 108 | 109 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)]) 110 | self.bn = nn.BatchNorm2d(c2) 111 | self.act = nn.LeakyReLU(0.1, inplace=True) 112 | 113 | def forward(self, x): 114 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1))) 115 | 116 | 117 | class Ensemble(nn.ModuleList): 118 | # Ensemble of models 119 | def __init__(self): 120 | super(Ensemble, self).__init__() 121 | 122 | def forward(self, x, augment=False): 123 | y = [] 124 | for module in self: 125 | y.append(module(x, augment)[0]) 126 | # y = torch.stack(y).max(0)[0] # max ensemble 127 | # y = torch.cat(y, 1) # nms ensemble 128 | y = torch.stack(y).mean(0) # mean ensemble 129 | return y, None # inference, train output 130 | 131 | 132 | def attempt_load(weights, map_location=None): 133 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a 134 | model = Ensemble() 135 | for w in weights if isinstance(weights, list) else [weights]: 136 | attempt_download(w) 137 | model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model 138 | 139 | if len(model) == 1: 140 | return model[-1] # return model 141 | else: 142 | print('Ensemble created with %s\n' % weights) 143 | for k in ['names', 'stride']: 144 | setattr(model, k, getattr(model[-1], k)) 145 | return model # return ensemble 146 | -------------------------------------------------------------------------------- /yolov5/models/export.py: -------------------------------------------------------------------------------- 1 | """Exports a YOLOv5 *.pt model to ONNX and TorchScript formats 2 | 3 | Usage: 4 | $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 5 | """ 6 | 7 | import argparse 8 | 9 | import torch 10 | 11 | from utils.google_utils import attempt_download 12 | from utils.general import set_logging 13 | 14 | if __name__ == '__main__': 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') 17 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') 18 | parser.add_argument('--batch-size', type=int, default=1, help='batch size') 19 | opt = parser.parse_args() 20 | opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand 21 | print(opt) 22 | set_logging() 23 | 24 | # Input 25 | img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection 26 | 27 | # Load PyTorch model 28 | attempt_download(opt.weights) 29 | model = torch.load(opt.weights, map_location=torch.device('cpu'))['model'].float() 30 | model.eval() 31 | model.model[-1].export = False # set Detect() layer export=True 32 | y = model(img) # dry run 33 | 34 | # TorchScript export 35 | try: 36 | print('\nStarting TorchScript export with torch %s...' % torch.__version__) 37 | f = opt.weights.replace('.pt', '.torchscript.pt') # filename 38 | ts = torch.jit.trace(model, img) 39 | ts.save(f) 40 | print('TorchScript export success, saved as %s' % f) 41 | except Exception as e: 42 | print('TorchScript export failure: %s' % e) 43 | 44 | # ONNX export 45 | try: 46 | import onnx 47 | 48 | print('\nStarting ONNX export with onnx %s...' % onnx.__version__) 49 | f = opt.weights.replace('.pt', '.onnx') # filename 50 | model.fuse() # only for ONNX 51 | # torch.onnx.export(model, img, f, verbose=False, opset_version=10, input_names=['images'], 52 | # output_names=['classes', 'boxes'] if y is None else ['output']) 53 | torch.onnx.export(model, img, f, verbose=False, opset_version=10, input_names=['data'], 54 | output_names=['prob']if y is None else ['output']) 55 | 56 | # Checks 57 | onnx_model = onnx.load(f) # load onnx model 58 | onnx.checker.check_model(onnx_model) # check onnx model 59 | print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model 60 | print('ONNX export success, saved as %s' % f) 61 | except Exception as e: 62 | print('ONNX export failure: %s' % e) 63 | 64 | # CoreML export 65 | try: 66 | import coremltools as ct 67 | 68 | print('\nStarting CoreML export with coremltools %s...' % ct.__version__) 69 | # convert model from torchscript and apply pixel scaling as per detect.py 70 | model = ct.convert(ts, inputs=[ct.ImageType(name='images', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) 71 | f = opt.weights.replace('.pt', '.mlmodel') # filename 72 | model.save(f) 73 | print('CoreML export success, saved as %s' % f) 74 | except Exception as e: 75 | print('CoreML export failure: %s' % e) 76 | 77 | # Finish 78 | print('\nExport complete. Visualize with https://github.com/lutzroeder/netron.') 79 | -------------------------------------------------------------------------------- /yolov5/models/hub/yolov3-spp.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # darknet53 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Conv, [32, 3, 1]], # 0 16 | [-1, 1, Conv, [64, 3, 2]], # 1-P1/2 17 | [-1, 1, Bottleneck, [64]], 18 | [-1, 1, Conv, [128, 3, 2]], # 3-P2/4 19 | [-1, 2, Bottleneck, [128]], 20 | [-1, 1, Conv, [256, 3, 2]], # 5-P3/8 21 | [-1, 8, Bottleneck, [256]], 22 | [-1, 1, Conv, [512, 3, 2]], # 7-P4/16 23 | [-1, 8, Bottleneck, [512]], 24 | [-1, 1, Conv, [1024, 3, 2]], # 9-P5/32 25 | [-1, 4, Bottleneck, [1024]], # 10 26 | ] 27 | 28 | # YOLOv3-SPP head 29 | head: 30 | [[-1, 1, Bottleneck, [1024, False]], 31 | [-1, 1, SPP, [512, [5, 9, 13]]], 32 | [-1, 1, Conv, [1024, 3, 1]], 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 1, Conv, [1024, 3, 1]], # 15 (P5/32-large) 35 | 36 | [-2, 1, Conv, [256, 1, 1]], 37 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 38 | [[-1, 8], 1, Concat, [1]], # cat backbone P4 39 | [-1, 1, Bottleneck, [512, False]], 40 | [-1, 1, Bottleneck, [512, False]], 41 | [-1, 1, Conv, [256, 1, 1]], 42 | [-1, 1, Conv, [512, 3, 1]], # 22 (P4/16-medium) 43 | 44 | [-2, 1, Conv, [128, 1, 1]], 45 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 46 | [[-1, 6], 1, Concat, [1]], # cat backbone P3 47 | [-1, 1, Bottleneck, [256, False]], 48 | [-1, 2, Bottleneck, [256, False]], # 27 (P3/8-small) 49 | 50 | [[27, 22, 15], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 51 | ] 52 | -------------------------------------------------------------------------------- /yolov5/models/hub/yolov5-fpn.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, Bottleneck, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 6, BottleneckCSP, [1024]], # 9 25 | ] 26 | 27 | # YOLOv5 FPN head 28 | head: 29 | [[-1, 3, BottleneckCSP, [1024, False]], # 10 (P5/32-large) 30 | 31 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 32 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 33 | [-1, 1, Conv, [512, 1, 1]], 34 | [-1, 3, BottleneckCSP, [512, False]], # 14 (P4/16-medium) 35 | 36 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 37 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 38 | [-1, 1, Conv, [256, 1, 1]], 39 | [-1, 3, BottleneckCSP, [256, False]], # 18 (P3/8-small) 40 | 41 | [[18, 14, 10], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 42 | ] 43 | -------------------------------------------------------------------------------- /yolov5/models/hub/yolov5-panet.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [116,90, 156,198, 373,326] # P5/32 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [10,13, 16,30, 33,23] # P3/8 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 PANet head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P5, P4, P3) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5/models/yolo.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import logging 4 | from copy import deepcopy 5 | from pathlib import Path 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat 11 | from models.experimental import MixConv2d, CrossConv, C3 12 | from utils.general import check_anchor_order, make_divisible, check_file, set_logging 13 | from utils.torch_utils import ( 14 | time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, select_device) 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | class Detect(nn.Module): 20 | stride = None # strides computed during build 21 | export = False # onnx export 22 | 23 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 24 | super(Detect, self).__init__() 25 | self.nc = nc # number of classes 26 | self.no = nc + 5 # number of outputs per anchor 27 | self.nl = len(anchors) # number of detection layers 28 | self.na = len(anchors[0]) // 2 # number of anchors 29 | self.grid = [torch.zeros(1)] * self.nl # init grid 30 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 31 | self.register_buffer('anchors', a) # shape(nl,na,2) 32 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 33 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 34 | 35 | def forward(self, x): 36 | # x = x.copy() # for profiling 37 | z = [] # inference output 38 | self.training |= self.export 39 | for i in range(self.nl): 40 | x[i] = self.m[i](x[i]) # conv 41 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 42 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 43 | 44 | if not self.training: # inference 45 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 46 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 47 | 48 | y = x[i].sigmoid() 49 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 50 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 51 | z.append(y.view(bs, -1, self.no)) 52 | 53 | return x if self.training else (torch.cat(z, 1), x) 54 | 55 | @staticmethod 56 | def _make_grid(nx=20, ny=20): 57 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 58 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 59 | 60 | 61 | class Model(nn.Module): 62 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes 63 | super(Model, self).__init__() 64 | if isinstance(cfg, dict): 65 | self.yaml = cfg # model dict 66 | else: # is *.yaml 67 | import yaml # for torch hub 68 | self.yaml_file = Path(cfg).name 69 | with open(cfg) as f: 70 | self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict 71 | 72 | # Define model 73 | if nc and nc != self.yaml['nc']: 74 | print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)) 75 | self.yaml['nc'] = nc # override yaml value 76 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out 77 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 78 | 79 | # Build strides, anchors 80 | m = self.model[-1] # Detect() 81 | if isinstance(m, Detect): 82 | s = 128 # 2x min stride 83 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 84 | m.anchors /= m.stride.view(-1, 1, 1) 85 | check_anchor_order(m) 86 | self.stride = m.stride 87 | self._initialize_biases() # only run once 88 | # print('Strides: %s' % m.stride.tolist()) 89 | 90 | # Init weights, biases 91 | initialize_weights(self) 92 | self.info() 93 | print('') 94 | 95 | def forward(self, x, augment=False, profile=False): 96 | if augment: 97 | img_size = x.shape[-2:] # height, width 98 | s = [1, 0.83, 0.67] # scales 99 | f = [None, 3, None] # flips (2-ud, 3-lr) 100 | y = [] # outputs 101 | for si, fi in zip(s, f): 102 | xi = scale_img(x.flip(fi) if fi else x, si) 103 | yi = self.forward_once(xi)[0] # forward 104 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 105 | yi[..., :4] /= si # de-scale 106 | if fi == 2: 107 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 108 | elif fi == 3: 109 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 110 | y.append(yi) 111 | return torch.cat(y, 1), None # augmented inference, train 112 | else: 113 | return self.forward_once(x, profile) # single-scale inference, train 114 | 115 | def forward_once(self, x, profile=False): 116 | y, dt = [], [] # outputs 117 | for m in self.model: 118 | if m.f != -1: # if not from previous layer 119 | x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers 120 | 121 | if profile: 122 | try: 123 | import thop 124 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS 125 | except: 126 | o = 0 127 | t = time_synchronized() 128 | for _ in range(10): 129 | _ = m(x) 130 | dt.append((time_synchronized() - t) * 100) 131 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 132 | 133 | x = m(x) # run 134 | y.append(x if m.i in self.save else None) # save output 135 | 136 | if profile: 137 | print('%.1fms total' % sum(dt)) 138 | return x 139 | 140 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 141 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 142 | m = self.model[-1] # Detect() module 143 | for mi, s in zip(m.m, m.stride): # from 144 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 145 | b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 146 | b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 147 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 148 | 149 | def _print_biases(self): 150 | m = self.model[-1] # Detect() module 151 | for mi in m.m: # from 152 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 153 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 154 | 155 | # def _print_weights(self): 156 | # for m in self.model.modules(): 157 | # if type(m) is Bottleneck: 158 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 159 | 160 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 161 | print('Fusing layers... ') 162 | for m in self.model.modules(): 163 | if type(m) is Conv: 164 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability 165 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 166 | m.bn = None # remove batchnorm 167 | m.forward = m.fuseforward # update forward 168 | self.info() 169 | return self 170 | 171 | def info(self): # print model information 172 | model_info(self) 173 | 174 | 175 | def parse_model(d, ch): # model_dict, input_channels(3) 176 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 177 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 178 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 179 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 180 | 181 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 182 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 183 | m = eval(m) if isinstance(m, str) else m # eval strings 184 | for j, a in enumerate(args): 185 | try: 186 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 187 | except: 188 | pass 189 | 190 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 191 | if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: 192 | c1, c2 = ch[f], args[0] 193 | 194 | # Normal 195 | # if i > 0 and args[0] != no: # channel expansion factor 196 | # ex = 1.75 # exponential (default 2.0) 197 | # e = math.log(c2 / ch[1]) / math.log(2) 198 | # c2 = int(ch[1] * ex ** e) 199 | # if m != Focus: 200 | 201 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 202 | 203 | # Experimental 204 | # if i > 0 and args[0] != no: # channel expansion factor 205 | # ex = 1 + gw # exponential (default 2.0) 206 | # ch1 = 32 # ch[1] 207 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n 208 | # c2 = int(ch1 * ex ** e) 209 | # if m != Focus: 210 | # c2 = make_divisible(c2, 8) if c2 != no else c2 211 | 212 | args = [c1, c2, *args[1:]] 213 | if m in [BottleneckCSP, C3]: 214 | args.insert(2, n) 215 | n = 1 216 | elif m is nn.BatchNorm2d: 217 | args = [ch[f]] 218 | elif m is Concat: 219 | c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) 220 | elif m is Detect: 221 | args.append([ch[x + 1] for x in f]) 222 | if isinstance(args[1], int): # number of anchors 223 | args[1] = [list(range(args[1] * 2))] * len(f) 224 | else: 225 | c2 = ch[f] 226 | 227 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 228 | t = str(m)[8:-2].replace('__main__.', '') # module type 229 | np = sum([x.numel() for x in m_.parameters()]) # number params 230 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 231 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print 232 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 233 | layers.append(m_) 234 | ch.append(c2) 235 | return nn.Sequential(*layers), sorted(save) 236 | 237 | 238 | if __name__ == '__main__': 239 | parser = argparse.ArgumentParser() 240 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 241 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 242 | opt = parser.parse_args() 243 | opt.cfg = check_file(opt.cfg) # check file 244 | set_logging() 245 | device = select_device(opt.device) 246 | 247 | # Create model 248 | model = Model(opt.cfg).to(device) 249 | model.train() 250 | 251 | # Profile 252 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 253 | # y = model(img, profile=True) 254 | 255 | # ONNX export 256 | # model.model[-1].export = True 257 | # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) 258 | 259 | # Tensorboard 260 | # from torch.utils.tensorboard import SummaryWriter 261 | # tb_writer = SummaryWriter() 262 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") 263 | # tb_writer.add_graph(model.model, img) # add model to tensorboard 264 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard 265 | -------------------------------------------------------------------------------- /yolov5/models/yolo_1.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import math 3 | import logging 4 | from copy import deepcopy 5 | from pathlib import Path 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat 11 | from models.experimental import MixConv2d, CrossConv, C3 12 | from utils.general import check_anchor_order, make_divisible, check_file, set_logging 13 | from utils.torch_utils import ( 14 | time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, select_device) 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | class Detect(nn.Module): 20 | stride = None # strides computed during build 21 | export = False # onnx export 22 | 23 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer 24 | super(Detect, self).__init__() 25 | self.nc = nc # number of classes 26 | self.no = nc + 5 # number of outputs per anchor 27 | self.nl = len(anchors) # number of detection layers 28 | self.na = len(anchors[0]) // 2 # number of anchors 29 | self.grid = [torch.zeros(1)] * self.nl # init grid 30 | a = torch.tensor(anchors).float().view(self.nl, -1, 2) 31 | self.register_buffer('anchors', a) # shape(nl,na,2) 32 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2) 33 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv 34 | 35 | def forward(self, x): 36 | # x = x.copy() # for profiling 37 | z = [] # inference output 38 | self.training |= self.export 39 | for i in range(self.nl): 40 | x[i] = self.m[i](x[i]) # conv 41 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 42 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() 43 | 44 | if not self.training: # inference 45 | if self.grid[i].shape[2:4] != x[i].shape[2:4]: 46 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device) 47 | 48 | y = x[i].sigmoid() 49 | # y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 50 | # y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 51 | c=(y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy 52 | d=(y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh 53 | # e=y[..., 4:6] 54 | e = y[..., 4:] 55 | f=torch.cat((c,d,e),4) 56 | z.append(f.view(bs, -1, self.no)) 57 | 58 | return x if self.training else torch.cat(z, 1) 59 | 60 | @staticmethod 61 | def _make_grid(nx=20, ny=20): 62 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)]) 63 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float() 64 | 65 | 66 | class Model(nn.Module): 67 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None): # model, input channels, number of classes 68 | super(Model, self).__init__() 69 | if isinstance(cfg, dict): 70 | self.yaml = cfg # model dict 71 | else: # is *.yaml 72 | import yaml # for torch hub 73 | self.yaml_file = Path(cfg).name 74 | with open(cfg) as f: 75 | self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict 76 | 77 | # Define model 78 | if nc and nc != self.yaml['nc']: 79 | print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)) 80 | self.yaml['nc'] = nc # override yaml value 81 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist, ch_out 82 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) 83 | 84 | # Build strides, anchors 85 | m = self.model[-1] # Detect() 86 | if isinstance(m, Detect): 87 | s = 128 # 2x min stride 88 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward 89 | m.anchors /= m.stride.view(-1, 1, 1) 90 | check_anchor_order(m) 91 | self.stride = m.stride 92 | self._initialize_biases() # only run once 93 | # print('Strides: %s' % m.stride.tolist()) 94 | 95 | # Init weights, biases 96 | initialize_weights(self) 97 | self.info() 98 | print('') 99 | 100 | def forward(self, x, augment=False, profile=False): 101 | if augment: 102 | img_size = x.shape[-2:] # height, width 103 | s = [1, 0.83, 0.67] # scales 104 | f = [None, 3, None] # flips (2-ud, 3-lr) 105 | y = [] # outputs 106 | for si, fi in zip(s, f): 107 | xi = scale_img(x.flip(fi) if fi else x, si) 108 | yi = self.forward_once(xi)[0] # forward 109 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 110 | yi[..., :4] /= si # de-scale 111 | if fi == 2: 112 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud 113 | elif fi == 3: 114 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr 115 | y.append(yi) 116 | return torch.cat(y, 1), None # augmented inference, train 117 | else: 118 | return self.forward_once(x, profile) # single-scale inference, train 119 | 120 | def forward_once(self, x, profile=False): 121 | y, dt = [], [] # outputs 122 | for m in self.model: 123 | if m.f != -1: # if not from previous layer 124 | x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers 125 | 126 | if profile: 127 | try: 128 | import thop 129 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS 130 | except: 131 | o = 0 132 | t = time_synchronized() 133 | for _ in range(10): 134 | _ = m(x) 135 | dt.append((time_synchronized() - t) * 100) 136 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) 137 | 138 | x = m(x) # run 139 | y.append(x if m.i in self.save else None) # save output 140 | 141 | if profile: 142 | print('%.1fms total' % sum(dt)) 143 | return x 144 | 145 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency 146 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 147 | m = self.model[-1] # Detect() module 148 | for mi, s in zip(m.m, m.stride): # from 149 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) 150 | b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) 151 | b[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls 152 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) 153 | 154 | def _print_biases(self): 155 | m = self.model[-1] # Detect() module 156 | for mi in m.m: # from 157 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) 158 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) 159 | 160 | # def _print_weights(self): 161 | # for m in self.model.modules(): 162 | # if type(m) is Bottleneck: 163 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights 164 | 165 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers 166 | print('Fusing layers... ') 167 | for m in self.model.modules(): 168 | if type(m) is Conv: 169 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatability 170 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 171 | m.bn = None # remove batchnorm 172 | m.forward = m.fuseforward # update forward 173 | self.info() 174 | return self 175 | 176 | def info(self): # print model information 177 | model_info(self) 178 | 179 | 180 | def parse_model(d, ch): # model_dict, input_channels(3) 181 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments')) 182 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 183 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors 184 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5) 185 | 186 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out 187 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args 188 | m = eval(m) if isinstance(m, str) else m # eval strings 189 | for j, a in enumerate(args): 190 | try: 191 | args[j] = eval(a) if isinstance(a, str) else a # eval strings 192 | except: 193 | pass 194 | 195 | n = max(round(n * gd), 1) if n > 1 else n # depth gain 196 | if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3]: 197 | c1, c2 = ch[f], args[0] 198 | 199 | # Normal 200 | # if i > 0 and args[0] != no: # channel expansion factor 201 | # ex = 1.75 # exponential (default 2.0) 202 | # e = math.log(c2 / ch[1]) / math.log(2) 203 | # c2 = int(ch[1] * ex ** e) 204 | # if m != Focus: 205 | 206 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 207 | 208 | # Experimental 209 | # if i > 0 and args[0] != no: # channel expansion factor 210 | # ex = 1 + gw # exponential (default 2.0) 211 | # ch1 = 32 # ch[1] 212 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n 213 | # c2 = int(ch1 * ex ** e) 214 | # if m != Focus: 215 | # c2 = make_divisible(c2, 8) if c2 != no else c2 216 | 217 | args = [c1, c2, *args[1:]] 218 | if m in [BottleneckCSP, C3]: 219 | args.insert(2, n) 220 | n = 1 221 | elif m is nn.BatchNorm2d: 222 | args = [ch[f]] 223 | elif m is Concat: 224 | c2 = sum([ch[-1 if x == -1 else x + 1] for x in f]) 225 | elif m is Detect: 226 | args.append([ch[x + 1] for x in f]) 227 | if isinstance(args[1], int): # number of anchors 228 | args[1] = [list(range(args[1] * 2))] * len(f) 229 | else: 230 | c2 = ch[f] 231 | 232 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module 233 | t = str(m)[8:-2].replace('__main__.', '') # module type 234 | np = sum([x.numel() for x in m_.parameters()]) # number params 235 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 236 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print 237 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist 238 | layers.append(m_) 239 | ch.append(c2) 240 | return nn.Sequential(*layers), sorted(save) 241 | 242 | 243 | if __name__ == '__main__': 244 | parser = argparse.ArgumentParser() 245 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') 246 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 247 | opt = parser.parse_args() 248 | opt.cfg = check_file(opt.cfg) # check file 249 | set_logging() 250 | device = select_device(opt.device) 251 | 252 | # Create model 253 | model = Model(opt.cfg).to(device) 254 | model.train() 255 | 256 | # Profile 257 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) 258 | # y = model(img, profile=True) 259 | 260 | # ONNX export 261 | # model.model[-1].export = True 262 | # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) 263 | 264 | # Tensorboard 265 | # from torch.utils.tensorboard import SummaryWriter 266 | # tb_writer = SummaryWriter() 267 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/") 268 | # tb_writer.add_graph(model.model, img) # add model to tensorboard 269 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard 270 | -------------------------------------------------------------------------------- /yolov5/models/yolov5l.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.0 # model depth multiple 4 | width_multiple: 1.0 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5/models/yolov5m.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 0.67 # model depth multiple 4 | width_multiple: 0.75 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5/models/yolov5s.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 2 # number of classes 3 | depth_multiple: 0.33 # model depth multiple 4 | width_multiple: 0.50 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5/models/yolov5x.yaml: -------------------------------------------------------------------------------- 1 | # parameters 2 | nc: 80 # number of classes 3 | depth_multiple: 1.33 # model depth multiple 4 | width_multiple: 1.25 # layer channel multiple 5 | 6 | # anchors 7 | anchors: 8 | - [10,13, 16,30, 33,23] # P3/8 9 | - [30,61, 62,45, 59,119] # P4/16 10 | - [116,90, 156,198, 373,326] # P5/32 11 | 12 | # YOLOv5 backbone 13 | backbone: 14 | # [from, number, module, args] 15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2 16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 17 | [-1, 3, BottleneckCSP, [128]], 18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 19 | [-1, 9, BottleneckCSP, [256]], 20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 21 | [-1, 9, BottleneckCSP, [512]], 22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 23 | [-1, 1, SPP, [1024, [5, 9, 13]]], 24 | [-1, 3, BottleneckCSP, [1024, False]], # 9 25 | ] 26 | 27 | # YOLOv5 head 28 | head: 29 | [[-1, 1, Conv, [512, 1, 1]], 30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4 32 | [-1, 3, BottleneckCSP, [512, False]], # 13 33 | 34 | [-1, 1, Conv, [256, 1, 1]], 35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']], 36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3 37 | [-1, 3, BottleneckCSP, [256, False]], # 17 (P3/8-small) 38 | 39 | [-1, 1, Conv, [256, 3, 2]], 40 | [[-1, 14], 1, Concat, [1]], # cat head P4 41 | [-1, 3, BottleneckCSP, [512, False]], # 20 (P4/16-medium) 42 | 43 | [-1, 1, Conv, [512, 3, 2]], 44 | [[-1, 10], 1, Concat, [1]], # cat head P5 45 | [-1, 3, BottleneckCSP, [1024, False]], # 23 (P5/32-large) 46 | 47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5) 48 | ] 49 | -------------------------------------------------------------------------------- /yolov5/requirements.txt: -------------------------------------------------------------------------------- 1 | # pip install -r requirements.txt 2 | 3 | # base ---------------------------------------- 4 | Cython 5 | matplotlib>=3.2.2 6 | numpy>=1.18.5 7 | opencv-python>=4.1.2 8 | pillow 9 | PyYAML>=5.3 10 | scipy>=1.4.1 11 | tensorboard>=2.2 12 | torch>=1.6.0 13 | torchvision>=0.7.0 14 | tqdm>=4.41.0 15 | 16 | # coco ---------------------------------------- 17 | # pycocotools>=2.0 18 | 19 | # export -------------------------------------- 20 | # coremltools==4.0b2 21 | # onnx>=1.7.0 22 | 23 | # extras -------------------------------------- 24 | # thop 25 | -------------------------------------------------------------------------------- /yolov5/test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import glob 3 | import json 4 | import os 5 | import shutil 6 | from pathlib import Path 7 | 8 | import numpy as np 9 | import torch 10 | import yaml 11 | from tqdm import tqdm 12 | 13 | from models.experimental import attempt_load 14 | from utils.datasets import create_dataloader 15 | from utils.general import ( 16 | coco80_to_coco91_class, check_dataset, check_file, check_img_size, compute_loss, non_max_suppression, scale_coords, 17 | xyxy2xywh, clip_coords, plot_images, xywh2xyxy, box_iou, output_to_target, ap_per_class, set_logging) 18 | from utils.torch_utils import select_device, time_synchronized 19 | 20 | 21 | def test(data, 22 | weights=None, 23 | batch_size=16, 24 | imgsz=640, 25 | conf_thres=0.001, 26 | iou_thres=0.6, # for NMS 27 | save_json=False, 28 | single_cls=False, 29 | augment=False, 30 | verbose=False, 31 | model=None, 32 | dataloader=None, 33 | save_dir='', 34 | merge=False, 35 | save_txt=False): 36 | # Initialize/load model and set device 37 | training = model is not None 38 | if training: # called by train.py 39 | device = next(model.parameters()).device # get model device 40 | 41 | else: # called directly 42 | set_logging() 43 | device = select_device(opt.device, batch_size=batch_size) 44 | merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels 45 | if save_txt: 46 | out = Path('inference/output') 47 | if os.path.exists(out): 48 | shutil.rmtree(out) # delete output folder 49 | os.makedirs(out) # make new output folder 50 | 51 | # Remove previous 52 | for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')): 53 | os.remove(f) 54 | 55 | # Load model 56 | model = attempt_load(weights, map_location=device) # load FP32 model 57 | imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size 58 | 59 | # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99 60 | # if device.type != 'cpu' and torch.cuda.device_count() > 1: 61 | # model = nn.DataParallel(model) 62 | 63 | # Half 64 | half = device.type != 'cpu' # half precision only supported on CUDA 65 | if half: 66 | model.half() 67 | 68 | # Configure 69 | model.eval() 70 | with open(data) as f: 71 | data = yaml.load(f, Loader=yaml.FullLoader) # model dict 72 | check_dataset(data) # check 73 | nc = 1 if single_cls else int(data['nc']) # number of classes 74 | iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 75 | niou = iouv.numel() 76 | 77 | # Dataloader 78 | if not training: 79 | img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img 80 | _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once 81 | path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images 82 | dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, 83 | hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0] 84 | 85 | seen = 0 86 | names = model.names if hasattr(model, 'names') else model.module.names 87 | coco91class = coco80_to_coco91_class() 88 | s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95') 89 | p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0. 90 | loss = torch.zeros(3, device=device) 91 | jdict, stats, ap, ap_class = [], [], [], [] 92 | for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)): 93 | img = img.to(device, non_blocking=True) 94 | img = img.half() if half else img.float() # uint8 to fp16/32 95 | img /= 255.0 # 0 - 255 to 0.0 - 1.0 96 | targets = targets.to(device) 97 | nb, _, height, width = img.shape # batch size, channels, height, width 98 | whwh = torch.Tensor([width, height, width, height]).to(device) 99 | 100 | # Disable gradients 101 | with torch.no_grad(): 102 | # Run model 103 | t = time_synchronized() 104 | inf_out, train_out = model(img, augment=augment) # inference and training outputs 105 | t0 += time_synchronized() - t 106 | 107 | # Compute loss 108 | if training: # if model has loss hyperparameters 109 | loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls 110 | 111 | # Run NMS 112 | t = time_synchronized() 113 | output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, merge=merge) 114 | t1 += time_synchronized() - t 115 | 116 | # Statistics per image 117 | for si, pred in enumerate(output): 118 | labels = targets[targets[:, 0] == si, 1:] 119 | nl = len(labels) 120 | tcls = labels[:, 0].tolist() if nl else [] # target class 121 | seen += 1 122 | 123 | if pred is None: 124 | if nl: 125 | stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls)) 126 | continue 127 | 128 | # Append to text file 129 | if save_txt: 130 | gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh 131 | x = pred.clone() 132 | x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4], shapes[si][0], shapes[si][1]) # to original 133 | for *xyxy, conf, cls in x: 134 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh 135 | with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f: 136 | f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format 137 | 138 | # Clip boxes to image bounds 139 | clip_coords(pred, (height, width)) 140 | 141 | # Append to pycocotools JSON dictionary 142 | if save_json: 143 | # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... 144 | image_id = Path(paths[si]).stem 145 | box = pred[:, :4].clone() # xyxy 146 | scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape 147 | box = xyxy2xywh(box) # xywh 148 | box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner 149 | for p, b in zip(pred.tolist(), box.tolist()): 150 | jdict.append({'image_id': int(image_id) if image_id.isnumeric() else image_id, 151 | 'category_id': coco91class[int(p[5])], 152 | 'bbox': [round(x, 3) for x in b], 153 | 'score': round(p[4], 5)}) 154 | 155 | # Assign all predictions as incorrect 156 | correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device) 157 | if nl: 158 | detected = [] # target indices 159 | tcls_tensor = labels[:, 0] 160 | 161 | # target boxes 162 | tbox = xywh2xyxy(labels[:, 1:5]) * whwh 163 | 164 | # Per target class 165 | for cls in torch.unique(tcls_tensor): 166 | ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices 167 | pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices 168 | 169 | # Search for detections 170 | if pi.shape[0]: 171 | # Prediction to target ious 172 | ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices 173 | 174 | # Append detections 175 | for j in (ious > iouv[0]).nonzero(as_tuple=False): 176 | d = ti[i[j]] # detected target 177 | if d not in detected: 178 | detected.append(d) 179 | correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn 180 | if len(detected) == nl: # all targets already located in image 181 | break 182 | 183 | # Append statistics (correct, conf, pcls, tcls) 184 | stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) 185 | 186 | # Plot images 187 | if batch_i < 1: 188 | f = Path(save_dir) / ('test_batch%g_gt.jpg' % batch_i) # filename 189 | plot_images(img, targets, paths, str(f), names) # ground truth 190 | f = Path(save_dir) / ('test_batch%g_pred.jpg' % batch_i) 191 | plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions 192 | 193 | # Compute statistics 194 | stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy 195 | if len(stats) and stats[0].any(): 196 | p, r, ap, f1, ap_class = ap_per_class(*stats) 197 | p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95] 198 | mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() 199 | nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class 200 | else: 201 | nt = torch.zeros(1) 202 | 203 | # Print results 204 | pf = '%20s' + '%12.3g' * 6 # print format 205 | print(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) 206 | 207 | # Print results per class 208 | if verbose and nc > 1 and len(stats): 209 | for i, c in enumerate(ap_class): 210 | print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) 211 | 212 | # Print speeds 213 | t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple 214 | if not training: 215 | print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t) 216 | 217 | # Save JSON 218 | if save_json and len(jdict): 219 | f = 'detections_val2017_%s_results.json' % \ 220 | (weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename 221 | print('\nCOCO mAP with pycocotools... saving %s...' % f) 222 | with open(f, 'w') as file: 223 | json.dump(jdict, file) 224 | 225 | try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb 226 | from pycocotools.coco import COCO 227 | from pycocotools.cocoeval import COCOeval 228 | 229 | imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] 230 | cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api 231 | cocoDt = cocoGt.loadRes(f) # initialize COCO pred api 232 | cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') 233 | cocoEval.params.imgIds = imgIds # image IDs to evaluate 234 | cocoEval.evaluate() 235 | cocoEval.accumulate() 236 | cocoEval.summarize() 237 | map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) 238 | except Exception as e: 239 | print('ERROR: pycocotools unable to run: %s' % e) 240 | 241 | # Return results 242 | model.float() # for training 243 | maps = np.zeros(nc) + map 244 | for i, c in enumerate(ap_class): 245 | maps[c] = ap[i] 246 | return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t 247 | 248 | 249 | if __name__ == '__main__': 250 | parser = argparse.ArgumentParser(prog='test.py') 251 | parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') 252 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path') 253 | parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch') 254 | parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') 255 | parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold') 256 | parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS') 257 | parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') 258 | parser.add_argument('--task', default='val', help="'val', 'test', 'study'") 259 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') 260 | parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') 261 | parser.add_argument('--augment', action='store_true', help='augmented inference') 262 | parser.add_argument('--merge', action='store_true', help='use Merge NMS') 263 | parser.add_argument('--verbose', action='store_true', help='report mAP by class') 264 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') 265 | opt = parser.parse_args() 266 | opt.save_json |= opt.data.endswith('coco.yaml') 267 | opt.data = check_file(opt.data) # check file 268 | print(opt) 269 | 270 | if opt.task in ['val', 'test']: # run normally 271 | test(opt.data, 272 | opt.weights, 273 | opt.batch_size, 274 | opt.img_size, 275 | opt.conf_thres, 276 | opt.iou_thres, 277 | opt.save_json, 278 | opt.single_cls, 279 | opt.augment, 280 | opt.verbose) 281 | 282 | elif opt.task == 'study': # run over a range of settings and save/plot 283 | for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: 284 | f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to 285 | x = list(range(320, 800, 64)) # x axis 286 | y = [] # y axis 287 | for i in x: # img-size 288 | print('\nRunning %s point %s...' % (f, i)) 289 | r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json) 290 | y.append(r + t) # results and times 291 | np.savetxt(f, y, fmt='%10.4g') # save 292 | os.system('zip -r study.zip study_*.txt') 293 | # utils.general.plot_study_txt(f, x) # plot 294 | -------------------------------------------------------------------------------- /yolov5/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/utils/__init__.py -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/utils/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/datasets.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/utils/__pycache__/datasets.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/general.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/utils/__pycache__/general.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/google_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/utils/__pycache__/google_utils.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/__pycache__/torch_utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/linhaoqi027/yolov5_openvino_sdk/0802ed9502bdfe7488207d9b32f77c0378f65cbc/yolov5/utils/__pycache__/torch_utils.cpython-36.pyc -------------------------------------------------------------------------------- /yolov5/utils/activations.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | # Swish https://arxiv.org/pdf/1905.02244.pdf --------------------------------------------------------------------------- 7 | class Swish(nn.Module): # 8 | @staticmethod 9 | def forward(x): 10 | return x * torch.sigmoid(x) 11 | 12 | 13 | class MemoryEfficientSwish(nn.Module): 14 | class F(torch.autograd.Function): 15 | @staticmethod 16 | def forward(ctx, x): 17 | ctx.save_for_backward(x) 18 | return x * torch.sigmoid(x) 19 | 20 | @staticmethod 21 | def backward(ctx, grad_output): 22 | x = ctx.saved_tensors[0] 23 | sx = torch.sigmoid(x) 24 | return grad_output * (sx * (1 + x * (1 - sx))) 25 | 26 | def forward(self, x): 27 | return self.F.apply(x) 28 | 29 | 30 | # Mish https://github.com/digantamisra98/Mish -------------------------------------------------------------------------- 31 | class Mish(nn.Module): 32 | @staticmethod 33 | def forward(x): 34 | return x * F.softplus(x).tanh() 35 | 36 | 37 | class MemoryEfficientMish(nn.Module): 38 | class F(torch.autograd.Function): 39 | @staticmethod 40 | def forward(ctx, x): 41 | ctx.save_for_backward(x) 42 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x))) 43 | 44 | @staticmethod 45 | def backward(ctx, grad_output): 46 | x = ctx.saved_tensors[0] 47 | sx = torch.sigmoid(x) 48 | fx = F.softplus(x).tanh() 49 | return grad_output * (fx + x * sx * (1 - fx * fx)) 50 | 51 | def forward(self, x): 52 | return self.F.apply(x) 53 | 54 | 55 | # FReLU https://arxiv.org/abs/2007.11824 ------------------------------------------------------------------------------- 56 | class FReLU(nn.Module): 57 | def __init__(self, c1, k=3): # ch_in, kernel 58 | super().__init__() 59 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1) 60 | self.bn = nn.BatchNorm2d(c1) 61 | 62 | def forward(self, x): 63 | return torch.max(x, self.bn(self.conv(x))) 64 | -------------------------------------------------------------------------------- /yolov5/utils/google_utils.py: -------------------------------------------------------------------------------- 1 | # This file contains google utils: https://cloud.google.com/storage/docs/reference/libraries 2 | # pip install --upgrade google-cloud-storage 3 | # from google.cloud import storage 4 | 5 | import os 6 | import platform 7 | import time 8 | from pathlib import Path 9 | import torch 10 | 11 | 12 | def attempt_download(weights): 13 | # Attempt to download pretrained weights if not found locally 14 | weights = weights.strip().replace("'", '') 15 | file = Path(weights).name 16 | 17 | msg = weights + ' missing, try downloading from https://github.com/ultralytics/yolov5/releases/' 18 | models = ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'] # available models 19 | 20 | if file in models and not os.path.isfile(weights): 21 | # Google Drive 22 | # d = {'yolov5s.pt': '1R5T6rIyy3lLwgFXNms8whc-387H0tMQO', 23 | # 'yolov5m.pt': '1vobuEExpWQVpXExsJ2w-Mbf3HJjWkQJr', 24 | # 'yolov5l.pt': '1hrlqD1Wdei7UT4OgT785BEk1JwnSvNEV', 25 | # 'yolov5x.pt': '1mM8aZJlWTxOg7BZJvNUMrTnA2AbeCVzS'} 26 | # r = gdrive_download(id=d[file], name=weights) if file in d else 1 27 | # if r == 0 and os.path.exists(weights) and os.path.getsize(weights) > 1E6: # check 28 | # return 29 | 30 | try: # GitHub 31 | url = 'https://github.com/ultralytics/yolov5/releases/download/v3.0/' + file 32 | print('Downloading %s to %s...' % (url, weights)) 33 | if platform.system() == 'Darwin': # avoid MacOS python requests certificate error 34 | r = os.system('curl -L %s -o %s' % (url, weights)) 35 | else: 36 | torch.hub.download_url_to_file(url, weights) 37 | assert os.path.exists(weights) and os.path.getsize(weights) > 1E6 # check 38 | except Exception as e: # GCP 39 | print('Download error: %s' % e) 40 | url = 'https://storage.googleapis.com/ultralytics/yolov5/ckpt/' + file 41 | print('Downloading %s to %s...' % (url, weights)) 42 | r = os.system('curl -L %s -o %s' % (url, weights)) # torch.hub.download_url_to_file(url, weights) 43 | finally: 44 | if not (os.path.exists(weights) and os.path.getsize(weights) > 1E6): # check 45 | os.remove(weights) if os.path.exists(weights) else None # remove partial downloads 46 | print('ERROR: Download failure: %s' % msg) 47 | print('') 48 | return 49 | 50 | 51 | def gdrive_download(id='1n_oKgR81BJtqk75b00eAjdv03qVCQn2f', name='coco128.zip'): 52 | # Downloads a file from Google Drive. from utils.google_utils import *; gdrive_download() 53 | t = time.time() 54 | 55 | print('Downloading https://drive.google.com/uc?export=download&id=%s as %s... ' % (id, name), end='') 56 | os.remove(name) if os.path.exists(name) else None # remove existing 57 | os.remove('cookie') if os.path.exists('cookie') else None 58 | 59 | # Attempt file download 60 | out = "NUL" if platform.system() == "Windows" else "/dev/null" 61 | os.system('curl -c ./cookie -s -L "drive.google.com/uc?export=download&id=%s" > %s ' % (id, out)) 62 | if os.path.exists('cookie'): # large file 63 | s = 'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm=%s&id=%s" -o %s' % (get_token(), id, name) 64 | else: # small file 65 | s = 'curl -s -L -o %s "drive.google.com/uc?export=download&id=%s"' % (name, id) 66 | r = os.system(s) # execute, capture return 67 | os.remove('cookie') if os.path.exists('cookie') else None 68 | 69 | # Error check 70 | if r != 0: 71 | os.remove(name) if os.path.exists(name) else None # remove partial 72 | print('Download error ') # raise Exception('Download error') 73 | return r 74 | 75 | # Unzip if archive 76 | if name.endswith('.zip'): 77 | print('unzipping... ', end='') 78 | os.system('unzip -q %s' % name) # unzip 79 | os.remove(name) # remove zip to free space 80 | 81 | print('Done (%.1fs)' % (time.time() - t)) 82 | return r 83 | 84 | 85 | def get_token(cookie="./cookie"): 86 | with open(cookie) as f: 87 | for line in f: 88 | if "download" in line: 89 | return line.split()[-1] 90 | return "" 91 | 92 | # def upload_blob(bucket_name, source_file_name, destination_blob_name): 93 | # # Uploads a file to a bucket 94 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python 95 | # 96 | # storage_client = storage.Client() 97 | # bucket = storage_client.get_bucket(bucket_name) 98 | # blob = bucket.blob(destination_blob_name) 99 | # 100 | # blob.upload_from_filename(source_file_name) 101 | # 102 | # print('File {} uploaded to {}.'.format( 103 | # source_file_name, 104 | # destination_blob_name)) 105 | # 106 | # 107 | # def download_blob(bucket_name, source_blob_name, destination_file_name): 108 | # # Uploads a blob from a bucket 109 | # storage_client = storage.Client() 110 | # bucket = storage_client.get_bucket(bucket_name) 111 | # blob = bucket.blob(source_blob_name) 112 | # 113 | # blob.download_to_filename(destination_file_name) 114 | # 115 | # print('Blob {} downloaded to {}.'.format( 116 | # source_blob_name, 117 | # destination_file_name)) 118 | -------------------------------------------------------------------------------- /yolov5/utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | import math 2 | import os 3 | import time 4 | import logging 5 | from copy import deepcopy 6 | 7 | import torch 8 | import torch.backends.cudnn as cudnn 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import torchvision.models as models 12 | 13 | logger = logging.getLogger(__name__) 14 | 15 | 16 | def init_seeds(seed=0): 17 | torch.manual_seed(seed) 18 | 19 | # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html 20 | if seed == 0: # slower, more reproducible 21 | cudnn.deterministic = True 22 | cudnn.benchmark = False 23 | else: # faster, less reproducible 24 | cudnn.deterministic = False 25 | cudnn.benchmark = True 26 | 27 | 28 | def select_device(device='', batch_size=None): 29 | # device = 'cpu' or '0' or '0,1,2,3' 30 | cpu_request = device.lower() == 'cpu' 31 | if device and not cpu_request: # if device requested other than 'cpu' 32 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable 33 | assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device # check availablity 34 | 35 | cuda = False if cpu_request else torch.cuda.is_available() 36 | if cuda: 37 | c = 1024 ** 2 # bytes to MB 38 | ng = torch.cuda.device_count() 39 | if ng > 1 and batch_size: # check that batch_size is compatible with device_count 40 | assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng) 41 | x = [torch.cuda.get_device_properties(i) for i in range(ng)] 42 | s = 'Using CUDA ' 43 | for i in range(0, ng): 44 | if i == 1: 45 | s = ' ' * len(s) 46 | logger.info("%sdevice%g _CudaDeviceProperties(name='%s', total_memory=%dMB)" % 47 | (s, i, x[i].name, x[i].total_memory / c)) 48 | else: 49 | logger.info('Using CPU') 50 | 51 | logger.info('') # skip a line 52 | return torch.device('cuda:0' if cuda else 'cpu') 53 | 54 | 55 | def time_synchronized(): 56 | torch.cuda.synchronize() if torch.cuda.is_available() else None 57 | return time.time() 58 | 59 | 60 | def is_parallel(model): 61 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel) 62 | 63 | 64 | def intersect_dicts(da, db, exclude=()): 65 | # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values 66 | 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} 67 | 68 | 69 | def initialize_weights(model): 70 | for m in model.modules(): 71 | t = type(m) 72 | if t is nn.Conv2d: 73 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 74 | elif t is nn.BatchNorm2d: 75 | m.eps = 1e-3 76 | m.momentum = 0.03 77 | elif t in [nn.LeakyReLU, nn.ReLU, nn.ReLU6]: 78 | m.inplace = True 79 | 80 | 81 | def find_modules(model, mclass=nn.Conv2d): 82 | # Finds layer indices matching module class 'mclass' 83 | return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)] 84 | 85 | 86 | def sparsity(model): 87 | # Return global model sparsity 88 | a, b = 0., 0. 89 | for p in model.parameters(): 90 | a += p.numel() 91 | b += (p == 0).sum() 92 | return b / a 93 | 94 | 95 | def prune(model, amount=0.3): 96 | # Prune model to requested global sparsity 97 | import torch.nn.utils.prune as prune 98 | print('Pruning model... ', end='') 99 | for name, m in model.named_modules(): 100 | if isinstance(m, nn.Conv2d): 101 | prune.l1_unstructured(m, name='weight', amount=amount) # prune 102 | prune.remove(m, 'weight') # make permanent 103 | print(' %.3g global sparsity' % sparsity(model)) 104 | 105 | 106 | def fuse_conv_and_bn(conv, bn): 107 | # https://tehnokv.com/posts/fusing-batchnorm-and-conv/ 108 | with torch.no_grad(): 109 | # init 110 | fusedconv = nn.Conv2d(conv.in_channels, 111 | conv.out_channels, 112 | kernel_size=conv.kernel_size, 113 | stride=conv.stride, 114 | padding=conv.padding, 115 | bias=True).to(conv.weight.device) 116 | 117 | # prepare filters 118 | w_conv = conv.weight.clone().view(conv.out_channels, -1) 119 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var))) 120 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size())) 121 | 122 | # prepare spatial bias 123 | b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias 124 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps)) 125 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn) 126 | 127 | return fusedconv 128 | 129 | 130 | def model_info(model, verbose=False): 131 | # Plots a line-by-line description of a PyTorch model 132 | n_p = sum(x.numel() for x in model.parameters()) # number parameters 133 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients 134 | if verbose: 135 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma')) 136 | for i, (name, p) in enumerate(model.named_parameters()): 137 | name = name.replace('module_list.', '') 138 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' % 139 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) 140 | 141 | try: # FLOPS 142 | from thop import profile 143 | flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, 64, 64),), verbose=False)[0] / 1E9 * 2 144 | fs = ', %.1f GFLOPS' % (flops * 100) # 640x640 FLOPS 145 | except: 146 | fs = '' 147 | 148 | logger.info( 149 | 'Model Summary: %g layers, %g parameters, %g gradients%s' % (len(list(model.parameters())), n_p, n_g, fs)) 150 | 151 | 152 | def load_classifier(name='resnet101', n=2): 153 | # Loads a pretrained model reshaped to n-class output 154 | model = models.__dict__[name](pretrained=True) 155 | 156 | # Display model properties 157 | input_size = [3, 224, 224] 158 | input_space = 'RGB' 159 | input_range = [0, 1] 160 | mean = [0.485, 0.456, 0.406] 161 | std = [0.229, 0.224, 0.225] 162 | for x in ['input_size', 'input_space', 'input_range', 'mean', 'std']: 163 | print(x + ' =', eval(x)) 164 | 165 | # Reshape output to n classes 166 | filters = model.fc.weight.shape[1] 167 | model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True) 168 | model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True) 169 | model.fc.out_features = n 170 | return model 171 | 172 | 173 | def scale_img(img, ratio=1.0, same_shape=False): # img(16,3,256,416), r=ratio 174 | # scales img(bs,3,y,x) by ratio 175 | if ratio == 1.0: 176 | return img 177 | else: 178 | h, w = img.shape[2:] 179 | s = (int(h * ratio), int(w * ratio)) # new size 180 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize 181 | if not same_shape: # pad/crop img 182 | gs = 32 # (pixels) grid size 183 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)] 184 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean 185 | 186 | 187 | def copy_attr(a, b, include=(), exclude=()): 188 | # Copy attributes from b to a, options to only include [...] and to exclude [...] 189 | for k, v in b.__dict__.items(): 190 | if (len(include) and k not in include) or k.startswith('_') or k in exclude: 191 | continue 192 | else: 193 | setattr(a, k, v) 194 | 195 | 196 | class ModelEMA: 197 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models 198 | Keep a moving average of everything in the model state_dict (parameters and buffers). 199 | This is intended to allow functionality like 200 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage 201 | A smoothed version of the weights is necessary for some training schemes to perform well. 202 | This class is sensitive where it is initialized in the sequence of model init, 203 | GPU assignment and distributed training wrappers. 204 | """ 205 | 206 | def __init__(self, model, decay=0.9999, updates=0): 207 | # Create EMA 208 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA 209 | # if next(model.parameters()).device.type != 'cpu': 210 | # self.ema.half() # FP16 EMA 211 | self.updates = updates # number of EMA updates 212 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) 213 | for p in self.ema.parameters(): 214 | p.requires_grad_(False) 215 | 216 | def update(self, model): 217 | # Update EMA parameters 218 | with torch.no_grad(): 219 | self.updates += 1 220 | d = self.decay(self.updates) 221 | 222 | msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict 223 | for k, v in self.ema.state_dict().items(): 224 | if v.dtype.is_floating_point: 225 | v *= d 226 | v += (1. - d) * msd[k].detach() 227 | 228 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): 229 | # Update EMA attributes 230 | copy_attr(self.ema, model, include, exclude) 231 | -------------------------------------------------------------------------------- /yolov5/weights/download_weights.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Download common models 3 | 4 | python -c " 5 | from utils.google_utils import *; 6 | attempt_download('weights/yolov5s.pt'); 7 | attempt_download('weights/yolov5m.pt'); 8 | attempt_download('weights/yolov5l.pt'); 9 | attempt_download('weights/yolov5x.pt') 10 | " 11 | --------------------------------------------------------------------------------