Python视频处理之噪声矩阵与并行计算

import os
import cv2
import numpy as np
from concurrent.futures import ProcessPoolExecutor

class NoiseWaveClass:
    def __init__(self, tmp_out_video, out_video, configs):
        self.tmp_out_video = tmp_out_video
        self.out_video = out_video
        self.noise_mean = int(configs['noise_mean'])
        self.noise_var = int(configs['noise_var'])

    @staticmethod
    def process_frame(frame, noise_mean, noise_var):
        row, col, ch = frame.shape
        sigma = noise_var ** 0.5
        gauss = np.random.normal(noise_mean, sigma, (row, col, ch)).astype('uint8')
        noisy = cv2.add(frame, gauss)
        return noisy

    # 定义一个可以被序列化的函数来包装参数传递
    @staticmethod
    def process_frame_with_params(args):
        frame, noise_mean, noise_var = args
        return NoiseWaveClass.process_frame(frame, noise_mean, noise_var)

    def process_noisewave(self):
        try:
            if os.path.exists(self.tmp_out_video):
                os.remove(self.tmp_out_video)
            os.rename(self.out_video, self.tmp_out_video)

            cap = cv2.VideoCapture(self.tmp_out_video)
            fps = int(cap.get(cv2.CAP_PROP_FPS))
            frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')

            out = cv2.VideoWriter(self.out_video, fourcc, fps, (frame_width, frame_height))

            def process_frames_in_parallel(frames):
                with ProcessPoolExecutor() as executor:
                    # 使用静态方法 process_frame_with_params 来包装参数传递
                    args = [(frame, self.noise_mean, self.noise_var) for frame in frames]
                    processed_frames = list(executor.map(NoiseWaveClass.process_frame_with_params, args))
                return processed_frames

            frames = []
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                frames.append(frame)

            processed_frames = process_frames_in_parallel(frames)

            for frame in processed_frames:
                out.write(frame)

        except Exception as e:
            print(f"Error processing video: {e}")
        finally:
            cap.release()
            out.release()
            cv2.destroyAllWindows()
            if os.path.exists(self.tmp_out_video):
                os.remove(self.tmp_out_video)