|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import functools
|
|
from .G1 import G1
|
|
from .G2 import G2
|
|
from .G3 import G3
|
|
from .zencoder import DeeperZencoder
|
|
|
|
|
|
class Generator(nn.Module):
|
|
def __init__(self, cfg):
|
|
super(Generator, self).__init__()
|
|
|
|
self.G1 = G1(cfg)
|
|
self.G2 = G2(cfg)
|
|
self.G3 = G3(cfg)
|
|
self.zencoder = DeeperZencoder(cfg)
|
|
self.up = nn.Upsample(scale_factor=2, mode='bilinear')
|
|
|
|
def forward(self, gt, input, segmap, inst_map, mask, cached_codes=None):
|
|
gt_G1 = F.interpolate(gt, size=(64, 64), mode='bilinear')
|
|
gt_G2 = F.interpolate(gt, size=(128, 128), mode='bilinear')
|
|
gt_G3 = F.interpolate(gt, size=(256, 256), mode='bilinear')
|
|
|
|
style_codes = self.zencoder(input, segmap, 1.0 - mask, inst_map, cached_codes=cached_codes)
|
|
|
|
input_G1 = F.interpolate(input, size=(64, 64), mode='bilinear')
|
|
fake_G1 = self.G1(input_G1, segmap, mask, style_codes)
|
|
mask_fake_G1 = self.masked_fake(gt_G1, fake_G1, mask)
|
|
input_G2 = self.next_img(gt_G2, fake_G1, mask)
|
|
|
|
fake_G2 = self.G2(input_G2, segmap, mask, style_codes)
|
|
mask_fake_G2 = self.masked_fake(gt_G2, fake_G2, mask)
|
|
input_G3 = self.next_img(gt_G3, fake_G2, mask)
|
|
|
|
fake_G3 = self.G3(input_G3, segmap, mask, style_codes)
|
|
mask_fake_G3 = self.masked_fake(gt_G3, fake_G3, mask)
|
|
|
|
return [gt_G1, gt_G2, gt_G3], [input_G1, input_G2, input_G3], [mask_fake_G1, mask_fake_G2, mask_fake_G3], [fake_G1, fake_G2, fake_G3]
|
|
|
|
def get_style_codes(self, input, segmap, inst_map, mask):
|
|
return self.zencoder.generate_style_codes(input, segmap, 1.0 - mask, inst_map)
|
|
|
|
def masked_fake(self, img, fake, mask):
|
|
mask = F.interpolate(mask, size=fake.size()[2:], mode='nearest')
|
|
combined = mask * fake + (1. - mask) * img
|
|
return combined
|
|
|
|
def next_img(self, img, prev_fake, mask):
|
|
fake = self.up(prev_fake)
|
|
mask = F.interpolate(mask, size=fake.size()[2:], mode='nearest')
|
|
combined = mask * fake + (1. - mask) * img
|
|
|
|
return combined |
