From 223f4221fc7547f6df314a8fadfa3da0a0582544 Mon Sep 17 00:00:00 2001 From: N00MKRAD Date: Sat, 27 Feb 2021 16:24:42 +0100 Subject: [PATCH] Updated RIFE-CUDA (scale arg, img piping tests) --- Pkgs/rife-cuda/model/IFNet_HD.py | 11 +++++------ Pkgs/rife-cuda/model/IFNet_HDv2.py | 27 +++++++++++++-------------- Pkgs/rife-cuda/model/RIFE_HD.py | 10 ++++------ Pkgs/rife-cuda/model/RIFE_HDv2.py | 10 ++++------ Pkgs/rife-cuda/rife.py | 18 +++++++++--------- 5 files changed, 35 insertions(+), 41 deletions(-) diff --git a/Pkgs/rife-cuda/model/IFNet_HD.py b/Pkgs/rife-cuda/model/IFNet_HD.py index fe315b2..6975679 100644 --- a/Pkgs/rife-cuda/model/IFNet_HD.py +++ b/Pkgs/rife-cuda/model/IFNet_HD.py @@ -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__': diff --git a/Pkgs/rife-cuda/model/IFNet_HDv2.py b/Pkgs/rife-cuda/model/IFNet_HDv2.py index f9b18cf..c71886c 100644 --- a/Pkgs/rife-cuda/model/IFNet_HDv2.py +++ b/Pkgs/rife-cuda/model/IFNet_HDv2.py @@ -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] diff --git a/Pkgs/rife-cuda/model/RIFE_HD.py b/Pkgs/rife-cuda/model/RIFE_HD.py index b96576f..461dc1b 100644 --- a/Pkgs/rife-cuda/model/RIFE_HD.py +++ b/Pkgs/rife-cuda/model/RIFE_HD.py @@ -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: diff --git a/Pkgs/rife-cuda/model/RIFE_HDv2.py b/Pkgs/rife-cuda/model/RIFE_HDv2.py index 9f19ae2..2de119d 100644 --- a/Pkgs/rife-cuda/model/RIFE_HDv2.py +++ b/Pkgs/rife-cuda/model/RIFE_HDv2.py @@ -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: diff --git a/Pkgs/rife-cuda/rife.py b/Pkgs/rife-cuda/rife.py index 93141d9..dd531f6 100644 --- a/Pkgs/rife-cuda/rife.py +++ b/Pkgs/rife-cuda/rife.py @@ -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)