import torch
import torch.nn as nn
import torch.nn.functional as F#涉及到 warplayer和refine代码 之后这两个也要看
from model.warplayer import warp
from model.refine import *#输入通道数：上一层的特征图有多少个通道 输出通道数：希望这层输出多少个通道
#kernel_size卷积核大小 默认是4*4 步长是2 填充是1
#这个就需要针对于输出尺寸的计算 这也就是空间尺寸的计算 跟输入输出通道无关 高度+2倍的填充-卷积核大小 这个除以步长 再加上1
#反卷积模块
def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):return nn.Sequential(torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),nn.PReLU(out_planes))
#dilation空洞卷积扩张率
def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):#sequential顺序容器 执行卷积之后 执行激活函数return nn.Sequential(nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,padding=padding, dilation=dilation, bias=True),#卷积层输出多少 激活层就激活多少nn.PReLU(out_planes))#IFBlock类定义
class IFBlock(nn.Module):def __init__(self, in_planes, c=64):super(IFBlock, self).__init__()self.conv0 = nn.Sequential(conv(in_planes, c//2, 3, 2, 1),conv(c//2, c, 3, 2, 1),)#连续八个卷积模块 每个模块包括卷积+Prelu激活self.convblock = nn.Sequential(conv(c, c),conv(c, c),conv(c, c),conv(c, c),conv(c, c),conv(c, c),conv(c, c),conv(c, c),)#kernel_size=4, stride=2, padding=1 → 这是一个标准的上采样配置#关于convtranspose2d可以去官方文档看 二维转置卷积层也叫做反卷积层（上采样卷积层）self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)#def forward(self, x, flow, scale):#F.interpolate 进行插值 使得图像放大或缩小if scale != 1:x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False)if flow != None:flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scalex = torch.cat((x, flow), 1)#将图像与flow在通道维度上进行拼接#上面的Init中定义的 conv0 convblock lastconv#x是输入的图像 b c h w张量会经过这些操作 变成 。。的张量x = self.conv0(x)x = self.convblock(x) + xtmp = self.lastconv(x)#到这里为止最后的c是5#tmp输入 输入的仍是 b c h w形式的张量#scale_factor是放大倍数 针对于双线性插值可以去找个视频看看tmp = F.interpolate(tmp, scale_factor = scale * 2, mode="bilinear", align_corners=False)#得到flow mask b 5 h wflow = tmp[:, :4] * scale * 2#提取光流前四个通道mask = tmp[:, 4:5]#最后一个通道return flow, maskclass IFNet(nn.Module):def __init__(self):super(IFNet, self).__init__()#这些block都依赖于前面IFBlock的定义#这些执行Init 每一个都封装了两层卷积 + 下采样（conv0）# 八层卷积块（convblock）一个转置卷积层用于上采样（lastconv）#这些构建好之后 准备在forward中进行调用self.block0 = IFBlock(6, c=240)self.block1 = IFBlock(13+4, c=150)self.block2 = IFBlock(13+4, c=90)self.block_tea = IFBlock(16+4, c=90)self.contextnet = Contextnet()#经典的编码器-解码器结构self.unet = Unet()def forward(self, x, scale=[4,2,1], timestep=0.5):#初始 结尾 中间真实 这个顺序不用感到疑惑 前面拼接的时候就是按照这个顺序拼接的img0 = x[:, :3]img1 = x[:, 3:6]gt = x[:, 6:] # In inference time, gt is None#这些定义数组之类的后面看都干啥了flow_list = []#生成merged = []mask_list = []warped_img0 = img0warped_img1 = img1flow = None#学生模型和教师模型生成中间帧差异损失函数loss_distill = 0#stu包括这三个 每个都是封装了IFBlcok中的模块stu = [self.block0, self.block1, self.block2]for i in range(3):if flow != None:#括号中作为输入的x#调用的是IFBlock中的forward方法#每次调用 stu[i](...)，实际上就是在调用 IFBlock.forward(...#返回flow mask 这个拼接的很多 3+3+3+3+1=13 后面的加4是针对于flowflow_d, mask_d = stu[i](torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i])flow = flow + flow_dmask = mask + mask_delse:#在flow=none 时候 传入img0 img1 return flow maskflow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i])#对每一个mask都进行sigmod激活 得到的值是0-1 然后加入到Mask_list中mask_list.append(torch.sigmoid(mask))#加入每一个flowflow_list.append(flow)#在这里是进行更新 看warp函数warped_img0 = warp(img0, flow[:, :2])warped_img1 = warp(img1, flow[:, 2:4])#被加入到merged当中merged_student = (warped_img0, warped_img1)merged.append(merged_student)#for循环结束之后#gt是真实的 如果c等于3 有gt是在train阶段if gt.shape[1] == 3:#在前面有这样的定义：self.block_tea = IFBlock(16+4, c=90)#最后更新得到的warped_img0, warped_img1参与 return flow mask#教师模型加上一个gt 维度又加3 那么加上这个gt拼接的优点在哪？#有gt 生成的更好#flow和mask都是前面的for循环中最后的flow和maskflow_d, mask_d = self.block_tea(torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1)#得到flow_teacherflow_teacher = flow + flow_d#得到。。warped_img0_teacher = warp(img0, flow_teacher[:, :2])warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])#对于掩码mask_teacher = torch.sigmoid(mask + mask_d)#得到merged_teachermerged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)else:flow_teacher = Nonemerged_teacher = None#用来指导for i in range(3):#merged_student = (warped_img0, warped_img1) merged.append(merged_student)#这个是前面的代码 指明merged中包含的是啥？每个元素是由warped_img0和warped_img1组成merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])#在训练阶段if gt.shape[1] == 3:#损失计算 loss_distill三阶段总和loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach()loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()#要是想知道光流的四个维度前两个和后两个代表啥意思 还需要看光流的作用含义之类的#反正现在知道的是flow = tmp[:, :4] * scale * 2#提取光流前四个通道#最主要是去看refine文件中的 这两个类的定义 后续继续看c0 = self.contextnet(img0, flow[:, :2])c1 = self.contextnet(img1, flow[:, 2:4])#关于unet和contextnet的定义在refine.py中有定义 refine.py这个代码也要看tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)#return torch.sigmoid(x) 这是在Unet最终的返回 那么维度上的值乘以2 是0到2范围 减去1 是-1 到 1范围# 在c的取出前三个res = tmp[:, :3] * 2 - 1#（是否经历for 看是不是训练阶段） 加到merged[2]上面#torch.clamp(tensor, min=0, max=1) 限制像素值在 0 1范围内merged[2] = torch.clamp(merged[2] + res, 0, 1)#最终返回的这些# flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(torch.cat((imgs, gt), 1), scale=[4, 2, 1])return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill
