Advertisement
Guest User

Create Keras YOLOv3 from scratch

a guest
Jun 24th, 2018
221
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 10.18 KB | None | 0 0
  1. #! /usr/bin/env python
  2. """
  3. Reads modified Darknet53 config and creates Keras model with TF backend.
  4.  
  5. Currently only supports layers in Darknet53 config.
  6. """
  7.  
  8. import argparse
  9. import configparser
  10. import io
  11. import os
  12. from collections import defaultdict
  13.  
  14. import numpy as np
  15. from keras import backend as K
  16. from keras.layers import (Conv2D, GlobalAveragePooling2D, Input, Reshape,
  17. ZeroPadding2D, UpSampling2D, Activation)
  18. from keras.layers.advanced_activations import LeakyReLU
  19. from keras.layers.merge import concatenate, add
  20. from keras.layers.normalization import BatchNormalization
  21. from keras.models import Model
  22. from keras.regularizers import l2
  23. from keras.utils.vis_utils import plot_model as plot
  24.  
  25. parser = argparse.ArgumentParser(
  26. description='Yet Another Darknet To Keras Converter.')
  27. parser.add_argument('config_path', help='Path to Darknet cfg file.')
  28. parser.add_argument('output_path', help='Path to output Keras model file.')
  29. parser.add_argument(
  30. '-p',
  31. '--plot_model',
  32. help='Plot generated Keras model and save as image.',
  33. action='store_true')
  34. parser.add_argument(
  35. '-flcl',
  36. '--fully_convolutional',
  37. help='Model is fully convolutional so set input shape to (None, None, 3). '
  38. 'WARNING: This experimental option does not work properly for YOLO_v2.',
  39. action='store_true')
  40.  
  41.  
  42. def unique_config_sections(config_file):
  43. """Convert all config sections to have unique names.
  44.  
  45. Adds unique suffixes to config sections for compability with configparser.
  46. """
  47. section_counters = defaultdict(int)
  48. output_stream = io.StringIO()
  49. with open(config_file) as fin:
  50. for line in fin:
  51. if line.startswith('['):
  52. section = line.strip().strip('[]')
  53. _section = section + '_' + str(section_counters[section])
  54. section_counters[section] += 1
  55. line = line.replace(section, _section)
  56. output_stream.write(line)
  57. output_stream.seek(0)
  58. return output_stream
  59.  
  60.  
  61. def _main(args):
  62. config_path = os.path.expanduser(args.config_path)
  63. assert config_path.endswith('.cfg'), '{} is not a .cfg file'.format(
  64. config_path)
  65.  
  66. output_path = os.path.expanduser(args.output_path)
  67. assert output_path.endswith(
  68. '.h5'), 'output path {} is not a .h5 file'.format(output_path)
  69. output_root = os.path.splitext(output_path)[0]
  70.  
  71. # Load weights and config.
  72. # TODO: Check transpose flag when implementing fully connected layers.
  73. # transpose = (weight_header[0] > 1000) or (weight_header[1] > 1000)
  74.  
  75. print('Parsing Darknet config.')
  76. unique_config_file = unique_config_sections(config_path)
  77. cfg_parser = configparser.ConfigParser()
  78. cfg_parser.read_file(unique_config_file)
  79.  
  80. print('Creating Keras model.')
  81. if args.fully_convolutional:
  82. image_height, image_width = None, None
  83. else:
  84. image_height = int(cfg_parser['net_0']['height'])
  85. image_width = int(cfg_parser['net_0']['width'])
  86.  
  87. channels = int(cfg_parser['net_0']['channels'])
  88. prev_layer = Input(shape=(image_height, image_width, channels))
  89. all_layers = [prev_layer]
  90. outputs = []
  91.  
  92. weight_decay = float(cfg_parser['net_0']['decay']
  93. ) if 'net_0' in cfg_parser.sections() else 5e-4
  94. count = 0
  95.  
  96. for section in cfg_parser.sections():
  97. print('Parsing section {}'.format(section))
  98. if section.startswith('convolutional'):
  99. filters = int(cfg_parser[section]['filters'])
  100. size = int(cfg_parser[section]['size'])
  101. stride = int(cfg_parser[section]['stride'])
  102. pad = int(cfg_parser[section]['pad'])
  103. activation = cfg_parser[section]['activation']
  104. batch_normalize = 'batch_normalize' in cfg_parser[section]
  105.  
  106. # Setting weights.
  107. # Darknet serializes convolutional weights as:
  108. # [bias/beta, [gamma, mean, variance], conv_weights]
  109. prev_layer_shape = K.int_shape(prev_layer)
  110.  
  111. # TODO: This assumes channel last dim_ordering.
  112. weights_shape = (size, size, prev_layer_shape[-1], filters)
  113. darknet_w_shape = (filters, weights_shape[2], size, size)
  114. weights_size = np.product(weights_shape)
  115.  
  116. print('conv2d', 'bn'
  117. if batch_normalize else ' ', activation, weights_shape)
  118.  
  119. conv_bias = np.ndarray(
  120. shape=(filters, ),
  121. dtype='float32')
  122. count += filters
  123.  
  124. if batch_normalize:
  125. bn_weights = np.ndarray(
  126. shape=(3, filters),
  127. dtype='float32')
  128. count += 3 * filters
  129.  
  130. # TODO: Keras BatchNormalization mistakenly refers to var
  131. # as std.
  132. bn_weight_list = [
  133. bn_weights[0], # scale gamma
  134. conv_bias, # shift beta
  135. bn_weights[1], # running mean
  136. bn_weights[2] # running var
  137. ]
  138.  
  139. conv_weights = np.ndarray(
  140. shape=darknet_w_shape,
  141. dtype='float32')
  142. count += weights_size
  143.  
  144. # DarkNet conv_weights are serialized Caffe-style:
  145. # (out_dim, in_dim, height, width)
  146. # We would like to set these to Tensorflow order:
  147. # (height, width, in_dim, out_dim)
  148. # TODO: Add check for Theano dim ordering.
  149. conv_weights = np.transpose(conv_weights, [2, 3, 1, 0])
  150. conv_weights = [conv_weights] if batch_normalize else [
  151. conv_weights, conv_bias
  152. ]
  153.  
  154. # Handle activation.
  155. act_fn = None
  156. if activation == 'leaky':
  157. pass # Add advanced activation later.
  158. elif activation != 'linear':
  159. raise ValueError(
  160. 'Unknown activation function `{}` in section {}'.format(
  161. activation, section))
  162.  
  163. padding = 'same' if pad == 1 and stride == 1 else 'valid'
  164. # Adjust padding model for darknet.
  165. if stride == 2:
  166. prev_layer = ZeroPadding2D(((1, 0), (1, 0)))(prev_layer)
  167.  
  168. # Create Conv2D layer
  169. conv_layer = (Conv2D(
  170. filters, (size, size),
  171. strides=(stride, stride),
  172. kernel_regularizer=l2(weight_decay),
  173. use_bias=not batch_normalize,
  174. weights=conv_weights,
  175. activation=act_fn,
  176. padding=padding))(prev_layer)
  177.  
  178. if batch_normalize:
  179. conv_layer = (BatchNormalization(
  180. weights=bn_weight_list))(conv_layer)
  181.  
  182. prev_layer = conv_layer
  183.  
  184. if activation == 'linear':
  185. all_layers.append(prev_layer)
  186. elif activation == 'leaky':
  187. act_layer = LeakyReLU(alpha=0.1)(prev_layer)
  188. prev_layer = act_layer
  189. all_layers.append(act_layer)
  190.  
  191. elif section.startswith('avgpool'):
  192. if cfg_parser.items(section) != []:
  193. raise ValueError('{} with params unsupported.'.format(section))
  194. all_layers.append(GlobalAveragePooling2D()(prev_layer))
  195. prev_layer = all_layers[-1]
  196.  
  197. elif section.startswith('route'):
  198. ids = [int(i) for i in cfg_parser[section]['layers'].split(',')]
  199. layers = [all_layers[i] for i in ids]
  200.  
  201. if len(layers) > 1:
  202. print('Concatenating route layers:', layers)
  203. concatenate_layer = concatenate(layers)
  204. all_layers.append(concatenate_layer)
  205. prev_layer = concatenate_layer
  206. else:
  207. skip_layer = layers[0] # only one layer to route
  208. all_layers.append(skip_layer)
  209. prev_layer = skip_layer
  210.  
  211. elif section.startswith('shortcut'):
  212. ids = [int(i) for i in cfg_parser[section]['from'].split(',')][0]
  213. activation = cfg_parser[section]['activation']
  214. shortcut = add([all_layers[ids], prev_layer])
  215. if activation == 'linear':
  216. shortcut = Activation('linear')(shortcut)
  217. all_layers.append(shortcut)
  218. prev_layer = all_layers[-1]
  219.  
  220. elif section.startswith('upsample'):
  221. stride = int(cfg_parser[section]['stride'])
  222. all_layers.append(
  223. UpSampling2D(
  224. size=(stride, stride))(prev_layer))
  225. prev_layer = all_layers[-1]
  226.  
  227. elif section.startswith('yolo'):
  228. classes = int(cfg_parser[section]['classes'])
  229. # num = int(cfg_parser[section]['num'])
  230. # mask = int(cfg_parser[section]['mask'])
  231. n1, n2 = int(prev_layer.shape[1]), int(prev_layer.shape[2])
  232. n3 = 3
  233. n4 = (4 + 1 + classes)
  234. yolo = Reshape((n1, n2, n3, n4))(prev_layer)
  235. all_layers.append(yolo)
  236.  
  237. prev_layer = all_layers[-1]
  238. outputs.append(len(all_layers) - 1)
  239.  
  240. elif (section.startswith('net')):
  241. pass # Configs not currently handled during model definition.
  242. else:
  243. raise ValueError(
  244. 'Unsupported section header type: {}'.format(section))
  245.  
  246. # Create and save model.
  247. model = Model(inputs=all_layers[0],
  248. outputs=[all_layers[i] for i in outputs])
  249. print(model.summary())
  250. model.save('{}'.format(output_path))
  251. print('Saved Keras model to {}'.format(output_path))
  252. # Check to see if all weights have been read.
  253. remaining_weights = 0
  254. print('Read {} of {} from Darknet weights.'.format(count, count +
  255. remaining_weights))
  256. if remaining_weights > 0:
  257. print('Warning: {} unused weights'.format(remaining_weights))
  258.  
  259. if args.plot_model:
  260. plot(model, to_file='{}.png'.format(output_root), show_shapes=True)
  261. print('Saved model plot to {}.png'.format(output_root))
  262.  
  263.  
  264. if __name__ == '__main__':
  265. _main(parser.parse_args())
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement