def main():
args = get_args()
# Load parameter
_ = nn.load_parameters(args.weights)
# Build a YOLO v2 network
x = nn.Variable((1, 3, args.height, args.width))
y = yolov2.yolov2(x / 255.0, args.num_anchors, args.classes, test=True)
y = yolov2.yolov2_activate(y, args.num_anchors, args.anchors)
y = F.nms_detection2d(y, args.thresh, args.nms, args.nms_per_class)
# Save NNP file (used in C++ inference later.).
runtime_contents = {
'networks': [{
'name': 'runtime',
'batch_size': 1,
'outputs': {
'y': y
},
'names': {
'x': x
}
}],
'executors': [{
'name': 'runtime',
'network': 'runtime',
'data': ['x'],
'output': ['y']
}]
}
import nnabla.utils.save
nnabla.utils.save.save(args.nnp,
runtime_contents,
variable_batch_size=False)
def main():
args = get_args()
names = np.genfromtxt(args.class_names, dtype=str, delimiter='?')
rng = np.random.RandomState(1223)
colors = rng.randint(0, 256, (args.classes, 3)).astype(np.uint8)
colors = [tuple(c.tolist()) for c in colors]
# Set context
from nnabla.ext_utils import get_extension_context
ctx = get_extension_context(args.context,
device_id=args.device_id,
type_config=args.type_config)
nn.set_default_context(ctx)
# Load parameter
_ = nn.load_parameters(args.weights)
# Build a YOLO v2 network
feature_dict = {}
x = nn.Variable((1, 3, args.width, args.width))
y = yolov2.yolov2(x,
args.anchors,
args.classes,
test=True,
feature_dict=feature_dict)
y = yolov2.yolov2_activate(y, args.anchors, args.biases)
y = F.nms_detection2d(y, args.thresh, args.nms, args.nms_per_class)
# Read image
img_orig = imread(args.input)
im_h, im_w, _ = img_orig.shape
# letterbox
w = args.width
h = args.width
if (w * 1.0 / im_w) < (h * 1. / im_h):
new_w = w
new_h = int((im_h * w) / im_w)
else:
new_h = h
new_w = int((im_w * h) / im_h)
patch = imresize(img_orig, (new_h, new_w)) / 255.
img = np.ones((h, w, 3), np.float32) * 0.5
# resize
x0 = int((w - new_w) / 2)
y0 = int((h - new_h) / 2)
img[y0:y0 + new_h, x0:x0 + new_w] = patch
# Execute YOLO v2
print("forward")
in_img = img.transpose(2, 0, 1).reshape(1, 3, args.width, args.width)
x.d = in_img
y.forward(clear_buffer=True)
print("done")
bboxes = y.d[0]
img_draw = draw_bounding_boxes(img_orig, bboxes, im_w, im_h, names, colors,
new_w * 1.0 / w, new_h * 1.0 / h,
args.thresh)
imsave(args.output, img_draw)
# Timing
s = time.time()
n_time = 10
for i in range(n_time):
x.d = in_img
y.forward(clear_buffer=True)
# Invoking device-to-host copy if CUDA
# so that time contains data transfer.
_ = y.d
print("Processing time: {:.1f} [ms/image]".format(
(time.time() - s) / n_time * 1000))