| 376 | |
| 377 | |
| 378 | class FlowNet2C(FlowNetBase): |
| 379 | def graph_structure(self, x1x2): |
| 380 | """ |
| 381 | Architecture of FlowNetCorr in Figure 2 of FlowNet 1.0. |
| 382 | Args: |
| 383 | x: 2CHW. |
| 384 | """ |
| 385 | with argscope([tf.layers.conv2d], activation=lambda x: tf.nn.leaky_relu(x, 0.1), |
| 386 | padding='valid', strides=2, kernel_size=3, |
| 387 | data_format='channels_first'), \ |
| 388 | argscope([tf.layers.conv2d_transpose], padding='same', activation=tf.identity, |
| 389 | data_format='channels_first', strides=2, kernel_size=4): |
| 390 | |
| 391 | # extract features |
| 392 | x = tf.layers.conv2d(pad(x1x2, 3), 64, kernel_size=7, name='conv1') |
| 393 | conv2 = tf.layers.conv2d(pad(x, 2), 128, kernel_size=5, name='conv2') |
| 394 | conv3 = tf.layers.conv2d(pad(conv2, 2), 256, kernel_size=5, name='conv3') |
| 395 | |
| 396 | conv2a, _ = tf.split(conv2, 2, axis=0) |
| 397 | conv3a, conv3b = tf.split(conv3, 2, axis=0) |
| 398 | |
| 399 | corr = correlation(conv3a, conv3b, |
| 400 | kernel_size=1, |
| 401 | max_displacement=20, |
| 402 | stride_1=1, |
| 403 | stride_2=2, |
| 404 | pad=20, data_format='NCHW') |
| 405 | corr = tf.nn.leaky_relu(corr, 0.1) |
| 406 | |
| 407 | conv_redir = tf.layers.conv2d(conv3a, 32, kernel_size=1, strides=1, name='conv_redir') |
| 408 | |
| 409 | in_conv3_1 = tf.concat([conv_redir, corr], axis=1, name='in_conv3_1') |
| 410 | conv3_1 = tf.layers.conv2d(pad(in_conv3_1, 1), 256, name='conv3_1', strides=1) |
| 411 | |
| 412 | x = tf.layers.conv2d(pad(conv3_1, 1), 512, name='conv4') |
| 413 | conv4 = tf.layers.conv2d(pad(x, 1), 512, name='conv4_1', strides=1) |
| 414 | x = tf.layers.conv2d(pad(conv4, 1), 512, name='conv5') |
| 415 | conv5 = tf.layers.conv2d(pad(x, 1), 512, name='conv5_1', strides=1) |
| 416 | x = tf.layers.conv2d(pad(conv5, 1), 1024, name='conv6') |
| 417 | conv6 = tf.layers.conv2d(pad(x, 1), 1024, name='conv6_1', strides=1) |
| 418 | |
| 419 | flow6 = tf.layers.conv2d(pad(conv6, 1), 2, name='predict_flow6', strides=1, activation=tf.identity) |
| 420 | flow6_up = tf.layers.conv2d_transpose(flow6, 2, name='upsampled_flow6_to_5') |
| 421 | x = tf.layers.conv2d_transpose(conv6, 512, name='deconv5', activation=lambda x: tf.nn.leaky_relu(x, 0.1)) |
| 422 | |
| 423 | # return flow6 |
| 424 | concat5 = tf.concat([conv5, x, flow6_up], axis=1, name='concat5') |
| 425 | flow5 = tf.layers.conv2d(pad(concat5, 1), 2, name='predict_flow5', strides=1, activation=tf.identity) |
| 426 | flow5_up = tf.layers.conv2d_transpose(flow5, 2, name='upsampled_flow5_to_4') |
| 427 | x = tf.layers.conv2d_transpose(concat5, 256, name='deconv4', activation=lambda x: tf.nn.leaky_relu(x, 0.1)) |
| 428 | |
| 429 | concat4 = tf.concat([conv4, x, flow5_up], axis=1, name='concat4') |
| 430 | flow4 = tf.layers.conv2d(pad(concat4, 1), 2, name='predict_flow4', strides=1, activation=tf.identity) |
| 431 | flow4_up = tf.layers.conv2d_transpose(flow4, 2, name='upsampled_flow4_to_3') |
| 432 | x = tf.layers.conv2d_transpose(concat4, 128, name='deconv3', activation=lambda x: tf.nn.leaky_relu(x, 0.1)) |
| 433 | |
| 434 | concat3 = tf.concat([conv3_1, x, flow4_up], axis=1, name='concat3') |
| 435 | flow3 = tf.layers.conv2d(pad(concat3, 1), 2, name='predict_flow3', strides=1, activation=tf.identity) |