import numpy as np import torch from torch import nn from torch.nn import functional as F from MT import FeatureTransformer from torch.cuda.amp import autocast as autocast from flow_tools import viz_img_seq, save_img_seq, plt_show_img_flow from copy import deepcopy from V1 import V1 import matplotlib.pyplot as plt from io import BytesIO from PIL import Image def conv(in_planes, out_planes, kernel_size=3, stride=1, dilation=1, isReLU=True): if isReLU: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True), nn.GELU() ) else: return nn.Sequential( nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size - 1) * dilation) // 2, bias=True) ) def plt_attention(attention, h, w): col = len(attention) // 2 fig = plt.figure(figsize=(10, 8)) for i in range(len(attention)): viz = attention[i][0, :, :, h, w].detach().cpu().numpy() # viz = viz[7:-7, 7:-7] if i == 0: viz_all = viz else: viz_all = viz_all + viz ax1 = fig.add_subplot(2, col, i + 1) img = ax1.imshow(viz, cmap="rainbow", interpolation="bilinear") ax1.scatter(w, h, color='grey', s=300, alpha=0.5) ax1.scatter(w, h, color='red', s=150, alpha=0.5) plt.title(" Iteration %d" % (i + 1)) if i == len(attention) - 1: plt.title(" Final Iteration") plt.xticks([]) plt.yticks([]) # tight layout plt.tight_layout() # save the figure buf = BytesIO() plt.savefig(buf, format='png') buf.seek(0) plt.close() # convert the figure to an array img = Image.open(buf) img = np.array(img) return img class FlowDecoder(nn.Module): # can reduce 25% of training time. def __init__(self, ch_in): super(FlowDecoder, self).__init__() self.conv1 = conv(ch_in, 256, kernel_size=1) self.conv2 = conv(256, 128, kernel_size=1) self.conv3 = conv(256 + 128, 96, kernel_size=1) self.conv4 = conv(96 + 128, 64, kernel_size=1) self.conv5 = conv(96 + 64, 32, kernel_size=1) self.feat_dim = 32 self.predict_flow = conv(64 + 32, 2, isReLU=False) def forward(self, x): x1 = self.conv1(x) x2 = self.conv2(x1) x3 = self.conv3(torch.cat([x1, x2], dim=1)) x4 = self.conv4(torch.cat([x2, x3], dim=1)) x5 = self.conv5(torch.cat([x3, x4], dim=1)) flow = self.predict_flow(torch.cat([x4, x5], dim=1)) return flow class FFV1DNN(nn.Module): def __init__(self, num_scales=8, num_cells=256, upsample_factor=8, feature_channels=256, scale_factor=16, num_layers=6, ): super(FFV1DNN, self).__init__() self.ffv1 = V1(spatial_num=num_cells // num_scales, scale_num=num_scales, scale_factor=scale_factor, kernel_radius=7, num_ft=num_cells // num_scales, kernel_size=6, average_time=True) self.v1_kz = 7 self.scale_factor = scale_factor scale_each_level = np.exp(1 / (num_scales - 1) * np.log(1 / scale_factor)) self.scale_num = num_scales self.scale_each_level = scale_each_level v1_channel = self.ffv1.num_after_st self.num_scales = num_scales self.MT_channel = feature_channels assert self.MT_channel == v1_channel self.feature_channels = feature_channels self.upsample_factor = upsample_factor self.num_layers = num_layers # convex upsampling: concat feature0 and flow as input self.upsampler_1 = nn.Sequential(nn.Conv2d(2 + feature_channels, 256, 3, 1, 1), nn.ReLU(inplace=True), nn.Conv2d(256, 256, 3, 1, 1), nn.ReLU(inplace=True), nn.Conv2d(256, upsample_factor ** 2 * 9, 3, 1, 1)) self.decoder = FlowDecoder(feature_channels) self.conv_feat = nn.ModuleList([conv(v1_channel, feature_channels, 1) for i in range(num_scales)]) self.MT = FeatureTransformer(d_model=feature_channels, num_layers=self.num_layers) # 2*2*8*scale` def upsample_flow(self, flow, feature, upsampler=None, bilinear=False, upsample_factor=4): if bilinear: up_flow = F.interpolate(flow, scale_factor=upsample_factor, mode='bilinear', align_corners=True) * upsample_factor else: # convex upsampling concat = torch.cat((flow, feature), dim=1) mask = upsampler(concat) b, flow_channel, h, w = flow.shape mask = mask.view(b, 1, 9, upsample_factor, upsample_factor, h, w) # [B, 1, 9, K, K, H, W] mask = torch.softmax(mask, dim=2) up_flow = F.unfold(upsample_factor * flow, [3, 3], padding=1) up_flow = up_flow.view(b, flow_channel, 9, 1, 1, h, w) # [B, 2, 9, 1, 1, H, W] up_flow = torch.sum(mask * up_flow, dim=2) # [B, 2, K, K, H, W] up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) # [B, 2, K, H, K, W] up_flow = up_flow.reshape(b, flow_channel, upsample_factor * h, upsample_factor * w) # [B, 2, K*H, K*W] return up_flow def forward(self, image_list, mix_enable=True, layer=6): if layer is not None: self.MT.num_layers = layer self.num_layers = layer results_dict = {} padding = self.v1_kz * self.scale_factor with torch.no_grad(): if image_list[0].max() > 10: image_list = [img / 255.0 for img in image_list] # [B, 1, H, W] 0-1 if image_list[0].shape[1] == 3: # convert to gray using transform Gray = R*0.299 + G*0.587 + B*0.114 image_list = [img[:, 0, :, :] * 0.299 + img[:, 1, :, :] * 0.587 + img[:, 2, :, :] * 0.114 for img in image_list] image_list = [img.unsqueeze(1) for img in image_list] B, _, H, W = image_list[0].shape MT_size = (H // 8, W // 8) with autocast(enabled=mix_enable): # with torch.no_grad(): # TODO: only for test wheather a trainable V1 is needed. st_component = self.ffv1(image_list) # viz_img_seq(image_scale, if_debug=True) if self.num_layers == 0: motion_feature = [st_component] flows = [self.decoder(feature) for feature in motion_feature] flows_up = [self.upsample_flow(flow, feature=None, bilinear=True, upsample_factor=8) for flow in flows] results_dict["flow_seq"] = flows_up return results_dict motion_feature, attn = self.MT.forward_save_mem(st_component) flow_v1 = self.decoder(st_component) flows = [flow_v1] + [self.decoder(feature) for feature in motion_feature] flows_bi = [self.upsample_flow(flow, feature=None, bilinear=True, upsample_factor=8) for flow in flows] flows_up = [flows_bi[0]] + \ [self.upsample_flow(flows, upsampler=self.upsampler_1, feature=attn, upsample_factor=8) for flows, attn in zip(flows[1:], attn)] assert len(flows_bi) == len(flows_up) results_dict["flow_seq"] = flows_up results_dict["flow_seq_bi"] = flows_bi return results_dict def forward_test(self, image_list, mix_enable=True, layer=6): if layer is not None: self.MT.num_layers = layer self.num_layers = layer results_dict = {} padding = self.v1_kz * self.scale_factor with torch.no_grad(): if image_list[0].max() > 10: image_list = [img / 255.0 for img in image_list] # [B, 1, H, W] 0-1 B, _, H, W = image_list[0].shape MT_size = (H // 8, W // 8) with autocast(enabled=mix_enable): st_component = self.ffv1(image_list) # viz_img_seq(image_scale, if_debug=True) if self.num_layers == 0: motion_feature = [st_component] flows = [self.decoder(feature) for feature in motion_feature] flows_up = [self.upsample_flow(flow, feature=None, bilinear=True, upsample_factor=8) for flow in flows] results_dict["flow_seq"] = flows_up return results_dict motion_feature, attn, _ = self.MT.forward_save_mem(st_component) flow_v1 = self.decoder(st_component) flows = [flow_v1] + [self.decoder(feature) for feature in motion_feature] flows_bi = [self.upsample_flow(flow, feature=None, bilinear=True, upsample_factor=8) for flow in flows] flows_up = [flows_bi[0]] + \ [self.upsample_flow(flows, upsampler=self.upsampler_1, feature=attn, upsample_factor=8) for flows, attn in zip(flows[1:], attn)] assert len(flows_bi) == len(flows_up) results_dict["flow_seq"] = flows_up results_dict["flow_seq_bi"] = flows_bi return results_dict def forward_viz(self, image_list, layer=None, x=50, y=50): x = x / 100 y = y / 100 if layer is not None: self.MT.num_layers = layer results_dict = {} padding = self.v1_kz * self.scale_factor with torch.no_grad(): if image_list[0].max() > 10: image_list = [img / 255.0 for img in image_list] # [B, 1, H, W] 0-1 if image_list[0].shape[1] == 3: # convert to gray using transform Gray = R*0.299 + G*0.587 + B*0.114 image_list = [img[:, 0, :, :] * 0.299 + img[:, 1, :, :] * 0.587 + img[:, 2, :, :] * 0.114 for img in image_list] image_list = [img.unsqueeze(1) for img in image_list] image_list_ori = deepcopy(image_list) B, _, H, W = image_list[0].shape MT_size = (H // 8, W // 8) with autocast(enabled=True): st_component = self.ffv1(image_list) activation = self.ffv1.visualize_activation(st_component) # viz_img_seq(image_scale, if_debug=True) motion_feature, attn, attn_viz = self.MT(st_component) flow_v1 = self.decoder(st_component) flows = [flow_v1] + [self.decoder(feature) for feature in motion_feature] flows_bi = [self.upsample_flow(flow, feature=None, bilinear=True, upsample_factor=8) for flow in flows] flows_up = [flows_bi[0]] + \ [self.upsample_flow(flows, upsampler=self.upsampler_1, feature=attn, upsample_factor=8) for flows, attn in zip(flows[1:], attn)] assert len(flows_bi) == len(flows_up) results_dict["flow_seq"] = flows_up # select 1,3,5,7 flows_up = [flows_up[i] for i in [0, 2, 4]] + [flows_up[-1]] attn_viz = [attn_viz[i] for i in [0, 2, 4]] + [attn_viz[-1]] flow = plt_show_img_flow(image_list_ori, flows_up) h = int(MT_size[0] * y) w = int(MT_size[1] * x) attention = plt_attention(attn_viz, h=h, w=w) print("done") results_dict["activation"] = activation results_dict["attention"] = attention results_dict["flow"] = flow return results_dict def num_parameters(self): return sum( [p.data.nelement() if p.requires_grad else 0 for p in self.parameters()]) def init_weights(self): for layer in self.named_modules(): if isinstance(layer, nn.Conv2d): nn.init.kaiming_normal_(layer.weight) if layer.bias is not None: nn.init.constant_(layer.bias, 0) if isinstance(layer, nn.Conv1d): nn.init.kaiming_normal_(layer.weight) if layer.bias is not None: nn.init.constant_(layer.bias, 0) elif isinstance(layer, nn.ConvTranspose2d): nn.init.kaiming_normal_(layer.weight) if layer.bias is not None: nn.init.constant_(layer.bias, 0) @staticmethod def demo(file=None): import time from utils import torch_utils as utils frame_list = [torch.randn([4, 1, 512, 512], device="cuda")] * 11 model = FFV1DNN(num_scales=8, scale_factor=16, num_cells=256, upsample_factor=8, num_layers=6, feature_channels=256).cuda() if file is not None: model = utils.restore_model(model, file) print(model.num_parameters()) for i in range(100): start = time.time() output = model.forward_viz(frame_list, layer=7) # print(output["flow_seq"][-1]) torch.mean(output["flow_seq"][-1]).backward() print(torch.any(torch.isnan(output["flow_seq"][-1]))) end = time.time() print(end - start) print("#================================++#") if __name__ == '__main__': FFV1DNN.demo(None)