Покадровое снижение напряжения видео с помощью MFDNet на Python ⇐ Python

[Anonymous]

Как уже упоминалось в вопросе CodeReview, я пытаюсь изменить код для обработки покадрового удаления полос дождя в видео. В этом коде используется пакет FFmpeg.

Код: Выделить всё

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)

Давайте сосредоточимся на части цикла while:
Версия приведенного выше фрагмента кода может быть выполнена без ошибок. На следующем шаге я пытаюсь следовать ответу 301_Moved_Permanency, чтобы использовать torch.save. Таким образом, содержимое цикла while выглядит следующим образом:

Код: Выделить всё

        # 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)

Произошла ошибка нехватки памяти со следующим сообщением:

torch.OutOfMemoryError: CUDA не хватает памяти. Пытался выделить 676,00 МБ. Графический процессор 0 имеет общую емкость 23,99 ГиБ, из которых 0 байт свободно. Из выделенной памяти 84,09 ГиБ выделяется PyTorch, а 1,21 ГиБ зарезервировано PyTorch, но нераспределено.

Для диагностики ошибки я удалил последние два строки кода:

Код: Выделить всё

        # 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)

Ошибка нехватки памяти все равно произошла. Для меня это странно. Насколько я понимаю код исполняемой версии, строка «восстановленный = восстановленный.cpu().detach().numpy()» предназначена для передачи восстановленных данных из памяти графического процессора в основную память, а затем преобразования их в нулевой формат. Почему я удалил эту строку кода, а затем произошла ошибка нехватки памяти?
Спецификация аппаратного и программного обеспечения, которую я использовал, следующая:

• ЦП: Intel(R) Core(TM) i9-12900K 12-го поколения, 3,20 ГГц
• ОЗУ: 128 ГБ (128 ГБ свободного места)
• Видеокарта: NVIDIA GeForce RTX 4090
• ОС: Windows 11 Pro 22H2, сборка ОС: 22621.4317
• Версия Pytorch:

  Код: Выделить всё

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

Подробнее здесь: https://stackoverflow.com/questions/791 ... -in-python