Spaces:
Running
Running
File size: 5,495 Bytes
0c87db7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
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
|