Yuv inference ready for sr generation

Author: RobHam99

basicsr/archs/spynet_arch.py

@@ -37,7 +37,6 @@ def __init__(self, load_path=None):
         super(SpyNet, self).__init__()
         self.basic_module = nn.ModuleList([BasicModule() for _ in range(6)])
         if load_path:
-            print('KEYS',torch.load(load_path).keys())
             self.load_state_dict(torch.load(load_path, map_location=lambda storage, loc: storage)['params'], strict=True)
 
         self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))

bvi-aom

@@ -0,0 +1,3 @@
+python3 yuv_swinir.py --input /mnt/e/datasets/bvi-aom/ds-yuv/BFireS21Mitch_1920x1088_24fps_10bit_420.yuv --output /mnt/e/datasets/bvi-aom/excluded_sequences/mitchx2.yuv --num_frames 60 --width 1920 --height 1088 --scale 2
+python3 yuv_swinir.py --input /mnt/e/datasets/bvi-aom/ds-yuv/CFireS21Mitch_960x544_24fps_10bit_420.yuv --output /mnt/e/datasets/bvi-aom/excluded_sequences/mitchx4.yuv --num_frames 60 --width 960 --height 544 --scale 4
+python3 yuv_swinir.py --input /mnt/e/datasets/bvi-aom/ds-yuv/BFireS21Mitch_480x272_24fps_10bit_420.yuv --output /mnt/e/datasets/bvi-aom/excluded_sequences/mitchx8.yuv --num_frames 60 --width 480 --height 272 --scale 8

bvi-sr

@@ -4,6 +4,7 @@
 import os
 import shutil
 import torch
+from tqdm import tqdm
 
 from basicsr.archs.basicvsrpp_arch import BasicVSRPlusPlus
 from basicsr.data.data_util import read_img_seq
@@ -18,52 +19,51 @@ def inference(frames_tensor, model, save_path):
     outputs = outputs.squeeze()
     outputs = outputs.permute(0, 2, 3, 1)
     outputs = outputs.cpu()
-    print("Output video shape: ", outputs.shape)
     rgb_to_yuv420p10bit(outputs, save_path)
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model_path', type=str, default='experiments/pretrained_models/basicvsr_plusplus_c64n7_8x1_600k_reds4_20210217-db622b2f.pth')
-    parser.add_argument(
-        '--input_path', type=str, default='', help='input yuv video')
-    parser.add_argument('--save_path', type=str, default='results/BasicVSRPP', help='save image path')
+    parser.add_argument('--input', type=str, default='', help='input yuv video')
+    parser.add_argument('--output', type=str, default='results/BasicVSRPP', help='save image path')
+    parser.add_argument('--num_frames', type=int, default=60, help='Number of frames to process')
+    parser.add_argument('--width', type=int, default=960, help='Width of the video')
+    parser.add_argument('--height', type=int, default=544, help='Height of the video')
     parser.add_argument('--interval', type=int, default=15, help='interval size')
     args = parser.parse_args()
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+    model_path = '/home/sk24938/source/sr/BasicSR/experiments/pretrained_models/basic_vsr_pp/basicvsr_plusplus_reds4.pth'
+
     # set up model
     model = BasicVSRPlusPlus(mid_channels=64, num_blocks=7, spynet_path='/home/sk24938/source/sr/BasicSR/experiments/pretrained_models/spynet_20210409-c6c1bd09.pth')
-    chkpt = torch.load(args.model_path)
-    print(chkpt.keys())
-    model.load_state_dict(torch.load(args.model_path)['state_dict'], strict=True)
+    model.load_state_dict(torch.load(model_path, weights_only=True)['state_dict'], strict=True)
     model.eval()
     model = model.to(device)
 
     # want to process yuv input frames
     # load yuv frames in rgb format
     # convert to tensor
     frames_np = load_yuv_frames(
-        video_file_path=args.input_path,
+        video_file_path=args.input,
         start_idx=0,
-        num_frames=12,
-        width=256,
-        height=256,
+        num_frames=args.num_frames,
+        width=args.width,
+        height=args.height,
         bit_depth=10,
         pixel_format='yuv420p'
         )
     frames_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2).float()
 
     # load data and inference
