File size: 3,485 Bytes
439c523 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from .layers_utils import warp, make_grid
from .generative_network import Generative_Encoder, Generative_Decoder
from .evolution_network import Evolution_Network
from .noise_projector import Noise_Projector
class Net(nn.Module):
def __init__(self, configs):
super(Net, self).__init__()
self.configs = configs
self.pred_length = self.configs.total_length - self.configs.input_length
self.evo_net = Evolution_Network(self.configs.input_length, self.pred_length, base_c=32)
self.gen_enc = Generative_Encoder(self.configs.total_length, base_c=self.configs.ngf)
self.gen_dec = Generative_Decoder(self.configs)
self.proj = Noise_Projector(self.configs.ngf, configs)
sample_tensor = torch.zeros(1, 1, self.configs.img_height, self.configs.img_width)
self.register_buffer("grid", make_grid(sample_tensor), persistent=False)
def forward(self, all_frames):
if all_frames.ndim == 4:
all_frames = all_frames.unsqueeze(-1)
if all_frames.ndim != 5 or all_frames.shape[-1] < 1:
raise ValueError("Expected frames with shape [B,T,H,W,C], C>=1")
all_frames = all_frames[:, :, :, :, :1]
frames = all_frames.permute(0, 1, 4, 2, 3)
batch = frames.shape[0]
height = frames.shape[3]
width = frames.shape[4]
# Input Frames
input_frames = frames[:, :self.configs.input_length]
input_frames = input_frames.reshape(batch, self.configs.input_length, height, width)
# Evolution Network
intensity, motion = self.evo_net(input_frames)
motion_ = motion.reshape(batch, self.pred_length, 2, height, width)
intensity_ = intensity.reshape(batch, self.pred_length, 1, height, width)
series = []
last_frames = all_frames[:, (self.configs.input_length - 1):self.configs.input_length, :, :, 0]
grid = self.grid.to(frames.device).repeat(batch, 1, 1, 1)
for i in range(self.pred_length):
last_frames = warp(last_frames, motion_[:, i], grid, mode="nearest", padding_mode="border")
last_frames = last_frames + intensity_[:, i]
series.append(last_frames)
evo_result = torch.cat(series, dim=1)
evo_result = evo_result/128
# Generative Network
evo_feature = self.gen_enc(torch.cat([input_frames, evo_result], dim=1))
noise = torch.randn(batch, self.configs.ngf, max(1, height // 32), max(1, width // 32), device=frames.device)
projected = self.proj(noise)
# The official reshape/permute is a 4x pixel shuffle: 32*ngf
# projected channels become 2*ngf channels at H/8 x W/8. Expressing
# it directly preserves the checkpoint layout and also works at the
# official 512x512 resolution (where the released literal reshape is
# dimensionally inconsistent).
if projected.shape[2] * 4 == height // 8 and projected.shape[3] * 4 == width // 8:
noise_feature = F.pixel_shuffle(projected, 4)
else:
target_hw = (max(1, height // 8), max(1, width // 8))
noise_feature = F.interpolate(projected[:, : 2 * self.configs.ngf], size=target_hw, mode="nearest")
feature = torch.cat([evo_feature, noise_feature], dim=1)
gen_result = self.gen_dec(feature, evo_result)
return gen_result.unsqueeze(-1)
|