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