LearnablePoolingMethods/model_utils.py at master · pomonam/LearnablePoolingMethods · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
# Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS-IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Contains a collection of util functions for model construction.
"""
# noinspection PyUnresolvedReferences
import pathmagic
import numpy
import tensorflow as tf
from tensorflow import logging
from tensorflow import flags
import tensorflow.contrib.slim as slim


def SampleRandomSequence(model_input, num_frames, num_samples):
    """ Samples a random sequence of frames of size num_samples.

    Args:
        model_input: A tensor of size batch_size x max_frames x feature_size
        num_frames: A tensor of size batch_size x 1
        num_samples: A scalar

    Returns:
        `model_input`: A tensor of size batch_size x num_samples x feature_size
      """

    batch_size = tf.shape(model_input)[0]
    frame_index_offset = tf.tile(
        tf.expand_dims(tf.range(num_samples), 0), [batch_size, 1])
    max_start_frame_index = tf.maximum(num_frames - num_samples, 0)
    start_frame_index = tf.cast(
        tf.multiply(
            tf.random_uniform([batch_size, 1]),
            tf.cast(max_start_frame_index + 1, tf.float32)), tf.int32)
    frame_index = tf.minimum(start_frame_index + frame_index_offset,
                             tf.cast(num_frames - 1, tf.int32))
    batch_index = tf.tile(
        tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
    index = tf.stack([batch_index, frame_index], 2)
    return tf.gather_nd(model_input, index)


def SampleRandomFrames(model_input, num_frames, num_samples):
    """ Samples a random set of frames of size num_samples.

      Args:
        model_input: A tensor of size batch_size x max_frames x feature_size
        num_frames: A tensor of size batch_size x 1
        num_samples: A scalar

      Returns:
        `model_input`: A tensor of size batch_size x num_samples x feature_size
      """
    batch_size = tf.shape(model_input)[0]
    frame_index = tf.cast(
        tf.multiply(
            tf.random_uniform([batch_size, num_samples]),
            tf.tile(tf.cast(num_frames, tf.float32), [1, num_samples])), tf.int32)
    batch_index = tf.tile(
        tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
    index = tf.stack([batch_index, frame_index], 2)
    return tf.gather_nd(model_input, index)


def FramePooling(frames, method, **unused_params):
    """Pools over the frames of a video.

    Args:
    frames: A tensor.
    method: "average", "max", "attention", or "none".
    Returns:
    A tensor with shape [batch_size, feature_size] for average, max, or
    attention pooling. A tensor with shape [batch_size*num_frames, feature_size]
    for none pooling.

    Raises:
    ValueError: if method is other than "average", "max", "attention", or
    "none".
    """
    if method == "average":
        return tf.reduce_mean(frames, 1)
    elif method == "max":
        return tf.reduce_max(frames, 1)
    elif method == "none":
        feature_size = frames.shape_as_list()[2]
        return tf.reshape(frames, [-1, feature_size])
    else:
        raise ValueError("Unrecognized pooling method: %s" % method)

def SampleUniformFrames(model_input, num_frames, num_samples):
    """ Uniformally samples (deterministically) a set of frames of size num_samples.

          Args:
            model_input: A tensor of size batch_size x max_frames x feature_size
            num_frames: A tensor of size batch_size x 1
            num_samples: A scalar

          Returns:
            `model_input`: A tensor of size batch_size x num_samples x feature_size
          """
    batch_size      = tf.shape(model_input)[0]
    even_dist_samp  = tf.expand_dims(tf.linspace(0.0, 1.0, num_samples+1), axis=0)
    even_dist_samp  = tf.slice(even_dist_samp, [0, 0], [1, num_samples])
    frame_index = tf.cast(
        tf.multiply(
            tf.tile(even_dist_samp, [batch_size, 1]),
            tf.tile(tf.cast(num_frames, tf.float32), [1, num_samples])), tf.int32)
    batch_index = tf.tile(
        tf.expand_dims(tf.range(batch_size), 1), [1, num_samples])
    index = tf.stack([batch_index, frame_index], 2)
    return tf.gather_nd(model_input, index)