# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, unused-variable, unused-argument, no-init
"""
Yolo detection boxes helper functions
====================
DarkNet helper functions for yolo and image loading.
This functions will not be loaded by default.
These are utility functions used for testing and tutorial file.
"""
from __future__ import division
import math
from collections import namedtuple
from functools import cmp_to_key
import numpy as np
Box = namedtuple('Box', ['x', 'y', 'w', 'h'])
def nms_comparator(a, b):
if 'sort_class' in b and b['sort_class'] >= 0:
diff = a['prob'][b['sort_class']] - b['prob'][b['sort_class']]
else:
diff = a['objectness'] - b['objectness']
return diff
def _correct_boxes(dets, w, h, netw, neth, relative):
new_w, new_h = (netw, (h*netw)//w) if (netw/w < neth/h) else ((w*neth//h), neth)
for det in dets:
b = det['bbox']
b = b._replace(x=(b.x - (netw - new_w)/2/netw) / (new_w/netw))
b = b._replace(y=(b.y - (neth - new_h)/2/neth) / (new_h/neth))
b = b._replace(w=b.w * netw/new_w)
b = b._replace(h=b.h * neth/new_h)
if not relative:
b = b._replace(x=b.x * w)
b = b._replace(w=b.w * w)
b = b._replace(y=b.y * h)
b = b._replace(h=b.h * h)
det['bbox'] = b
return dets
def _overlap(x1, w1, x2, w2):
l1 = x1 - w1/2
l2 = x2 - w2/2
left = l1 if l1 > l2 else l2
r1 = x1 + w1/2
r2 = x2 + w2/2
right = r1 if r1 < r2 else r2
return right - left
def _box_intersection(a, b):
w = _overlap(a.x, a.w, b.x, b.w)
h = _overlap(a.y, a.h, b.y, b.h)
if w < 0 or h < 0:
return 0
return w*h
def _box_union(a, b):
i = _box_intersection(a, b)
u = a.w*a.h + b.w*b.h - i
return u
def _box_iou(a, b):
return _box_intersection(a, b)/_box_union(a, b)
def _get_box(data, biases, n, location, lw, lh, w, h):
bx = (location[2] + data[location[0]][0][location[1]][location[2]]) / lw
by = (location[1] + data[location[0]][1][location[1]][location[2]]) / lh
bw = np.exp(data[location[0]][2][location[1]][location[2]]) * biases[2*n] / w
bh = np.exp(data[location[0]][3][location[1]][location[2]]) * biases[2*n+1] / h
return Box(bx, by, bw, bh)
def _get_yolo_detections(l, im_shape, net_shape, thresh, relative, dets):
data = l['output']
active_data_loc = np.asarray(np.where(data[:, 4, :, :] > thresh))
before_correct_dets = []
for i in range(active_data_loc.shape[1]):
location = [active_data_loc[0][i], active_data_loc[1][i], active_data_loc[2][i]]
box_b = _get_box(data, l['biases'], np.asarray(l['mask'])[location[0]], location,
data.shape[2], data.shape[3], net_shape[0], net_shape[1])
objectness = data[location[0]][4][location[1]][location[2]]
classes = l['classes']
prob = objectness*data[location[0], 5:5 + 1 + classes, location[1], location[2]]
prob[prob < thresh] = 0
detection = {}
detection['bbox'] = box_b
detection['classes'] = classes
detection['prob'] = prob
detection['objectness'] = objectness
before_correct_dets.append(detection)
dets.extend(_correct_boxes(before_correct_dets, im_shape[0], im_shape[1],
net_shape[0], net_shape[1], relative))
def _get_region_detections(l, im_shape, net_shape, thresh, relative, dets):
data = l['output']
before_correct_dets = []
for row in range(data.shape[2]):
for col in range(data.shape[3]):
for n in range(data.shape[0]):
prob = [0]*l['classes']
scale = data[n, l['coords'], row, col] if not l['background'] else 1
location = [n, row, col]
box_b = _get_box(data, l['biases'], n, location,
data.shape[2], data.shape[3], data.shape[2], data.shape[3])
objectness = scale if scale > thresh else 0
if objectness:
prob = scale * data[n, l['coords']+1: l['coords']+1+l['classes'],
row, col]
prob[prob < thresh] = 0
detection = {}
detection['bbox'] = box_b
detection['prob'] = prob
detection['objectness'] = objectness
before_correct_dets.append(detection)
_correct_boxes(before_correct_dets, im_shape[0], im_shape[1],
net_shape[0], net_shape[1], relative)
dets.extend(before_correct_dets)
def fill_network_boxes(net_shape, im_shape,
thresh, relative, tvm_out):
dets = []
for layer in tvm_out:
if layer['type'] == 'Yolo':
_get_yolo_detections(layer, im_shape, net_shape, thresh, relative, dets)
elif layer['type'] == 'Region':
_get_region_detections(layer, im_shape, net_shape, thresh, relative, dets)
return dets
def do_nms_sort(dets, classes, thresh):
"Does the sorting based on the threshold values"
k = len(dets)-1
cnt = 0
while cnt < k:
if dets[cnt]['objectness'] == 0:
dets[k], dets[cnt] = dets[cnt], dets[k]
k = k - 1
else:
cnt = cnt + 1
total = k+1
for k in range(classes):
for i in range(total):
dets[i]['sort_class'] = k
dets[0:total] = sorted(dets[0:total],
key=cmp_to_key(nms_comparator), reverse=True)
for i in range(total):
if dets[i]['prob'][k] == 0:
continue
a = dets[i]['bbox']
for j in range(i+1, total):
b = dets[j]['bbox']
if _box_iou(a, b) > thresh:
dets[j]['prob'][k] = 0
def draw_detections(font_path, im, dets, thresh, names, classes):
"Draw the markings around the detected region"
for det in dets:
labelstr = []
category = -1
for j in range(classes):
if det['prob'][j] > thresh:
if category == -1:
category = j
labelstr.append(names[j] + " " + str(round(det['prob'][j], 4)))
if category > -1:
imc, imh, imw = im.shape
width = int(imh * 0.006)
offset = category*123457 % classes
red = _get_color(2, offset, classes)
green = _get_color(1, offset, classes)
blue = _get_color(0, offset, classes)
rgb = [red, green, blue]
b = det['bbox']
left = int((b.x-b.w/2.)*imw)
right = int((b.x+b.w/2.)*imw)
top = int((b.y-b.h/2.)*imh)
bot = int((b.y+b.h/2.)*imh)
if left < 0:
left = 0
if right > imw-1:
right = imw-1
if top < 0:
top = 0
if bot > imh-1:
bot = imh-1
_draw_box_width(im, left, top, right, bot, width, red, green, blue)
label = _get_label(font_path, ''.join(labelstr), rgb)
_draw_label(im, top + width, left, label, rgb)
def _get_pixel(im, x, y, c):
return im[c][y][x]
def _set_pixel(im, x, y, c, val):
if x < 0 or y < 0 or c < 0 or x >= im.shape[2] or y >= im.shape[1] or c >= im.shape[0]:
return
im[c][y][x] = val
def _draw_label(im, r, c, label, rgb):
w = label.shape[2]
h = label.shape[1]
if (r - h) >= 0:
r = r - h
for j in range(h):
if j < h and (j + r) < im.shape[1]:
for i in range(w):
if i < w and (i + c) < im.shape[2]:
for k in range(label.shape[0]):
val = _get_pixel(label, i, j, k)
_set_pixel(im, i+c, j+r, k, val)#rgb[k] * val)
def _get_label(font_path, labelstr, rgb):
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
text = labelstr
colorText = "black"
testDraw = ImageDraw.Draw(Image.new('RGB', (1, 1)))
font = ImageFont.truetype(font_path, 25)
width, height = testDraw.textsize(labelstr, font=font)
img = Image.new('RGB', (width, height), color=(int(rgb[0]*255), int(rgb[1]*255),
int(rgb[2]*255)))
d = ImageDraw.Draw(img)
d.text((0, 0), text, fill=colorText, font=font)
opencvImage = np.divide(np.asarray(img), 255)
return opencvImage.transpose(2, 0, 1)
def _get_color(c, x, max_value):
c = int(c)
colors = [[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]]
ratio = (float(x)/float(max_value)) * 5
i = int(math.floor(ratio))
j = int(math.ceil(ratio))
ratio -= i
r = (1-ratio) * colors[i][c] + ratio*colors[j][c]
return r
def _draw_box(im, x1, y1, x2, y2, r, g, b):
y1 = int(y1)
y2 = int(y2)
x1 = int(x1)
x2 = int(x2)
ac, ah, aw = im.shape
if x1 < 0:
x1 = 0
if x1 >= aw:
y1 = 0
if y1 >= ah:
y1 = ah - 1
if y2 < 0:
y2 = 0
if y2 >= ah:
y2 = ah - 1
for i in range(x1, x2):
im[0][y1][i] = r
im[0][y2][i] = r
im[1][y1][i] = g
im[1][y2][i] = g
im[2][y1][i] = b
im[2][y2][i] = b
for i in range(y1, y2):
im[0][i][x1] = r
im[0][i][x2] = r
im[1][i][x1] = g
im[1][i][x2] = g
im[2][i][x1] = b
im[2][i][x2] = b
def _draw_box_width(im, x1, y1, x2, y2, w, r, g, b):
for i in range(int(w)):
_draw_box(im, x1+i, y1+i, x2-i, y2-i, r, g, b)