#!/usr/bin/env python3

"""Test/benchmark deep_dream.py."""

import argparse
from pathlib import Path
import time

from PIL import Image

import deep_dream as dd


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--gpu', type=int, help='the CUDA device ID to use')
    parser.add_argument('--max-tile-size', type=int, default=512,
                        help='the maximum dimension of a tile')
    parser.add_argument('--save', action='store_true', help='save result to output.png')
    args = parser.parse_args()

    pwd = Path(__file__).parent
    cnn = dd.CNN(dd.GOOGLENET_BVLC, gpu=args.gpu)
    input_img = Image.open(str(pwd/'kodim/img0022.jpg')).resize((768, 512), Image.LANCZOS)
    print('Input image classes:')
    for c in cnn.classify(input_img, 5):
        print('%.3f %s' % c)
    cnn.dream(input_img, 'inception_3a/3x3', scales=4, n=1, max_tile_size=args.max_tile_size)
    t1 = time.perf_counter()
    img = cnn.dream(input_img, 'inception_3a/3x3', scales=4, n=10, max_tile_size=args.max_tile_size)
    print('Input image classes:')
    for c in cnn.classify(img, 5):
        print('%.3f %s' % c)
    t2 = time.perf_counter()
    print('Time taken: %.3f s' % (t2-t1))
    if args.save:
        dd.to_image(img).save('output.png')

if __name__ == '__main__':
    main()

@@ -227,16 +227,145 @@
 _doc__)%0A
+    parser.add_argument('--cpu-workers', type=int, default=0,%0A                        help='the number of CPU workers to start')%0A
     pars
@@ -386,16 +386,17 @@
 t('--gpu
+s
 ', type=
@@ -399,16 +399,63 @@
 ype=int,
+ nargs='+', default=%5B%5D,%0A                       
  help='t
@@ -471,16 +471,17 @@
 evice ID
+s
  to use'
@@ -765,16 +765,92 @@
 .parent%0A
+    if not args.gpus and not args.cpu_workers:%0A        args.cpu_workers = 1%0A
     cnn 
@@ -877,19 +877,50 @@
 T_BVLC, 
-gpu
+cpu_workers=args.cpu_workers, gpus
 =args.gp
@@ -920,16 +920,17 @@
 args.gpu
+s
 )%0A    in
@@ -1338,24 +1338,53 @@
 _tile_size)%0A
+    t2 = time.perf_counter()%0A
     print('I
@@ -1432,32 +1432,32 @@
 assify(img, 5):%0A
+
         print('%25
@@ -1473,37 +1473,8 @@
  c)%0A
-    t2 = time.perf_counter()%0A
     

Update deep_dream_test.py for multiprocess support