|
- #! /usr/bin/env python
- # -*- coding: utf-8 -*-
- import os
- import sys
- import cv2
- import time
- import numpy as np
- import matplotlib.pyplot as plt
- import core.utils as utils
- import tensorflow as tf
- from PIL import Image
-
-
- if __name__ == '__main__':
- argv = sys.argv
- if len(argv) < 5:
- print('usage: python show_layer_feature_map.py gpu_id pb_file img_file out_path')
- sys.exit()
-
- gpu_id = argv[1]
- os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
-
- pb_file = argv[2]
- if not os.path.exists(pb_file):
- print('pb_file=%s not exist' % pb_file)
- sys.exit()
-
- img_file = argv[3]
- if not os.path.exists(img_file):
- print('img_file=%s not exist' % img_file)
- sys.exit()
-
- out_path = argv[4]
- if not os.path.exists(out_path):
- os.makedirs(out_path)
- print('show_layer_feature_map gpu_id=%s, pb_file=%s, img_file=%s, out_path=%s' %
- (gpu_id, pb_file, img_file, out_path))
-
- input_size = 416
- img = cv2.imread(img_file)
- image_data = utils.image_preporcess(np.copy(img), [input_size, input_size])
- image_data = image_data[np.newaxis, ...]
-
- graph = tf.Graph()
- return_elements = ['input/input_data:0', 'pred_sbbox/concat_2:0', 'pred_mbbox/concat_2:0', 'pred_lbbox/concat_2:0']
- return_tensors = utils.read_pb_return_tensors(graph, pb_file, return_elements)
-
- with tf.Session(graph=graph) as sess:
- tensor_names = [tensor.name for tensor in tf.get_default_graph().as_graph_def().node]
- conv_layer_names = []
- for idx, tensor_name in enumerate(tensor_names):
- if 'Conv2D' in tensor_name:
- conv_layer_names.append(tensor_name)
- print('conv_layer_names=', conv_layer_names)
-
- for idx, layer_name in enumerate(conv_layer_names):
- conv = sess.graph.get_tensor_by_name('%s:0' % layer_name)
- features = np.array(conv.eval({return_tensors[0]: image_data}))
- print('\n[%d/%d] %s' % (idx, len(conv_layer_names), layer_name), ' features.shape=', features.shape)
-
- out_layer_path = os.path.join(out_path, '%s-%sx%sx%s' % (layer_name.replace('/', '_'), str(features.shape[1]),
- str(features.shape[2]), str(features.shape[3])))
- if not os.path.exists(out_layer_path):
- os.makedirs(out_layer_path)
-
- plt.figure(idx, figsize=(10, 10))
- for jdx in range(features.shape[3]):
- plt.matshow(features[0, :, :, jdx], cmap=plt.cm.gray, fignum=idx) #remove cmap=plt.cm.gray to show RGBA image
- plt.title('' + layer_name + '_' + str(jdx))
-
- out_file = os.path.join(out_layer_path, 'img_%s.jpg' % str(jdx))
- plt.savefig(out_file)
- print('idx=', idx, ' layer_name=', layer_name, ' jdx=', jdx, ' out_file=', out_file)
|