Update net.py

net.py

@@ -31,36 +31,22 @@ def data_loader_bd_rm_from_tfrecord(batch_size=1):
 	return loader_dict, num_batch
 
 class Network(object):
-	"""docstring for Network"""
-	def __init__(self, dtype=tf.float32):
-		print('Initial nn network object...')
-		self.dtype = dtype
-		self.pre_train_restore_map = {'vgg_16/conv1/conv1_1/weights':'FNet/conv1_1/W', # {'checkpoint_scope_var_name':'current_scope_var_name'} shape must be the same
-									'vgg_16/conv1/conv1_1/biases':'FNet/conv1_1/b',	
-									'vgg_16/conv1/conv1_2/weights':'FNet/conv1_2/W',
-									'vgg_16/conv1/conv1_2/biases':'FNet/conv1_2/b',	
-									'vgg_16/conv2/conv2_1/weights':'FNet/conv2_1/W',
-									'vgg_16/conv2/conv2_1/biases':'FNet/conv2_1/b',	
-									'vgg_16/conv2/conv2_2/weights':'FNet/conv2_2/W',
-									'vgg_16/conv2/conv2_2/biases':'FNet/conv2_2/b',	
-									'vgg_16/conv3/conv3_1/weights':'FNet/conv3_1/W',
-									'vgg_16/conv3/conv3_1/biases':'FNet/conv3_1/b',	
-									'vgg_16/conv3/conv3_2/weights':'FNet/conv3_2/W',
-									'vgg_16/conv3/conv3_2/biases':'FNet/conv3_2/b',	
-									'vgg_16/conv3/conv3_3/weights':'FNet/conv3_3/W',
-									'vgg_16/conv3/conv3_3/biases':'FNet/conv3_3/b',	
-									'vgg_16/conv4/conv4_1/weights':'FNet/conv4_1/W',
-									'vgg_16/conv4/conv4_1/biases':'FNet/conv4_1/b',	
-									'vgg_16/conv4/conv4_2/weights':'FNet/conv4_2/W',
-									'vgg_16/conv4/conv4_2/biases':'FNet/conv4_2/b',	
-									'vgg_16/conv4/conv4_3/weights':'FNet/conv4_3/W',
-									'vgg_16/conv4/conv4_3/biases':'FNet/conv4_3/b',	
-									'vgg_16/conv5/conv5_1/weights':'FNet/conv5_1/W',
-									'vgg_16/conv5/conv5_1/biases':'FNet/conv5_1/b',	
-									'vgg_16/conv5/conv5_2/weights':'FNet/conv5_2/W',
-									'vgg_16/conv5/conv5_2/biases':'FNet/conv5_2/b',	
-									'vgg_16/conv5/conv5_3/weights':'FNet/conv5_3/W',
-									'vgg_16/conv5/conv5_3/biases':'FNet/conv5_3/b'} 
+    def __init__(self, dtype=tf.float32):
+        print('Initial nn network object...')
+        self.dtype = dtype
+        # ... existing code ...
+
+    def convert_one_hot_to_image(self, one_hot, dtype='float', act=None):
+        import tensorflow.compat.v1 as tf
+        if act == 'softmax':
+            one_hot = tf.nn.softmax(one_hot, axis=-1)
+        [n, h, w, c] = one_hot.shape.as_list()
+        im = tf.reshape(tf.argmax(one_hot, axis=-1), [n, h, w, 1])
+        if dtype == 'int':
+            im = tf.cast(im, dtype=tf.uint8)
+        else:
+            im = tf.cast(im, dtype=tf.float32)
+        return im
 
 	# basic layer 
 	def _he_uniform(self, shape, regularizer=None, trainable=None, name=None):