texture_nets.py

import torch
import torch.nn as nn
from .common import * 


normalization = nn.BatchNorm2d


def conv(in_f, out_f, kernel_size, stride=1, bias=True, pad='zero'):
    if pad == 'zero':
        return nn.Conv2d(in_f, out_f, kernel_size, stride, padding=(kernel_size - 1) / 2, bias=bias)
    elif pad == 'reflection':
        layers = [nn.ReflectionPad2d((kernel_size - 1) / 2),
                  nn.Conv2d(in_f, out_f, kernel_size, stride, padding=0, bias=bias)]
        return nn.Sequential(*layers)

def get_texture_nets(inp=3, ratios = [32, 16, 8, 4, 2, 1], fill_noise=False, pad='zero', need_sigmoid=False, conv_num=8, upsample_mode='nearest'):


    for i in range(len(ratios)):
        j = i + 1

        seq = nn.Sequential()

        tmp =  nn.AvgPool2d(ratios[i], ratios[i])

        seq.add(tmp)
        if fill_noise:
            seq.add(GenNoise(inp))

        seq.add(conv(inp, conv_num, 3, pad=pad))
        seq.add(normalization(conv_num))
        seq.add(act())

        seq.add(conv(conv_num, conv_num, 3, pad=pad))
        seq.add(normalization(conv_num))
        seq.add(act())

        seq.add(conv(conv_num, conv_num, 1, pad=pad))
        seq.add(normalization(conv_num))
        seq.add(act())

        if i == 0:
            seq.add(nn.Upsample(scale_factor=2, mode=upsample_mode))
            cur = seq
        else:

            cur_temp = cur

            cur = nn.Sequential()

            # Batch norm before merging 
            seq.add(normalization(conv_num))
            cur_temp.add(normalization(conv_num * (j - 1)))

            cur.add(Concat(1, cur_temp, seq))

            cur.add(conv(conv_num * j, conv_num * j, 3, pad=pad))
            cur.add(normalization(conv_num * j))
            cur.add(act())

            cur.add(conv(conv_num * j, conv_num * j, 3, pad=pad))
            cur.add(normalization(conv_num * j))
            cur.add(act())

            cur.add(conv(conv_num * j, conv_num * j, 1, pad=pad))
            cur.add(normalization(conv_num * j))
            cur.add(act())

            if i == len(ratios) - 1: 
                cur.add(conv(conv_num * j, 3, 1, pad=pad))
            else:
                cur.add(nn.Upsample(scale_factor=2, mode=upsample_mode)) 
            
    model = cur
    if need_sigmoid:
        model.add(nn.Sigmoid())

    return model