1# Licensed to the Apache Software Foundation (ASF) under one 2# or more contributor license agreements. See the NOTICE file 3# distributed with this work for additional information 4# regarding copyright ownership. The ASF licenses this file 5# to you under the Apache License, Version 2.0 (the 6# "License"); you may not use this file except in compliance 7# with the License. You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, 12# software distributed under the License is distributed on an 13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14# KIND, either express or implied. See the License for the 15# specific language governing permissions and limitations 16# under the License. 17""" 18Compile YOLO-V2 and YOLO-V3 in DarkNet Models 19================================= 20**Author**: `Siju Samuel <https://siju-samuel.github.io/>`_ 21 22This article is an introductory tutorial to deploy darknet models with NNVM. 23All the required models and libraries will be downloaded from the internet by the script. 24This script runs the YOLO-V2 and YOLO-V3 Model with the bounding boxes 25Darknet parsing have dependancy with CFFI and CV2 library 26Please install CFFI and CV2 before executing this script 27 28.. code-block:: bash 29 30 pip install cffi 31 pip install opencv-python 32""" 33 34import nnvm 35import nnvm.frontend.darknet 36import tvm.relay.testing.yolo_detection 37import tvm.relay.testing.darknet 38import matplotlib.pyplot as plt 39import numpy as np 40import tvm 41import sys 42 43from ctypes import * 44from tvm.contrib.download import download_testdata 45from tvm.relay.testing.darknet import __darknetffi__ 46 47# Model name 48MODEL_NAME = 'yolov3' 49 50###################################################################### 51# Download required files 52# ----------------------- 53# Download cfg and weights file if first time. 54CFG_NAME = MODEL_NAME + '.cfg' 55WEIGHTS_NAME = MODEL_NAME + '.weights' 56REPO_URL = 'https://github.com/siju-samuel/darknet/blob/master/' 57CFG_URL = REPO_URL + 'cfg/' + CFG_NAME + '?raw=true' 58WEIGHTS_URL = 'https://pjreddie.com/media/files/' + WEIGHTS_NAME 59 60cfg_path = download_testdata(CFG_URL, CFG_NAME, module="darknet") 61weights_path = download_testdata(WEIGHTS_URL, WEIGHTS_NAME, module="darknet") 62 63# Download and Load darknet library 64if sys.platform in ['linux', 'linux2']: 65 DARKNET_LIB = 'libdarknet2.0.so' 66 DARKNET_URL = REPO_URL + 'lib/' + DARKNET_LIB + '?raw=true' 67elif sys.platform == 'darwin': 68 DARKNET_LIB = 'libdarknet_mac2.0.so' 69 DARKNET_URL = REPO_URL + 'lib_osx/' + DARKNET_LIB + '?raw=true' 70else: 71 err = "Darknet lib is not supported on {} platform".format(sys.platform) 72 raise NotImplementedError(err) 73 74lib_path = download_testdata(DARKNET_URL, DARKNET_LIB, module="darknet") 75 76DARKNET_LIB = __darknetffi__.dlopen(lib_path) 77net = DARKNET_LIB.load_network(cfg_path.encode('utf-8'), weights_path.encode('utf-8'), 0) 78dtype = 'float32' 79batch_size = 1 80 81print("Converting darknet to nnvm symbols...") 82sym, params = nnvm.frontend.darknet.from_darknet(net, dtype) 83 84###################################################################### 85# Compile the model on NNVM 86# ------------------------- 87# compile the model 88target = 'llvm' 89ctx = tvm.cpu(0) 90data = np.empty([batch_size, net.c, net.h, net.w], dtype) 91shape = {'data': data.shape} 92print("Compiling the model...") 93dtype_dict = {} 94with nnvm.compiler.build_config(opt_level=2): 95 graph, lib, params = nnvm.compiler.build(sym, target, shape, dtype_dict, params) 96 97[neth, netw] = shape['data'][2:] # Current image shape is 608x608 98###################################################################### 99# Load a test image 100# -------------------------------------------------------------------- 101test_image = 'dog.jpg' 102print("Loading the test image...") 103img_url = 'https://github.com/siju-samuel/darknet/blob/master/data/' + \ 104 test_image + '?raw=true' 105img_path = download_testdata(img_url, test_image, "data") 106 107data = tvm.relay.testing.darknet.load_image(img_path, netw, neth) 108###################################################################### 109# Execute on TVM Runtime 110# ---------------------- 111# The process is no different from other examples. 112from tvm.contrib import graph_runtime 113 114m = graph_runtime.create(graph, lib, ctx) 115 116# set inputs 117m.set_input('data', tvm.nd.array(data.astype(dtype))) 118m.set_input(**params) 119# execute 120print("Running the test image...") 121 122m.run() 123# get outputs 124tvm_out = [] 125if MODEL_NAME == 'yolov2': 126 layer_out = {} 127 layer_out['type'] = 'Region' 128 # Get the region layer attributes (n, out_c, out_h, out_w, classes, coords, background) 129 layer_attr = m.get_output(2).asnumpy() 130 layer_out['biases'] = m.get_output(1).asnumpy() 131 out_shape = (layer_attr[0], layer_attr[1]//layer_attr[0], 132 layer_attr[2], layer_attr[3]) 133 layer_out['output'] = m.get_output(0).asnumpy().reshape(out_shape) 134 layer_out['classes'] = layer_attr[4] 135 layer_out['coords'] = layer_attr[5] 136 layer_out['background'] = layer_attr[6] 137 tvm_out.append(layer_out) 138 139elif MODEL_NAME == 'yolov3': 140 for i in range(3): 141 layer_out = {} 142 layer_out['type'] = 'Yolo' 143 # Get the yolo layer attributes (n, out_c, out_h, out_w, classes, total) 144 layer_attr = m.get_output(i*4+3).asnumpy() 145 layer_out['biases'] = m.get_output(i*4+2).asnumpy() 146 layer_out['mask'] = m.get_output(i*4+1).asnumpy() 147 out_shape = (layer_attr[0], layer_attr[1]//layer_attr[0], 148 layer_attr[2], layer_attr[3]) 149 layer_out['output'] = m.get_output(i*4).asnumpy().reshape(out_shape) 150 layer_out['classes'] = layer_attr[4] 151 tvm_out.append(layer_out) 152 153# do the detection and bring up the bounding boxes 154thresh = 0.5 155nms_thresh = 0.45 156img = tvm.relay.testing.darknet.load_image_color(img_path) 157_, im_h, im_w = img.shape 158dets = tvm.relay.testing.yolo_detection.fill_network_boxes((netw, neth), (im_w, im_h), thresh, 159 1, tvm_out) 160last_layer = net.layers[net.n - 1] 161tvm.relay.testing.yolo_detection.do_nms_sort(dets, last_layer.classes, nms_thresh) 162 163coco_name = 'coco.names' 164coco_url = 'https://github.com/siju-samuel/darknet/blob/master/data/' + coco_name + '?raw=true' 165font_name = 'arial.ttf' 166font_url = 'https://github.com/siju-samuel/darknet/blob/master/data/' + font_name + '?raw=true' 167coco_path = download_testdata(coco_url, coco_name, module='data') 168font_path = download_testdata(font_url, font_name, module='data') 169 170with open(coco_path) as f: 171 content = f.readlines() 172 173names = [x.strip() for x in content] 174 175tvm.relay.testing.yolo_detection.draw_detections(font_path, img, dets, thresh, names, last_layer.classes) 176plt.imshow(img.transpose(1, 2, 0)) 177plt.show() 178