-    num_frames = len(frames_tensor)
-    if num_frames <= args.interval:  # too many images may cause CUDA out of memory
+    if args.num_frames <= args.interval:  # too many images may cause CUDA out of memory
         frames_tensor = frames_tensor.unsqueeze(0).to(device)
-        inference(frames_tensor, model, args.save_path)
+        inference(frames_tensor, model, args.output)
     else:
-        for idx in range(0, num_frames, args.interval):
-            interval = min(args.interval, num_frames - idx)
+        for idx in tqdm(range(0, args.num_frames, args.interval), desc='BasicVSR++'):
+            interval = min(args.interval, args.num_frames - idx)
             frames_tensor = frames_tensor.unsqueeze(0).to(device)
-            inference(frames_tensor, model, args.save_path)
+            inference(frames_tensor, model, args.output)
 
 
 if __name__ == '__main__':

inference/triple_run.sh

@@ -2,6 +2,7 @@
 from yuv_utils import *
 import cv2
 import argparse
+from tqdm import tqdm
 
 
 def write_yuv_file(y, u, v, output_file_path):
@@ -14,8 +15,6 @@ def write_yuv_file(y, u, v, output_file_path):
             f.write(u[i].tobytes())
             f.write(v[i].tobytes())
 
-    print(f"YUV file written to {output_file_path}")
-
 
 def rescale_frame(frame, scaling_factor=2, bit_depth=10, method='bicubic'):
     interpolation_methods = {
@@ -88,7 +87,7 @@ def main():
     y_arr = []
     u_arr = []
     v_arr = []
-    for i, frame in enumerate(yuv_frame):
+    for i, frame in tqdm(enumerate(yuv_frame), desc=f'{args.method}'):
         y, u, v = rescale_frame(
             frame=frame,
             scaling_factor=args.scale,

inference/yuv_basic_vsrpp.py

@@ -19,7 +19,7 @@ def main():
     parser.add_argument('--width', type=int, default=960, help='Width of the video')
     parser.add_argument('--height', type=int, default=544, help='Height of the video')
     parser.add_argument('--scale', type=int, default=4, help='Scaling factor')
-
+    # Before inference
     args = parser.parse_args()
 
     if args.scale == 4:
@@ -42,7 +42,7 @@ def main():
         img_range=1.0,
         rgb_mean=(0.4488, 0.4371, 0.4040)
     )
-    model.load_state_dict(torch.load(model_path)['params'], strict=True)
+    model.load_state_dict(torch.load(model_path, weights_only=True)['params'], strict=True)
     model.eval()
     model = model.to(device)
 
@@ -64,25 +64,23 @@ def main():
     frames_tensor = frames_tensor.to(device)
 
     frame_list = []
-    for i in tqdm(range(frames_tensor.shape[0]), desc='Processing frames'):
+    for i in tqdm(range(frames_tensor.shape[0]), desc='EDSR'):
         frame = frames_tensor[i, :, :, :]
         # inference
         try:
             with torch.no_grad():
                 output = model(frame)
+
         except Exception as error:
             print('Error', error, i)
         else:
             # save image
             output = output.data.squeeze().cpu()
             frame_list.append(output)
 
-    print('Saving video...')
     video = torch.stack(frame_list, dim=0)
     video = video.permute(0, 2, 3, 1)
-    print('Upsampled video shape: ', video.shape)
     rgb_to_yuv420p10bit(video, args.output)
-    print('Done!')
 
 
 if __name__ == '__main__':

inference/yuv_conventional.py

@@ -54,7 +54,7 @@ def main():
         num_reconstruct_block=10,
         hr_in=False,
     )
