Mirrors/flowframes
1
0
Fork 0
mirror of https://github.com/n00mkrad/flowframes.git synced 2026-09-01 19:51:28 +02:00
Code Issues Packages Projects Releases Wiki Activity

Updated RIFE-CUDA (scale arg, img piping tests)

Browse Source
This commit is contained in:
N00MKRAD
2021-02-27 16:24:42 +01:00
parent 27ca73a900
commit 223f4221fc
5 changed files with 35 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
Pkgs/rife-cuda/model/IFNet_HD.py
View File

@@ -91,12 +91,9 @@ class IFNet(nn.Module):
self.block2 = IFBlock(8, scale=2, c=96)
self.block3 = IFBlock(8, scale=1, c=48)
def forward(self, x, UHD=False):
if UHD:
x = F.interpolate(x, scale_factor=0.25, mode="bilinear", align_corners=False)
else:
x = F.interpolate(x, scale_factor=0.5, mode="bilinear",
align_corners=False)
def forward(self, x, scale=1.0):
x = F.interpolate(x, scale_factor=0.5 * scale, mode="bilinear",
align_corners=False)
flow0 = self.block0(x)
F1 = flow0
warped_img0 = warp(x[:, :3], F1)
@@ -111,6 +108,8 @@ class IFNet(nn.Module):
warped_img1 = warp(x[:, 3:], -F3)
flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3), 1))
F4 = (flow0 + flow1 + flow2 + flow3)
F4 = F.interpolate(F4, scale_factor=1 / scale, mode="bilinear",
align_corners=False) / scale
return F4, [F1, F2, F3, F4]
if __name__ == '__main__':

27
Pkgs/rife-cuda/model/IFNet_HDv2.py
View File

@@ -39,18 +39,19 @@ class IFBlock(nn.Module):
)
self.conv1 = nn.ConvTranspose2d(2*c, 4, 4, 2, 1)
def forward(self, x):
if self.scale != 1:
x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear",
def forward(self, x, scale=1.0):
infer_scale = self.scale / scale
if infer_scale != 1.0:
x = F.interpolate(x, scale_factor=1. / infer_scale, mode="bilinear",
align_corners=False)
x = self.conv0(x)
x = self.convblock(x)
x = self.conv1(x)
flow = x
if self.scale != 1:
flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear",
align_corners=False)
return flow
if infer_scale != 1.0:
flow = F.interpolate(flow, scale_factor=infer_scale, mode="bilinear",
align_corners=False)
return flow / scale
class IFNet(nn.Module):
@@ -61,25 +62,23 @@ class IFNet(nn.Module):
self.block2 = IFBlock(10, scale=2, c=96)
self.block3 = IFBlock(10, scale=1, c=48)
def forward(self, x, UHD=False):
if UHD:
x = F.interpolate(x, scale_factor=0.5, mode="bilinear", align_corners=False)
flow0 = self.block0(x)
def forward(self, x, scale=1.0):
flow0 = self.block0(x, scale)
F1 = flow0
F1_large = F.interpolate(F1, scale_factor=2.0, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 2.0
warped_img0 = warp(x[:, :3], F1_large[:, :2])
warped_img1 = warp(x[:, 3:], F1_large[:, 2:4])
flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1_large), 1))
flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1_large), 1), scale)
F2 = (flow0 + flow1)
F2_large = F.interpolate(F2, scale_factor=2.0, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 2.0
warped_img0 = warp(x[:, :3], F2_large[:, :2])
warped_img1 = warp(x[:, 3:], F2_large[:, 2:4])
flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2_large), 1))
flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2_large), 1), scale)
F3 = (flow0 + flow1 + flow2)
F3_large = F.interpolate(F3, scale_factor=2.0, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 2.0
warped_img0 = warp(x[:, :3], F3_large[:, :2])
warped_img1 = warp(x[:, 3:], F3_large[:, 2:4])
flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3_large), 1))
flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3_large), 1), scale)
F4 = (flow0 + flow1 + flow2 + flow3)
return F4, [F1, F2, F3, F4]

10
Pkgs/rife-cuda/model/RIFE_HD.py
View File

