chain_img_processor/video.py

import roop.globals

from threading import Thread
from chain_img_processor import ChainImgProcessor

class ThreadWithReturnValue(Thread):

    def __init__(self, group=None, target=None, name=None,
                 args=(), kwargs={}, Verbose=None):
        Thread.__init__(self, group, target, name, args, kwargs)
        self._return = None

    def run(self):
        if self._target is not None:
            self._return = self._target(*self._args,
                                        **self._kwargs)

    def join(self, *args):
        Thread.join(self, *args)
        return self._return


# in beta
class ChainVideoProcessor(ChainImgProcessor):
    def __init__(self):
        ChainImgProcessor.__init__(self)

        self.video_save_codec = "libx264"
        self.video_save_crf = 14

    def init_with_plugins(self):
        self.init_plugins(["core","core_video"])
        self.display_init_info()

        init_on_start_arr = self.init_on_start.split(",")
        for proc_id in init_on_start_arr:
            self.init_processor(proc_id)

    def run_video_chain(self, source_video, target_video, fps, threads:int = 1, chain = None, params_frame_gen_func = None, video_audio = None):
        import cv2
        from tqdm import tqdm
        from chain_img_processor.ffmpeg_writer import FFMPEG_VideoWriter # ffmpeg install needed

        cap = cv2.VideoCapture(source_video)
        # width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        # height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        # first frame do manually - because upscale may happen, we need to estimate width/height
        ret, frame = cap.read()
        if params_frame_gen_func is not None:
            params = params_frame_gen_func(self, frame)
        else:
            params = {}
        params["original_frame"] = frame
        frame_processed, params = self.run_chain(frame,params,chain)
        height, width, channels = frame_processed.shape

        self.fill_processors_for_thread_chains(threads,chain)
        #print(self.processors_objects)
        #import threading
        #locks:list[threading.Lock] = []
        locks: list[bool] = []
        for i in range(threads):
            #locks.append(threading.Lock())
            locks.append(False)

        temp = []
        with FFMPEG_VideoWriter(target_video, (width, height), fps, codec=roop.globals.video_encoder, crf=roop.globals.video_quality, audiofile=video_audio) as output_video_ff:
            with tqdm(total=frame_count, desc='Processing', unit="frame", dynamic_ncols=True,
                      bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]') as progress:

                # do first frame
                output_video_ff.write_frame(frame_processed)
                progress.update(1) #
                cnt_frames = 0

                # do rest frames
                while True:
                    # getting frame
                    ret, frame = cap.read()

                    if not ret:
                        break
                    cnt_frames+=1
                    thread_ind = cnt_frames % threads
                    # we are having an array of length %gpu_threads%, running in parallel
                    # so if array is equal or longer than gpu threads, waiting
                    #while len(temp) >= threads:
                    while locks[thread_ind]:
                        #print('WAIT', thread_ind)
                        # we are order dependent, so we are forced to wait for first element to finish. When finished removing thread from the list
                        frame_processed, params = temp.pop(0).join()
                        locks[params["_thread_index"]] = False
                        #print('OFF',cnt_frames,locks[params["_thread_index"]],locks)
                        # writing into output
                        output_video_ff.write_frame(frame_processed)
                        # updating the status
                        progress.update(1)

                    # calc params for frame
                    if params_frame_gen_func is not None:
                        params = params_frame_gen_func(self,frame)
                    else:
                        params = {}

                    # adding new frame to the list and starting it
                    locks[thread_ind] = True
                    #print('ON', cnt_frames, thread_ind, locks)
                    params["original_frame"] = frame
                    temp.append(
                        ThreadWithReturnValue(target=self.run_chain, args=(frame, params, chain, thread_ind)))
                    temp[-1].start()

                while len(temp) > 0:
                    # we are order dependent, so we are forced to wait for first element to finish. When finished removing thread from the list
                    frame_processed, params = temp.pop(0).join()
                    locks[params["_thread_index"]] = False
                    # writing into output
                    output_video_ff.write_frame(frame_processed)

                    progress.update(1)

                #print("FINAL", locks)

_video_processor:ChainVideoProcessor = None
def get_single_video_processor() -> ChainVideoProcessor:
    global _video_processor
    if _video_processor is None:
        _video_processor = ChainVideoProcessor()
        _video_processor.init_with_plugins()
    return _video_processor