-    model.load_state_dict(torch.load(model_path)['params'], strict=True)
+    model.load_state_dict(torch.load(model_path, weights_only=True)['params'], strict=True)
     model.eval()
     model = model.to(device)
 
@@ -78,20 +78,16 @@ def main():
     frames_tensor = pad_frames(frames_tensor, pad_size)
 
     frame_list = []
-    for idx in tqdm(range(0 + pad_size, args.num_frames + pad_size, 1), desc='Processing frames'):
+    for idx in tqdm(range(0 + pad_size, args.num_frames + pad_size, 1), desc='EDVR'):
         start_idx = idx - pad_size
         end_idx = idx + pad_size
         frames_tensor_chunk = frames_tensor[:, start_idx:end_idx+1, :, :, :].to(device)
         output = inference(frames_tensor_chunk, model, args.output)
         frame_list.append(output)
         torch.cuda.empty_cache()
 
-    print('Saving video...')
     video = torch.stack(frame_list, dim=0)
-    print('Upsampled video shape: ', video.shape)
     rgb_to_yuv420p10bit(video, args.output)
-    print('Done!')
-
 
 if __name__ == '__main__':
     main()

inference/yuv_edsr.py

@@ -16,23 +16,22 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--input', type=str, default='datasets/swin.yuv', help='input yuv file')
     parser.add_argument('--output', type=str, default='results/swin.yuv', help='output yuv file')
-    parser.add_argument('--num_frames', type=int, default=64, help='number of frames to process')
+    parser.add_argument('--num_frames', type=int, default=60, help='number of frames to process')
     parser.add_argument('--width', type=int, default=960, help='width of the frames')
     parser.add_argument('--height', type=int, default=544, help='height of the frames')
-    parser.add_argument('--patch_size', type=int, default=64, help='training patch size')
     parser.add_argument('--scale', type=int, default=4, help='scale factor: 2, 4, 8')
+    parser.add_argument('--patch_size', type=int, default=64, help='patch size')
     args = parser.parse_args()
 
-    model_path = f"experiments/pretrained_models/swinir/001_classicalSR_DF2K_s64w8_SwinIR-M_x{args.scale}.pth"
+    model_path = f"/home/sk24938/source/sr/BasicSR/experiments/pretrained_models/swinir/001_classicalSR_DF2K_s64w8_SwinIR-M_x{args.scale}.pth"
+    window_size = 8
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     # set up model
     model = define_model(args, model_path)
     model.eval()
     model = model.to(device)
 
-    window_size = 8
-
     frames_np = load_yuv_frames(
         video_file_path=args.input,
         start_idx=0,
@@ -45,7 +44,7 @@ def main():
     frames_tensor = torch.from_numpy(frames_np).permute(0, 3, 1, 2).float().unsqueeze(0)
 
     frames_list = []
-    for idx in tqdm(range(frames_tensor.shape[1]), desc='Processing frames', leave=False):
+    for idx in tqdm(range(frames_tensor.shape[1]), desc='SwinIR'):
         frame = frames_tensor[:, idx].to(device)
         # inference
         with torch.no_grad():
@@ -67,12 +66,9 @@ def main():
         output = output.data.squeeze().float().cpu()
         frames_list.append(output)
 
-    tqdm.write('Saving video')
     video = torch.stack(frames_list, dim=0)
     video = video.permute(0, 2, 3, 1)
-    tqdm.write(f'Upsampled video shape: {video.shape}')
     rgb_to_yuv420p10bit(video, args.output)
-    tqdm.write('Done!')
 
 def define_model(args, model_path):
     # 001 classical image sr
@@ -90,7 +86,7 @@ def define_model(args, model_path):
         resi_connection='1conv')
 
 
-    loadnet = torch.load(model_path)
+    loadnet = torch.load(model_path, weights_only=True)
     if 'params_ema' in loadnet:
         keyname = 'params_ema'
     else: