Recherche avancée

Sur d’autres sites (11763)

• Video Frame-by-Frame Deraining with MFDNet in Python

  28 octobre 2024, par JimmyHu

  As this CodeReview question mentioned, I am trying to modify the code to process frame-by-frame rain streaks removal in a video. FFmpeg package is used in this code.

  


  import argparse
import os
import time

import cv2
import ffmpeg
import numpy as np
import torch
from skimage import img_as_ubyte
from torch.utils.data import DataLoader
from tqdm import tqdm

import utils
from data_RGB import get_test_data
from MFDNet import HPCNet as mfdnet


def process_video_frame_by_frame(input_file, output_file, model_restoration):
    """
    Decodes a video frame by frame, processes each frame,
    and re-encodes to a new video.

    Args:
        input_file: Path to the input video file.
        output_file: Path to the output video file.
    """
    try:
        # Probe for video information
        probe = ffmpeg.probe(input_file)
        video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        width = int(video_stream['width'])
        height = int(video_stream['height'])

        # Input
        process1 = (
            ffmpeg
            .input(input_file)
            .output('pipe:', format='rawvideo', pix_fmt='rgb24')
            .run_async(pipe_stdout=True)
        )

        # Output
        process2 = (
            ffmpeg
            .input('pipe:', format='rawvideo', pix_fmt='rgb24', s='{}x{}'.format(width, height))
            .output(output_file, vcodec='libx264', pix_fmt='yuv420p')
            .overwrite_output()
            .run_async(pipe_stdin=True)
        )

        # Process frame (deraining processing)
        while in_bytes := process1.stdout.read(width * height * 3):
            in_frame = torch.frombuffer(in_bytes, dtype=torch.uint8).float().reshape((1, 3, width, height))
            restored = model_restoration(torch.div(in_frame, 255).to(device='cuda'))
            restored = torch.clamp(restored[0], 0, 1)
            restored = restored.cpu().detach().numpy()
            restored *= 255
            out_frame = restored
            np.reshape(out_frame, (3, width, height))

            # Encode and write the frame
            process2.stdin.write(
                out_frame
                .astype(np.uint8)
                .tobytes()
            )
            
        # Close streams
        process1.stdout.close()
        process2.stdin.close()
        process1.wait()
        process2.wait()

    except ffmpeg.Error as e:
        print('stdout:', e.stdout.decode('utf8'))
        print('stderr:', e.stderr.decode('utf8'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Image Deraining using MPRNet')

    parser.add_argument('--weights', default='./checkpoints/checkpoints_mfd.pth', type=str,
                        help='Path to weights')
    parser.add_argument('--gpus', default='0', type=str, help='CUDA_VISIBLE_DEVICES')

    args = parser.parse_args()

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus

    model_restoration = mfdnet()
    utils.load_checkpoint(model_restoration, args.weights)
    print("===>Testing using weights: ", args.weights)

    model_restoration.eval().cuda()
    
    input_video = "Input_video.mp4"
    output_video = 'output_video.mp4'

    process_video_frame_by_frame(input_video, output_video, model_restoration)


  


  Let's focus on the while loop part :

  


  The version of the code snippet above can be executed without error. In the next step, I am trying to follow 301_Moved_Permanently's answer to make the usage of torch.save. Therefore, the contents of while loop comes as the following code :

  


          # Process frame (deraining processing)
        while in_bytes := process1.stdout.read(width * height * 3):
            in_frame = torch.frombuffer(in_bytes, dtype=torch.uint8).float().reshape((1, 3, width, height))
            restored = model_restoration(torch.div(in_frame, 255).to(device='cuda'))
            restored = torch.clamp(restored[0], 0, 1)
            out_frame = torch.mul(restored.cpu().detach(), 255).reshape(3, width, height).byte()
            torch.save(out_frame, process2.stdin)


  


  Out of memory error happened with the following message :

  


  


  torch.OutOfMemoryError : CUDA out of memory. Tried to allocate 676.00 MiB. GPU 0 has a total capacity of 23.99 GiB of which 0 bytes is free. Of the allocated memory 84.09 GiB is allocated by PyTorch, and 1.21 GiB is reserved by PyTorch but unallocated.

  


  


  To diagnostics the error, I removed the last two lines of code :

  


          # Process frame (deraining processing)
        while in_bytes := process1.stdout.read(width * height * 3):
            in_frame = torch.frombuffer(in_bytes, dtype=torch.uint8).float().reshape((1, 3, width, height))
            restored = model_restoration(torch.div(in_frame, 255).to(device='cuda'))
            restored = torch.clamp(restored[0], 0, 1)


  


  The out of memory error still happened. This is weird to me. My understanding of the executable version code, the line restored = restored.cpu().detach().numpy() is to transfer the restored data in GPU memory to main memory and then convert it to numpy format. Why I remove this line of code then out of memory error happened ?

  


  The hardware and software specification I used is as follows :

  


  


  • CPU : 12th Gen Intel(R) Core(TM) i9-12900K 3.20 GHz

    


  


  • RAM : 128 GB (128 GB usable)

    


  


  • Graphic card : NVIDIA GeForce RTX 4090

    


  


  • OS : Windows 11 Pro 22H2, OS build : 22621.4317

    


  


  • Pytorch version :

    


    > python -c "import torch; print(torch.__version__)"
2.5.0+cu124


    


  


  


• avcodec/mpegvideo_enc : Don't do unnecessary work for AMV

  25 février, par Andreas Rheinhardt

  avcodec/mpegvideo_enc : Don't do unnecessary work for AMV
  Up until now, the initialization of AMV's matrices happened
  
  after the initialization for MJPEG matrices, overwriting
  
  the earlier data. This commit changes this.
  Signed-off-by : Andreas Rheinhardt <andreas.rheinhardt@outlook.com>
  

  • [DH] libavcodec/mpegvideo_enc.c

• avcodec/get_bits : Remove GetBitContext.buffer_end

  4 juillet, par Andreas Rheinhardt

  avcodec/get_bits : Remove GetBitContext.buffer_end
  It is unused. Furthermore, this automatically fixes the issue
  
  that init_get_bits() failure would lead to NULL + 0 (when
  
  setting buffer_end) which is UB before C23. This happened
  
  in the fic-avi and fic-avi-skip_cursor FATE-tests.
  This saved 7296B of .text here.
  Signed-off-by : Andreas Rheinhardt <andreas.rheinhardt@outlook.com>
  

  • [DH] libavcodec/get_bits.h