zR commited on
Commit
cc979ab
1 Parent(s): 1aceaa0
rife/IFNet.py ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,127 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
File without changes
rife/laplacian.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,104 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ )