コード例 #1
0
ファイル: yolov2_nnp.py プロジェクト: Nob100/nnabla-examples
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)
コード例 #2
0
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))