@@ -188,11 +188,9 @@ class Model:
torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path))
torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path))
def predict(self, imgs, flow, training=True, flow_gt=None, UHD=False):
def predict(self, imgs, flow, training=True, flow_gt=None):
img0 = imgs[:, :3]
img1 = imgs[:, 3:]
if UHD:
flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0
c0 = self.contextnet(img0, flow)
c1 = self.contextnet(img1, -flow)
flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear",
@@ -209,10 +207,10 @@ class Model:
else:
return pred
def inference(self, img0, img1, UHD=False):
def inference(self, img0, img1, scale=1.0):
imgs = torch.cat((img0, img1), 1)
flow, _ = self.flownet(imgs, UHD)
return self.predict(imgs, flow, training=False, UHD=UHD)
flow, _ = self.flownet(imgs, scale)
return self.predict(imgs, flow, training=False)
def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
for param_group in self.optimG.param_groups:

10
Pkgs/rife-cuda/model/RIFE_HDv2.py
View File

@@ -173,11 +173,9 @@ class Model:
torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path))
torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path))
def predict(self, imgs, flow, training=True, flow_gt=None, UHD=False):
def predict(self, imgs, flow, training=True, flow_gt=None):
img0 = imgs[:, :3]
img1 = imgs[:, 3:]
if UHD:
flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0
c0 = self.contextnet(img0, flow[:, :2])
c1 = self.contextnet(img1, flow[:, 2:4])
flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear",
@@ -194,10 +192,10 @@ class Model:
else:
return pred
def inference(self, img0, img1, UHD=False):
def inference(self, img0, img1, scale=1.0):
imgs = torch.cat((img0, img1), 1)
flow, _ = self.flownet(imgs, UHD)
return self.predict(imgs, flow, training=False, UHD=UHD)
flow, _ = self.flownet(imgs, scale)
return self.predict(imgs, flow, training=False)
def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
for param_group in self.optimG.param_groups:

18
Pkgs/rife-cuda/rife.py
View File

@@ -10,6 +10,7 @@ import _thread
import skvideo.io
from queue import Queue, Empty
import shutil
import base64
warnings.filterwarnings("ignore")
abspath = os.path.abspath(__file__)
@@ -29,6 +30,7 @@ parser.add_argument('--rbuffer', dest='rbuffer', type=int, default=200)
parser.add_argument('--wthreads', dest='wthreads', type=int, default=4)
parser.add_argument('--fp16', dest='fp16', action='store_true', help='half-precision mode')
parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video')
parser.add_argument('--scale', dest='scale', type=float, default=1.0, help='Try scale=0.5 for 4k video')
parser.add_argument('--exp', dest='exp', type=int, default=1)
args = parser.parse_args()
assert (not args.input is None)
@@ -92,19 +94,20 @@ def clear_write_buffer(user_args, write_buffer, thread_id):
frameNum = item[0]
img = item[1]
print('[T{}] => {:0>8d}.{}'.format(thread_id, frameNum, args.imgformat))
#imgBytes = base64.b64encode(cv2.imencode(f'.{args.imgformat}', img[:, :, ::-1], [cv2.IMWRITE_PNG_COMPRESSION, 2])[1].tostring())
#print(f"{frameNum:08}:"+ imgBytes.decode('utf-8') + "\n\n\n\n")
cv2.imwrite('{}/{:0>8d}.{}'.format(interp_output_path, frameNum, args.imgformat), img[:, :, ::-1], [cv2.IMWRITE_PNG_COMPRESSION, 2])
def build_read_buffer(user_args, read_buffer, videogen):
for frame in videogen:
if not user_args.input is None:
#print("Loading input frame " + str(frame))
frame = cv2.imread(os.path.join(user_args.input, frame))[:, :, ::-1].copy()
read_buffer.put(frame)
read_buffer.put(None)
def make_inference(I0, I1, exp):
global model
middle = model.inference(I0, I1, args.UHD)
middle = model.inference(I0, I1, args.scale)
if exp == 1:
return [middle]
first_half = make_inference(I0, middle, exp=exp - 1)
@@ -117,13 +120,10 @@ def pad_image(img):
else:
return F.pad(img, padding)
if args.UHD:
print("UHD mode enabled.")
ph = ((h - 1) // 64 + 1) * 64
pw = ((w - 1) // 64 + 1) * 64
else:
ph = ((h - 1) // 32 + 1) * 32
pw = ((w - 1) // 32 + 1) * 32
print(f"Scale: {args.scale}")
tmp = max(32, int(32 / args.scale))
ph = ((h - 1) // tmp + 1) * tmp
pw = ((w - 1) // tmp + 1) * tmp
padding = (0, pw - w, 0, ph - h)
write_buffer = Queue(maxsize=args.rbuffer)