zR commited on
Commit
09b490c
1 Parent(s): 8b2ce2a

for 5b demo

Browse files
rife/IFNet.py DELETED
@@ -1,123 +0,0 @@
1
- from .refine import *
2
-
3
-
4
- def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
5
- return nn.Sequential(
6
- torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),
7
- nn.PReLU(out_planes),
8
- )
9
-
10
-
11
- def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
12
- return nn.Sequential(
13
- nn.Conv2d(
14
- in_planes,
15
- out_planes,
16
- kernel_size=kernel_size,
17
- stride=stride,
18
- padding=padding,
19
- dilation=dilation,
20
- bias=True,
21
- ),
22
- nn.PReLU(out_planes),
23
- )
24
-
25
-
26
- class IFBlock(nn.Module):
27
- def __init__(self, in_planes, c=64):
28
- super(IFBlock, self).__init__()
29
- self.conv0 = nn.Sequential(
30
- conv(in_planes, c // 2, 3, 2, 1),
31
- conv(c // 2, c, 3, 2, 1),
32
- )
33
- self.convblock = nn.Sequential(
34
- conv(c, c),
35
- conv(c, c),
36
- conv(c, c),
37
- conv(c, c),
38
- conv(c, c),
39
- conv(c, c),
40
- conv(c, c),
41
- conv(c, c),
42
- )
43
- self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
44
-
45
- def forward(self, x, flow, scale):
46
- if scale != 1:
47
- x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
48
- if flow != None:
49
- flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) * 1.0 / scale
50
- x = torch.cat((x, flow), 1)
51
- x = self.conv0(x)
52
- x = self.convblock(x) + x
53
- tmp = self.lastconv(x)
54
- tmp = F.interpolate(tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False)
55
- flow = tmp[:, :4] * scale * 2
56
- mask = tmp[:, 4:5]
57
- return flow, mask
58
-
59
-
60
- class IFNet(nn.Module):
61
- def __init__(self):
62
- super(IFNet, self).__init__()
63
- self.block0 = IFBlock(6, c=240)
64
- self.block1 = IFBlock(13 + 4, c=150)
65
- self.block2 = IFBlock(13 + 4, c=90)
66
- self.block_tea = IFBlock(16 + 4, c=90)
67
- self.contextnet = Contextnet()
68
- self.unet = Unet()
69
-
70
- def forward(self, x, scale=[4, 2, 1], timestep=0.5):
71
- img0 = x[:, :3]
72
- img1 = x[:, 3:6]
73
- gt = x[:, 6:] # In inference time, gt is None
74
- flow_list = []
75
- merged = []
76
- mask_list = []
77
- warped_img0 = img0
78
- warped_img1 = img1
79
- flow = None
80
- loss_distill = 0
81
- stu = [self.block0, self.block1, self.block2]
82
- for i in range(3):
83
- if flow != None:
84
- flow_d, mask_d = stu[i](
85
- torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i]
86
- )
87
- flow = flow + flow_d
88
- mask = mask + mask_d
89
- else:
90
- flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i])
91
- mask_list.append(torch.sigmoid(mask))
92
- flow_list.append(flow)
93
- warped_img0 = warp(img0, flow[:, :2])
94
- warped_img1 = warp(img1, flow[:, 2:4])
95
- merged_student = (warped_img0, warped_img1)
96
- merged.append(merged_student)
97
- if gt.shape[1] == 3:
98
- flow_d, mask_d = self.block_tea(
99
- torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1
100
- )
101
- flow_teacher = flow + flow_d
102
- warped_img0_teacher = warp(img0, flow_teacher[:, :2])
103
- warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])
104
- mask_teacher = torch.sigmoid(mask + mask_d)
105
- merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)
106
- else:
107
- flow_teacher = None
108
- merged_teacher = None
109
- for i in range(3):
110
- merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
111
- if gt.shape[1] == 3:
112
- loss_mask = (
113
- ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01)
114
- .float()
115
- .detach()
116
- )
117
- loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()
118
- c0 = self.contextnet(img0, flow[:, :2])
119
- c1 = self.contextnet(img1, flow[:, 2:4])
120
- tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
121
- res = tmp[:, :3] * 2 - 1
122
- merged[2] = torch.clamp(merged[2] + res, 0, 1)
123
- return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/IFNet_2R.py DELETED
@@ -1,123 +0,0 @@
1
- from .refine_2R import *
2
-
3
-
4
- def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
5
- return nn.Sequential(
6
- torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),
7
- nn.PReLU(out_planes),
8
- )
9
-
10
-
11
- def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
12
- return nn.Sequential(
13
- nn.Conv2d(
14
- in_planes,
15
- out_planes,
16
- kernel_size=kernel_size,
17
- stride=stride,
18
- padding=padding,
19
- dilation=dilation,
20
- bias=True,
21
- ),
22
- nn.PReLU(out_planes),
23
- )
24
-
25
-
26
- class IFBlock(nn.Module):
27
- def __init__(self, in_planes, c=64):
28
- super(IFBlock, self).__init__()
29
- self.conv0 = nn.Sequential(
30
- conv(in_planes, c // 2, 3, 1, 1),
31
- conv(c // 2, c, 3, 2, 1),
32
- )
33
- self.convblock = nn.Sequential(
34
- conv(c, c),
35
- conv(c, c),
36
- conv(c, c),
37
- conv(c, c),
38
- conv(c, c),
39
- conv(c, c),
40
- conv(c, c),
41
- conv(c, c),
42
- )
43
- self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
44
-
45
- def forward(self, x, flow, scale):
46
- if scale != 1:
47
- x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
48
- if flow != None:
49
- flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) * 1.0 / scale
50
- x = torch.cat((x, flow), 1)
51
- x = self.conv0(x)
52
- x = self.convblock(x) + x
53
- tmp = self.lastconv(x)
54
- tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False)
55
- flow = tmp[:, :4] * scale
56
- mask = tmp[:, 4:5]
57
- return flow, mask
58
-
59
-
60
- class IFNet(nn.Module):
61
- def __init__(self):
62
- super(IFNet, self).__init__()
63
- self.block0 = IFBlock(6, c=240)
64
- self.block1 = IFBlock(13 + 4, c=150)
65
- self.block2 = IFBlock(13 + 4, c=90)
66
- self.block_tea = IFBlock(16 + 4, c=90)
67
- self.contextnet = Contextnet()
68
- self.unet = Unet()
69
-
70
- def forward(self, x, scale=[4, 2, 1], timestep=0.5):
71
- img0 = x[:, :3]
72
- img1 = x[:, 3:6]
73
- gt = x[:, 6:] # In inference time, gt is None
74
- flow_list = []
75
- merged = []
76
- mask_list = []
77
- warped_img0 = img0
78
- warped_img1 = img1
79
- flow = None
80
- loss_distill = 0
81
- stu = [self.block0, self.block1, self.block2]
82
- for i in range(3):
83
- if flow != None:
84
- flow_d, mask_d = stu[i](
85
- torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i]
86
- )
87
- flow = flow + flow_d
88
- mask = mask + mask_d
89
- else:
90
- flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i])
91
- mask_list.append(torch.sigmoid(mask))
92
- flow_list.append(flow)
93
- warped_img0 = warp(img0, flow[:, :2])
94
- warped_img1 = warp(img1, flow[:, 2:4])
95
- merged_student = (warped_img0, warped_img1)
96
- merged.append(merged_student)
97
- if gt.shape[1] == 3:
98
- flow_d, mask_d = self.block_tea(
99
- torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1
100
- )
101
- flow_teacher = flow + flow_d
102
- warped_img0_teacher = warp(img0, flow_teacher[:, :2])
103
- warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])
104
- mask_teacher = torch.sigmoid(mask + mask_d)
105
- merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)
106
- else:
107
- flow_teacher = None
108
- merged_teacher = None
109
- for i in range(3):
110
- merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
111
- if gt.shape[1] == 3:
112
- loss_mask = (
113
- ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01)
114
- .float()
115
- .detach()
116
- )
117
- loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()
118
- c0 = self.contextnet(img0, flow[:, :2])
119
- c1 = self.contextnet(img1, flow[:, 2:4])
120
- tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
121
- res = tmp[:, :3] * 2 - 1
122
- merged[2] = torch.clamp(merged[2] + res, 0, 1)
123
- return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/IFNet_HDv3.py DELETED
@@ -1,138 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
- import torch.nn.functional as F
4
- from .warplayer import warp
5
-
6
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
-
8
-
9
- def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
10
- return nn.Sequential(
11
- nn.Conv2d(
12
- in_planes,
13
- out_planes,
14
- kernel_size=kernel_size,
15
- stride=stride,
16
- padding=padding,
17
- dilation=dilation,
18
- bias=True,
19
- ),
20
- nn.PReLU(out_planes),
21
- )
22
-
23
-
24
- def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
25
- return nn.Sequential(
26
- nn.Conv2d(
27
- in_planes,
28
- out_planes,
29
- kernel_size=kernel_size,
30
- stride=stride,
31
- padding=padding,
32
- dilation=dilation,
33
- bias=False,
34
- ),
35
- nn.BatchNorm2d(out_planes),
36
- nn.PReLU(out_planes),
37
- )
38
-
39
-
40
- class IFBlock(nn.Module):
41
- def __init__(self, in_planes, c=64):
42
- super(IFBlock, self).__init__()
43
- self.conv0 = nn.Sequential(
44
- conv(in_planes, c // 2, 3, 2, 1),
45
- conv(c // 2, c, 3, 2, 1),
46
- )
47
- self.convblock0 = nn.Sequential(conv(c, c), conv(c, c))
48
- self.convblock1 = nn.Sequential(conv(c, c), conv(c, c))
49
- self.convblock2 = nn.Sequential(conv(c, c), conv(c, c))
50
- self.convblock3 = nn.Sequential(conv(c, c), conv(c, c))
51
- self.conv1 = nn.Sequential(
52
- nn.ConvTranspose2d(c, c // 2, 4, 2, 1),
53
- nn.PReLU(c // 2),
54
- nn.ConvTranspose2d(c // 2, 4, 4, 2, 1),
55
- )
56
- self.conv2 = nn.Sequential(
57
- nn.ConvTranspose2d(c, c // 2, 4, 2, 1),
58
- nn.PReLU(c // 2),
59
- nn.ConvTranspose2d(c // 2, 1, 4, 2, 1),
60
- )
61
-
62
- def forward(self, x, flow, scale=1):
63
- x = F.interpolate(
64
- x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False
65
- )
66
- flow = (
67
- F.interpolate(
68
- flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False, recompute_scale_factor=False
69
- )
70
- * 1.0
71
- / scale
72
- )
73
- feat = self.conv0(torch.cat((x, flow), 1))
74
- feat = self.convblock0(feat) + feat
75
- feat = self.convblock1(feat) + feat
76
- feat = self.convblock2(feat) + feat
77
- feat = self.convblock3(feat) + feat
78
- flow = self.conv1(feat)
79
- mask = self.conv2(feat)
80
- flow = (
81
- F.interpolate(flow, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False)
82
- * scale
83
- )
84
- mask = F.interpolate(
85
- mask, scale_factor=scale, mode="bilinear", align_corners=False, recompute_scale_factor=False
86
- )
87
- return flow, mask
88
-
89
-
90
- class IFNet(nn.Module):
91
- def __init__(self):
92
- super(IFNet, self).__init__()
93
- self.block0 = IFBlock(7 + 4, c=90)
94
- self.block1 = IFBlock(7 + 4, c=90)
95
- self.block2 = IFBlock(7 + 4, c=90)
96
- self.block_tea = IFBlock(10 + 4, c=90)
97
- # self.contextnet = Contextnet()
98
- # self.unet = Unet()
99
-
100
- def forward(self, x, scale_list=[4, 2, 1], training=False):
101
- if training == False:
102
- channel = x.shape[1] // 2
103
- img0 = x[:, :channel]
104
- img1 = x[:, channel:]
105
- flow_list = []
106
- merged = []
107
- mask_list = []
108
- warped_img0 = img0
109
- warped_img1 = img1
110
- flow = (x[:, :4]).detach() * 0
111
- mask = (x[:, :1]).detach() * 0
112
- loss_cons = 0
113
- block = [self.block0, self.block1, self.block2]
114
- for i in range(3):
115
- f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), flow, scale=scale_list[i])
116
- f1, m1 = block[i](
117
- torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1),
118
- torch.cat((flow[:, 2:4], flow[:, :2]), 1),
119
- scale=scale_list[i],
120
- )
121
- flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
122
- mask = mask + (m0 + (-m1)) / 2
123
- mask_list.append(mask)
124
- flow_list.append(flow)
125
- warped_img0 = warp(img0, flow[:, :2])
126
- warped_img1 = warp(img1, flow[:, 2:4])
127
- merged.append((warped_img0, warped_img1))
128
- """
129
- c0 = self.contextnet(img0, flow[:, :2])
130
- c1 = self.contextnet(img1, flow[:, 2:4])
131
- tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
132
- res = tmp[:, 1:4] * 2 - 1
133
- """
134
- for i in range(3):
135
- mask_list[i] = torch.sigmoid(mask_list[i])
136
- merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
137
- # merged[i] = torch.clamp(merged[i] + res, 0, 1)
138
- return flow_list, mask_list[2], merged
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/IFNet_m.py DELETED
@@ -1,127 +0,0 @@
1
- from .refine import *
2
-
3
-
4
- def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
5
- return nn.Sequential(
6
- torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),
7
- nn.PReLU(out_planes),
8
- )
9
-
10
-
11
- def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
12
- return nn.Sequential(
13
- nn.Conv2d(
14
- in_planes,
15
- out_planes,
16
- kernel_size=kernel_size,
17
- stride=stride,
18
- padding=padding,
19
- dilation=dilation,
20
- bias=True,
21
- ),
22
- nn.PReLU(out_planes),
23
- )
24
-
25
-
26
- class IFBlock(nn.Module):
27
- def __init__(self, in_planes, c=64):
28
- super(IFBlock, self).__init__()
29
- self.conv0 = nn.Sequential(
30
- conv(in_planes, c // 2, 3, 2, 1),
31
- conv(c // 2, c, 3, 2, 1),
32
- )
33
- self.convblock = nn.Sequential(
34
- conv(c, c),
35
- conv(c, c),
36
- conv(c, c),
37
- conv(c, c),
38
- conv(c, c),
39
- conv(c, c),
40
- conv(c, c),
41
- conv(c, c),
42
- )
43
- self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
44
-
45
- def forward(self, x, flow, scale):
46
- if scale != 1:
47
- x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
48
- if flow != None:
49
- flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) * 1.0 / scale
50
- x = torch.cat((x, flow), 1)
51
- x = self.conv0(x)
52
- x = self.convblock(x) + x
53
- tmp = self.lastconv(x)
54
- tmp = F.interpolate(tmp, scale_factor=scale * 2, mode="bilinear", align_corners=False)
55
- flow = tmp[:, :4] * scale * 2
56
- mask = tmp[:, 4:5]
57
- return flow, mask
58
-
59
-
60
- class IFNet_m(nn.Module):
61
- def __init__(self):
62
- super(IFNet_m, self).__init__()
63
- self.block0 = IFBlock(6 + 1, c=240)
64
- self.block1 = IFBlock(13 + 4 + 1, c=150)
65
- self.block2 = IFBlock(13 + 4 + 1, c=90)
66
- self.block_tea = IFBlock(16 + 4 + 1, c=90)
67
- self.contextnet = Contextnet()
68
- self.unet = Unet()
69
-
70
- def forward(self, x, scale=[4, 2, 1], timestep=0.5, returnflow=False):
71
- timestep = (x[:, :1].clone() * 0 + 1) * timestep
72
- img0 = x[:, :3]
73
- img1 = x[:, 3:6]
74
- gt = x[:, 6:] # In inference time, gt is None
75
- flow_list = []
76
- merged = []
77
- mask_list = []
78
- warped_img0 = img0
79
- warped_img1 = img1
80
- flow = None
81
- loss_distill = 0
82
- stu = [self.block0, self.block1, self.block2]
83
- for i in range(3):
84
- if flow != None:
85
- flow_d, mask_d = stu[i](
86
- torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask), 1), flow, scale=scale[i]
87
- )
88
- flow = flow + flow_d
89
- mask = mask + mask_d
90
- else:
91
- flow, mask = stu[i](torch.cat((img0, img1, timestep), 1), None, scale=scale[i])
92
- mask_list.append(torch.sigmoid(mask))
93
- flow_list.append(flow)
94
- warped_img0 = warp(img0, flow[:, :2])
95
- warped_img1 = warp(img1, flow[:, 2:4])
96
- merged_student = (warped_img0, warped_img1)
97
- merged.append(merged_student)
98
- if gt.shape[1] == 3:
99
- flow_d, mask_d = self.block_tea(
100
- torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask, gt), 1), flow, scale=1
101
- )
102
- flow_teacher = flow + flow_d
103
- warped_img0_teacher = warp(img0, flow_teacher[:, :2])
104
- warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])
105
- mask_teacher = torch.sigmoid(mask + mask_d)
106
- merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)
107
- else:
108
- flow_teacher = None
109
- merged_teacher = None
110
- for i in range(3):
111
- merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
112
- if gt.shape[1] == 3:
113
- loss_mask = (
114
- ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01)
115
- .float()
116
- .detach()
117
- )
118
- loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()
119
- if returnflow:
120
- return flow
121
- else:
122
- c0 = self.contextnet(img0, flow[:, :2])
123
- c1 = self.contextnet(img1, flow[:, 2:4])
124
- tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
125
- res = tmp[:, :3] * 2 - 1
126
- merged[2] = torch.clamp(merged[2] + res, 0, 1)
127
- return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/RIFE.py DELETED
@@ -1,95 +0,0 @@
1
- from torch.optim import AdamW
2
- from torch.nn.parallel import DistributedDataParallel as DDP
3
- from .IFNet import *
4
- from .IFNet_m import *
5
- from .loss import *
6
- from .laplacian import *
7
- from .refine import *
8
-
9
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
-
11
-
12
- class Model:
13
- def __init__(self, local_rank=-1, arbitrary=False):
14
- if arbitrary == True:
15
- self.flownet = IFNet_m()
16
- else:
17
- self.flownet = IFNet()
18
- self.device()
19
- self.optimG = AdamW(
20
- self.flownet.parameters(), lr=1e-6, weight_decay=1e-3
21
- ) # use large weight decay may avoid NaN loss
22
- self.epe = EPE()
23
- self.lap = LapLoss()
24
- self.sobel = SOBEL()
25
- if local_rank != -1:
26
- self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank)
27
-
28
- def train(self):
29
- self.flownet.train()
30
-
31
- def eval(self):
32
- self.flownet.eval()
33
-
34
- def device(self):
35
- self.flownet.to(device)
36
-
37
- def load_model(self, path, rank=0):
38
- def convert(param):
39
- return {k.replace("module.", ""): v for k, v in param.items() if "module." in k}
40
-
41
- if rank <= 0:
42
- self.flownet.load_state_dict(convert(torch.load("{}/flownet.pkl".format(path))))
43
-
44
- def save_model(self, path, rank=0):
45
- if rank == 0:
46
- torch.save(self.flownet.state_dict(), "{}/flownet.pkl".format(path))
47
-
48
- def inference(self, img0, img1, scale=1, scale_list=[4, 2, 1], TTA=False, timestep=0.5):
49
- for i in range(3):
50
- scale_list[i] = scale_list[i] * 1.0 / scale
51
- imgs = torch.cat((img0, img1), 1)
52
- flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(
53
- imgs, scale_list, timestep=timestep
54
- )
55
- if TTA == False:
56
- return merged[2]
57
- else:
58
- flow2, mask2, merged2, flow_teacher2, merged_teacher2, loss_distill2 = self.flownet(
59
- imgs.flip(2).flip(3), scale_list, timestep=timestep
60
- )
61
- return (merged[2] + merged2[2].flip(2).flip(3)) / 2
62
-
63
- def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
64
- for param_group in self.optimG.param_groups:
65
- param_group["lr"] = learning_rate
66
- img0 = imgs[:, :3]
67
- img1 = imgs[:, 3:]
68
- if training:
69
- self.train()
70
- else:
71
- self.eval()
72
- flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(
73
- torch.cat((imgs, gt), 1), scale=[4, 2, 1]
74
- )
75
- loss_l1 = (self.lap(merged[2], gt)).mean()
76
- loss_tea = (self.lap(merged_teacher, gt)).mean()
77
- if training:
78
- self.optimG.zero_grad()
79
- loss_G = (
80
- loss_l1 + loss_tea + loss_distill * 0.01
81
- ) # when training RIFEm, the weight of loss_distill should be 0.005 or 0.002
82
- loss_G.backward()
83
- self.optimG.step()
84
- else:
85
- flow_teacher = flow[2]
86
- return merged[2], {
87
- "merged_tea": merged_teacher,
88
- "mask": mask,
89
- "mask_tea": mask,
90
- "flow": flow[2][:, :2],
91
- "flow_tea": flow_teacher,
92
- "loss_l1": loss_l1,
93
- "loss_tea": loss_tea,
94
- "loss_distill": loss_distill,
95
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/RIFE_HDv3.py DELETED
@@ -1,86 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
- import numpy as np
4
- from torch.optim import AdamW
5
- import torch.optim as optim
6
- import itertools
7
- from .warplayer import warp
8
- from torch.nn.parallel import DistributedDataParallel as DDP
9
- from .IFNet_HDv3 import *
10
- import torch.nn.functional as F
11
- from .loss import *
12
-
13
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
-
15
-
16
- class Model:
17
- def __init__(self, local_rank=-1):
18
- self.flownet = IFNet()
19
- self.device()
20
- self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4)
21
- self.epe = EPE()
22
- # self.vgg = VGGPerceptualLoss().to(device)
23
- self.sobel = SOBEL()
24
- if local_rank != -1:
25
- self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank)
26
-
27
- def train(self):
28
- self.flownet.train()
29
-
30
- def eval(self):
31
- self.flownet.eval()
32
-
33
- def device(self):
34
- self.flownet.to(device)
35
-
36
- def load_model(self, path, rank=0):
37
- def convert(param):
38
- if rank == -1:
39
- return {k.replace("module.", ""): v for k, v in param.items() if "module." in k}
40
- else:
41
- return param
42
-
43
- if rank <= 0:
44
- if torch.cuda.is_available():
45
- self.flownet.load_state_dict(convert(torch.load("{}/flownet.pkl".format(path))))
46
- else:
47
- self.flownet.load_state_dict(convert(torch.load("{}/flownet.pkl".format(path), map_location="cpu")))
48
-
49
- def save_model(self, path, rank=0):
50
- if rank == 0:
51
- torch.save(self.flownet.state_dict(), "{}/flownet.pkl".format(path))
52
-
53
- def inference(self, img0, img1, scale=1.0):
54
- imgs = torch.cat((img0, img1), 1)
55
- scale_list = [4 / scale, 2 / scale, 1 / scale]
56
- flow, mask, merged = self.flownet(imgs, scale_list)
57
- return merged[2]
58
-
59
- def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
60
- for param_group in self.optimG.param_groups:
61
- param_group["lr"] = learning_rate
62
- img0 = imgs[:, :3]
63
- img1 = imgs[:, 3:]
64
- if training:
65
- self.train()
66
- else:
67
- self.eval()
68
- scale = [4, 2, 1]
69
- flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training)
70
- loss_l1 = (merged[2] - gt).abs().mean()
71
- loss_smooth = self.sobel(flow[2], flow[2] * 0).mean()
72
- # loss_vgg = self.vgg(merged[2], gt)
73
- if training:
74
- self.optimG.zero_grad()
75
- loss_G = loss_cons + loss_smooth * 0.1
76
- loss_G.backward()
77
- self.optimG.step()
78
- else:
79
- flow_teacher = flow[2]
80
- return merged[2], {
81
- "mask": mask,
82
- "flow": flow[2][:, :2],
83
- "loss_l1": loss_l1,
84
- "loss_cons": loss_cons,
85
- "loss_smooth": loss_smooth,
86
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/__init__.py DELETED
File without changes
rife/laplacian.py DELETED
@@ -1,69 +0,0 @@
1
- import torch
2
- import numpy as np
3
- import torch.nn as nn
4
- import torch.nn.functional as F
5
-
6
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
-
8
- import torch
9
-
10
-
11
- def gauss_kernel(size=5, channels=3):
12
- kernel = torch.tensor(
13
- [
14
- [1.0, 4.0, 6.0, 4.0, 1],
15
- [4.0, 16.0, 24.0, 16.0, 4.0],
16
- [6.0, 24.0, 36.0, 24.0, 6.0],
17
- [4.0, 16.0, 24.0, 16.0, 4.0],
18
- [1.0, 4.0, 6.0, 4.0, 1.0],
19
- ]
20
- )
21
- kernel /= 256.0
22
- kernel = kernel.repeat(channels, 1, 1, 1)
23
- kernel = kernel.to(device)
24
- return kernel
25
-
26
-
27
- def downsample(x):
28
- return x[:, :, ::2, ::2]
29
-
30
-
31
- def upsample(x):
32
- cc = torch.cat([x, torch.zeros(x.shape[0], x.shape[1], x.shape[2], x.shape[3]).to(device)], dim=3)
33
- cc = cc.view(x.shape[0], x.shape[1], x.shape[2] * 2, x.shape[3])
34
- cc = cc.permute(0, 1, 3, 2)
35
- cc = torch.cat([cc, torch.zeros(x.shape[0], x.shape[1], x.shape[3], x.shape[2] * 2).to(device)], dim=3)
36
- cc = cc.view(x.shape[0], x.shape[1], x.shape[3] * 2, x.shape[2] * 2)
37
- x_up = cc.permute(0, 1, 3, 2)
38
- return conv_gauss(x_up, 4 * gauss_kernel(channels=x.shape[1]))
39
-
40
-
41
- def conv_gauss(img, kernel):
42
- img = torch.nn.functional.pad(img, (2, 2, 2, 2), mode="reflect")
43
- out = torch.nn.functional.conv2d(img, kernel, groups=img.shape[1])
44
- return out
45
-
46
-
47
- def laplacian_pyramid(img, kernel, max_levels=3):
48
- current = img
49
- pyr = []
50
- for level in range(max_levels):
51
- filtered = conv_gauss(current, kernel)
52
- down = downsample(filtered)
53
- up = upsample(down)
54
- diff = current - up
55
- pyr.append(diff)
56
- current = down
57
- return pyr
58
-
59
-
60
- class LapLoss(torch.nn.Module):
61
- def __init__(self, max_levels=5, channels=3):
62
- super(LapLoss, self).__init__()
63
- self.max_levels = max_levels
64
- self.gauss_kernel = gauss_kernel(channels=channels)
65
-
66
- def forward(self, input, target):
67
- pyr_input = laplacian_pyramid(img=input, kernel=self.gauss_kernel, max_levels=self.max_levels)
68
- pyr_target = laplacian_pyramid(img=target, kernel=self.gauss_kernel, max_levels=self.max_levels)
69
- return sum(torch.nn.functional.l1_loss(a, b) for a, b in zip(pyr_input, pyr_target))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/loss.py DELETED
@@ -1,130 +0,0 @@
1
- import torch
2
- import numpy as np
3
- import torch.nn as nn
4
- import torch.nn.functional as F
5
- import torchvision.models as models
6
-
7
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
-
9
-
10
- class EPE(nn.Module):
11
- def __init__(self):
12
- super(EPE, self).__init__()
13
-
14
- def forward(self, flow, gt, loss_mask):
15
- loss_map = (flow - gt.detach()) ** 2
16
- loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
17
- return loss_map * loss_mask
18
-
19
-
20
- class Ternary(nn.Module):
21
- def __init__(self):
22
- super(Ternary, self).__init__()
23
- patch_size = 7
24
- out_channels = patch_size * patch_size
25
- self.w = np.eye(out_channels).reshape((patch_size, patch_size, 1, out_channels))
26
- self.w = np.transpose(self.w, (3, 2, 0, 1))
27
- self.w = torch.tensor(self.w).float().to(device)
28
-
29
- def transform(self, img):
30
- patches = F.conv2d(img, self.w, padding=3, bias=None)
31
- transf = patches - img
32
- transf_norm = transf / torch.sqrt(0.81 + transf**2)
33
- return transf_norm
34
-
35
- def rgb2gray(self, rgb):
36
- r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
37
- gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
38
- return gray
39
-
40
- def hamming(self, t1, t2):
41
- dist = (t1 - t2) ** 2
42
- dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
43
- return dist_norm
44
-
45
- def valid_mask(self, t, padding):
46
- n, _, h, w = t.size()
47
- inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
48
- mask = F.pad(inner, [padding] * 4)
49
- return mask
50
-
51
- def forward(self, img0, img1):
52
- img0 = self.transform(self.rgb2gray(img0))
53
- img1 = self.transform(self.rgb2gray(img1))
54
- return self.hamming(img0, img1) * self.valid_mask(img0, 1)
55
-
56
-
57
- class SOBEL(nn.Module):
58
- def __init__(self):
59
- super(SOBEL, self).__init__()
60
- self.kernelX = torch.tensor(
61
- [
62
- [1, 0, -1],
63
- [2, 0, -2],
64
- [1, 0, -1],
65
- ]
66
- ).float()
67
- self.kernelY = self.kernelX.clone().T
68
- self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device)
69
- self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device)
70
-
71
- def forward(self, pred, gt):
72
- N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
73
- img_stack = torch.cat([pred.reshape(N * C, 1, H, W), gt.reshape(N * C, 1, H, W)], 0)
74
- sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
75
- sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
76
- pred_X, gt_X = sobel_stack_x[: N * C], sobel_stack_x[N * C :]
77
- pred_Y, gt_Y = sobel_stack_y[: N * C], sobel_stack_y[N * C :]
78
-
79
- L1X, L1Y = torch.abs(pred_X - gt_X), torch.abs(pred_Y - gt_Y)
80
- loss = L1X + L1Y
81
- return loss
82
-
83
-
84
- class MeanShift(nn.Conv2d):
85
- def __init__(self, data_mean, data_std, data_range=1, norm=True):
86
- c = len(data_mean)
87
- super(MeanShift, self).__init__(c, c, kernel_size=1)
88
- std = torch.Tensor(data_std)
89
- self.weight.data = torch.eye(c).view(c, c, 1, 1)
90
- if norm:
91
- self.weight.data.div_(std.view(c, 1, 1, 1))
92
- self.bias.data = -1 * data_range * torch.Tensor(data_mean)
93
- self.bias.data.div_(std)
94
- else:
95
- self.weight.data.mul_(std.view(c, 1, 1, 1))
96
- self.bias.data = data_range * torch.Tensor(data_mean)
97
- self.requires_grad = False
98
-
99
-
100
- class VGGPerceptualLoss(torch.nn.Module):
101
- def __init__(self, rank=0):
102
- super(VGGPerceptualLoss, self).__init__()
103
- blocks = []
104
- pretrained = True
105
- self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features
106
- self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
107
- for param in self.parameters():
108
- param.requires_grad = False
109
-
110
- def forward(self, X, Y, indices=None):
111
- X = self.normalize(X)
112
- Y = self.normalize(Y)
113
- indices = [2, 7, 12, 21, 30]
114
- weights = [1.0 / 2.6, 1.0 / 4.8, 1.0 / 3.7, 1.0 / 5.6, 10 / 1.5]
115
- k = 0
116
- loss = 0
117
- for i in range(indices[-1]):
118
- X = self.vgg_pretrained_features[i](X)
119
- Y = self.vgg_pretrained_features[i](Y)
120
- if (i + 1) in indices:
121
- loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1
122
- k += 1
123
- return loss
124
-
125
-
126
- if __name__ == "__main__":
127
- img0 = torch.zeros(3, 3, 256, 256).float().to(device)
128
- img1 = torch.tensor(np.random.normal(0, 1, (3, 3, 256, 256))).float().to(device)
129
- ternary_loss = Ternary()
130
- print(ternary_loss(img0, img1).shape)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/pytorch_msssim/__init__.py DELETED
@@ -1,203 +0,0 @@
1
- import torch
2
- import torch.nn.functional as F
3
- from math import exp
4
- import numpy as np
5
-
6
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
-
8
-
9
- def gaussian(window_size, sigma):
10
- gauss = torch.Tensor([exp(-((x - window_size // 2) ** 2) / float(2 * sigma**2)) for x in range(window_size)])
11
- return gauss / gauss.sum()
12
-
13
-
14
- def create_window(window_size, channel=1):
15
- _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
16
- _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device)
17
- window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
18
- return window
19
-
20
-
21
- def create_window_3d(window_size, channel=1):
22
- _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
23
- _2D_window = _1D_window.mm(_1D_window.t())
24
- _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t())
25
- window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device)
26
- return window
27
-
28
-
29
- def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
30
- # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
31
- if val_range is None:
32
- if torch.max(img1) > 128:
33
- max_val = 255
34
- else:
35
- max_val = 1
36
-
37
- if torch.min(img1) < -0.5:
38
- min_val = -1
39
- else:
40
- min_val = 0
41
- L = max_val - min_val
42
- else:
43
- L = val_range
44
-
45
- padd = 0
46
- (_, channel, height, width) = img1.size()
47
- if window is None:
48
- real_size = min(window_size, height, width)
49
- window = create_window(real_size, channel=channel).to(img1.device)
50
-
51
- # mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
52
- # mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
53
- mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=channel)
54
- mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=channel)
55
-
56
- mu1_sq = mu1.pow(2)
57
- mu2_sq = mu2.pow(2)
58
- mu1_mu2 = mu1 * mu2
59
-
60
- sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel) - mu1_sq
61
- sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel) - mu2_sq
62
- sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel) - mu1_mu2
63
-
64
- C1 = (0.01 * L) ** 2
65
- C2 = (0.03 * L) ** 2
66
-
67
- v1 = 2.0 * sigma12 + C2
68
- v2 = sigma1_sq + sigma2_sq + C2
69
- cs = torch.mean(v1 / v2) # contrast sensitivity
70
-
71
- ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
72
-
73
- if size_average:
74
- ret = ssim_map.mean()
75
- else:
76
- ret = ssim_map.mean(1).mean(1).mean(1)
77
-
78
- if full:
79
- return ret, cs
80
- return ret
81
-
82
-
83
- def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
84
- # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
85
- if val_range is None:
86
- if torch.max(img1) > 128:
87
- max_val = 255
88
- else:
89
- max_val = 1
90
-
91
- if torch.min(img1) < -0.5:
92
- min_val = -1
93
- else:
94
- min_val = 0
95
- L = max_val - min_val
96
- else:
97
- L = val_range
98
-
99
- padd = 0
100
- (_, _, height, width) = img1.size()
101
- if window is None:
102
- real_size = min(window_size, height, width)
103
- window = create_window_3d(real_size, channel=1).to(img1.device, dtype=img1.dtype)
104
- # Channel is set to 1 since we consider color images as volumetric images
105
-
106
- img1 = img1.unsqueeze(1)
107
- img2 = img2.unsqueeze(1)
108
-
109
- mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=1)
110
- mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=1)
111
-
112
- mu1_sq = mu1.pow(2)
113
- mu2_sq = mu2.pow(2)
114
- mu1_mu2 = mu1 * mu2
115
-
116
- sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1) - mu1_sq
117
- sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1) - mu2_sq
118
- sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1) - mu1_mu2
119
-
120
- C1 = (0.01 * L) ** 2
121
- C2 = (0.03 * L) ** 2
122
-
123
- v1 = 2.0 * sigma12 + C2
124
- v2 = sigma1_sq + sigma2_sq + C2
125
- cs = torch.mean(v1 / v2) # contrast sensitivity
126
-
127
- ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
128
-
129
- if size_average:
130
- ret = ssim_map.mean()
131
- else:
132
- ret = ssim_map.mean(1).mean(1).mean(1)
133
-
134
- if full:
135
- return ret, cs
136
- return ret
137
-
138
-
139
- def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False):
140
- device = img1.device
141
- weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device)
142
- levels = weights.size()[0]
143
- mssim = []
144
- mcs = []
145
- for _ in range(levels):
146
- sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)
147
- mssim.append(sim)
148
- mcs.append(cs)
149
-
150
- img1 = F.avg_pool2d(img1, (2, 2))
151
- img2 = F.avg_pool2d(img2, (2, 2))
152
-
153
- mssim = torch.stack(mssim)
154
- mcs = torch.stack(mcs)
155
-
156
- # Normalize (to avoid NaNs during training unstable models, not compliant with original definition)
157
- if normalize:
158
- mssim = (mssim + 1) / 2
159
- mcs = (mcs + 1) / 2
160
-
161
- pow1 = mcs**weights
162
- pow2 = mssim**weights
163
- # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/
164
- output = torch.prod(pow1[:-1] * pow2[-1])
165
- return output
166
-
167
-
168
- # Classes to re-use window
169
- class SSIM(torch.nn.Module):
170
- def __init__(self, window_size=11, size_average=True, val_range=None):
171
- super(SSIM, self).__init__()
172
- self.window_size = window_size
173
- self.size_average = size_average
174
- self.val_range = val_range
175
-
176
- # Assume 3 channel for SSIM
177
- self.channel = 3
178
- self.window = create_window(window_size, channel=self.channel)
179
-
180
- def forward(self, img1, img2):
181
- (_, channel, _, _) = img1.size()
182
-
183
- if channel == self.channel and self.window.dtype == img1.dtype:
184
- window = self.window
185
- else:
186
- window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)
187
- self.window = window
188
- self.channel = channel
189
-
190
- _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)
191
- dssim = (1 - _ssim) / 2
192
- return dssim
193
-
194
-
195
- class MSSSIM(torch.nn.Module):
196
- def __init__(self, window_size=11, size_average=True, channel=3):
197
- super(MSSSIM, self).__init__()
198
- self.window_size = window_size
199
- self.size_average = size_average
200
- self.channel = channel
201
-
202
- def forward(self, img1, img2):
203
- return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/refine.py DELETED
@@ -1,107 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
- from .warplayer import warp
4
- import torch.nn.functional as F
5
-
6
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
-
8
-
9
- def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
10
- return nn.Sequential(
11
- nn.Conv2d(
12
- in_planes,
13
- out_planes,
14
- kernel_size=kernel_size,
15
- stride=stride,
16
- padding=padding,
17
- dilation=dilation,
18
- bias=True,
19
- ),
20
- nn.PReLU(out_planes),
21
- )
22
-
23
-
24
- def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
25
- return nn.Sequential(
26
- torch.nn.ConvTranspose2d(
27
- in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True
28
- ),
29
- nn.PReLU(out_planes),
30
- )
31
-
32
-
33
- class Conv2(nn.Module):
34
- def __init__(self, in_planes, out_planes, stride=2):
35
- super(Conv2, self).__init__()
36
- self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
37
- self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
38
-
39
- def forward(self, x):
40
- x = self.conv1(x)
41
- x = self.conv2(x)
42
- return x
43
-
44
-
45
- c = 16
46
-
47
-
48
- class Contextnet(nn.Module):
49
- def __init__(self):
50
- super(Contextnet, self).__init__()
51
- self.conv1 = Conv2(3, c)
52
- self.conv2 = Conv2(c, 2 * c)
53
- self.conv3 = Conv2(2 * c, 4 * c)
54
- self.conv4 = Conv2(4 * c, 8 * c)
55
-
56
- def forward(self, x, flow):
57
- x = self.conv1(x)
58
- flow = (
59
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
60
- * 0.5
61
- )
62
- f1 = warp(x, flow)
63
- x = self.conv2(x)
64
- flow = (
65
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
66
- * 0.5
67
- )
68
- f2 = warp(x, flow)
69
- x = self.conv3(x)
70
- flow = (
71
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
72
- * 0.5
73
- )
74
- f3 = warp(x, flow)
75
- x = self.conv4(x)
76
- flow = (
77
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
78
- * 0.5
79
- )
80
- f4 = warp(x, flow)
81
- return [f1, f2, f3, f4]
82
-
83
-
84
- class Unet(nn.Module):
85
- def __init__(self):
86
- super(Unet, self).__init__()
87
- self.down0 = Conv2(17, 2 * c)
88
- self.down1 = Conv2(4 * c, 4 * c)
89
- self.down2 = Conv2(8 * c, 8 * c)
90
- self.down3 = Conv2(16 * c, 16 * c)
91
- self.up0 = deconv(32 * c, 8 * c)
92
- self.up1 = deconv(16 * c, 4 * c)
93
- self.up2 = deconv(8 * c, 2 * c)
94
- self.up3 = deconv(4 * c, c)
95
- self.conv = nn.Conv2d(c, 3, 3, 1, 1)
96
-
97
- def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
98
- s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1))
99
- s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
100
- s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
101
- s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
102
- x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
103
- x = self.up1(torch.cat((x, s2), 1))
104
- x = self.up2(torch.cat((x, s1), 1))
105
- x = self.up3(torch.cat((x, s0), 1))
106
- x = self.conv(x)
107
- return torch.sigmoid(x)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/refine_2R.py DELETED
@@ -1,104 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
- from .warplayer import warp
4
- import torch.nn.functional as F
5
-
6
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
-
8
-
9
- def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
10
- return nn.Sequential(
11
- nn.Conv2d(
12
- in_planes,
13
- out_planes,
14
- kernel_size=kernel_size,
15
- stride=stride,
16
- padding=padding,
17
- dilation=dilation,
18
- bias=True,
19
- ),
20
- nn.PReLU(out_planes),
21
- )
22
-
23
-
24
- def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
25
- return nn.Sequential(
26
- torch.nn.ConvTranspose2d(
27
- in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True
28
- ),
29
- nn.PReLU(out_planes),
30
- )
31
-
32
-
33
- class Conv2(nn.Module):
34
- def __init__(self, in_planes, out_planes, stride=2):
35
- super(Conv2, self).__init__()
36
- self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
37
- self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
38
-
39
- def forward(self, x):
40
- x = self.conv1(x)
41
- x = self.conv2(x)
42
- return x
43
-
44
-
45
- c = 16
46
-
47
-
48
- class Contextnet(nn.Module):
49
- def __init__(self):
50
- super(Contextnet, self).__init__()
51
- self.conv1 = Conv2(3, c, 1)
52
- self.conv2 = Conv2(c, 2 * c)
53
- self.conv3 = Conv2(2 * c, 4 * c)
54
- self.conv4 = Conv2(4 * c, 8 * c)
55
-
56
- def forward(self, x, flow):
57
- x = self.conv1(x)
58
- # flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
59
- f1 = warp(x, flow)
60
- x = self.conv2(x)
61
- flow = (
62
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
63
- * 0.5
64
- )
65
- f2 = warp(x, flow)
66
- x = self.conv3(x)
67
- flow = (
68
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
69
- * 0.5
70
- )
71
- f3 = warp(x, flow)
72
- x = self.conv4(x)
73
- flow = (
74
- F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False)
75
- * 0.5
76
- )
77
- f4 = warp(x, flow)
78
- return [f1, f2, f3, f4]
79
-
80
-
81
- class Unet(nn.Module):
82
- def __init__(self):
83
- super(Unet, self).__init__()
84
- self.down0 = Conv2(17, 2 * c, 1)
85
- self.down1 = Conv2(4 * c, 4 * c)
86
- self.down2 = Conv2(8 * c, 8 * c)
87
- self.down3 = Conv2(16 * c, 16 * c)
88
- self.up0 = deconv(32 * c, 8 * c)
89
- self.up1 = deconv(16 * c, 4 * c)
90
- self.up2 = deconv(8 * c, 2 * c)
91
- self.up3 = deconv(4 * c, c)
92
- self.conv = nn.Conv2d(c, 3, 3, 2, 1)
93
-
94
- def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
95
- s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1))
96
- s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
97
- s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
98
- s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
99
- x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
100
- x = self.up1(torch.cat((x, s2), 1))
101
- x = self.up2(torch.cat((x, s1), 1))
102
- x = self.up3(torch.cat((x, s0), 1))
103
- x = self.conv(x)
104
- return torch.sigmoid(x)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
rife/warplayer.py DELETED
@@ -1,34 +0,0 @@
1
- import torch
2
- import torch.nn as nn
3
-
4
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5
- backwarp_tenGrid = {}
6
-
7
-
8
- def warp(tenInput, tenFlow):
9
- k = (str(tenFlow.device), str(tenFlow.size()))
10
- if k not in backwarp_tenGrid:
11
- tenHorizontal = (
12
- torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device)
13
- .view(1, 1, 1, tenFlow.shape[3])
14
- .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
15
- )
16
- tenVertical = (
17
- torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device)
18
- .view(1, 1, tenFlow.shape[2], 1)
19
- .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
20
- )
21
- backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(device)
22
-
23
- tenFlow = torch.cat(
24
- [
25
- tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
26
- tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0),
27
- ],
28
- 1,
29
- )
30
-
31
- g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
32
- return torch.nn.functional.grid_sample(
33
- input=tenInput, grid=g, mode="bilinear", padding_mode="border", align_corners=True
34
- )