Single Hourglass Module

hg.py

def HourglassModule(inputs, order, filters, num_residual):
    """
    One Hourglass Module. Usually we stacked multiple of them together.
    https://github.com/princeton-vl/pose-hg-train/blob/master/src/models/hg.lua#L3

    inputs:
    order: The remaining order for HG modules to call itself recursively.
    num_residual: Number of residual layers for this HG module.
    """
    # Upper branch
    up1 = BottleneckBlock(inputs, filters, downsample=False)

    for i in range(num_residual):
        up1 = BottleneckBlock(up1, filters, downsample=False)

    # Lower branch
    low1 = MaxPool2D(pool_size=2, strides=2)(inputs)
    for i in range(num_residual):
        low1 = BottleneckBlock(low1, filters, downsample=False)

    low2 = low1
    if order > 1:
        low2 = HourglassModule(low1, order - 1, filters, num_residual)
    else:
        for i in range(num_residual):
            low2 = BottleneckBlock(low2, filters, downsample=False)

    low3 = low2
    for i in range(num_residual):
        low3 = BottleneckBlock(low3, filters, downsample=False)

    up2 = UpSampling2D(size=2)(low3)

    return up2 + up1