diff --git a/deepdreamer/deepdreamer.py b/deepdreamer/deepdreamer.py index e147d02..2522937 100644 --- a/deepdreamer/deepdreamer.py +++ b/deepdreamer/deepdreamer.py @@ -23,20 +23,20 @@ def _select_network(netname): if netname == 'bvlc_googlenet': - NET_FN = "deploy.prototxt" # Make sure force_backward: true - PARAM_FN = "bvlc_googlenet.caffemodel" - CHANNEL_SWAP = (2, 1, 0) + net_fn = "deploy.prototxt" # Make sure force_backward: true + param_fn = "bvlc_googlenet.caffemodel" + channel_swap = (2, 1, 0) # ImageNet mean, training set dependent - CAFFE_MEAN = np.float32([104.0, 116.0, 122.0]) - return NET_FN, PARAM_FN, CHANNEL_SWAP, CAFFE_MEAN + caffe_mean = np.float32([104.0, 116.0, 122.0]) + return net_fn, param_fn, channel_swap, caffe_mean elif netname == 'googlenet_place205': # TODO: refit SWAP and MEAN for places205? These work for now. - NET_FN = "deploy_places205.protxt" # Make sure force_backward: true - PARAM_FN = "googlelet_places205_train_iter_2400000.caffemodel" - CHANNEL_SWAP = (2, 1, 0) + net_fn = "deploy_places205.protxt" # Make sure force_backward: true + param_fn = "googlelet_places205_train_iter_2400000.caffemodel" + channel_swap = (2, 1, 0) # ImageNet mean, training set dependent - CAFFE_MEAN = np.float32([104.0, 116.0, 122.0]) - return NET_FN, PARAM_FN, CHANNEL_SWAP, CAFFE_MEAN + caffe_mean = np.float32([104.0, 116.0, 122.0]) + return net_fn, param_fn, channel_swap, caffe_mean else: print("Error: network {} not implemented".format(netname)) @@ -138,9 +138,9 @@ def _create_video(video, frame_rate=24): def list_layers(network="bvlc_googlenet"): # Load DNN model - NET_FN, PARAM_FN, CHANNEL_SWAP, CAFFE_MEAN = _select_network(network) + net_fn, param_fn, channel_swap, caffe_mean = _select_network(network) net = Classifier( - NET_FN, PARAM_FN, mean=CAFFE_MEAN, channel_swap=CHANNEL_SWAP) + net_fn, param_fn, mean=caffe_mean, channel_swap=channel_swap) net.blobs.keys() @@ -159,9 +159,9 @@ def deepdream( set_mode_gpu() # Select, load DNN model - NET_FN, PARAM_FN, CHANNEL_SWAP, CAFFE_MEAN = _select_network(network) + net_fn, param_fn, channel_swap, caffe_mean = _select_network(network) net = Classifier( - NET_FN, PARAM_FN, mean=CAFFE_MEAN, channel_swap=CHANNEL_SWAP) + net_fn, param_fn, mean=caffe_mean, channel_swap=channel_swap) img_pool = [img_path] @@ -205,9 +205,9 @@ def deepdream_video( frame_rate=24): # Select, load DNN model - NET_FN, PARAM_FN, CHANNEL_SWAP, CAFFE_MEAN = _select_network(network) + net_fn, param_fn, channel_swap, caffe_mean = _select_network(network) net = Classifier( - NET_FN, PARAM_FN, mean=CAFFE_MEAN, channel_swap=CHANNEL_SWAP) + net_fn, param_fn, mean=caffe_mean, channel_swap=channel_swap) print("Extracting video...") _extract_video(video)