def _get_writer(self, frame_dims: Tuple[int]) -> Generator[None, np.ndarray, None]: """ Add the requested encoding options and return the writer. Parameters ---------- frame_dims: tuple The (rows, colums) shape of the input image Returns ------- generator The imageio ffmpeg writer """ audio_codec = self._audio_codec audio_path = None if audio_codec is None else self._source_video logger.debug("writer config: %s, audio_path: '%s'", self.config, audio_path) retval = im_ffm.write_frames(self._output_filename, size=(frame_dims[1], frame_dims[0]), fps=self._video_fps, quality=None, codec=self.config["codec"], macro_block_size=8, ffmpeg_log_level="error", ffmpeg_timeout=10, output_params=self._output_params, audio_path=audio_path, audio_codec=audio_codec) logger.debug("FFMPEG Writer created: %s", retval) retval.send(None) return retval
def test_write1(): for n in (1, 9, 14, 279, 280, 281): # Prepare for writing gen = imageio_ffmpeg.write_frames(test_file2, (64, 64)) assert isinstance(gen, types.GeneratorType) gen.send(None) # seed # Write n frames for i in range(n): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * 3) gen.send(data) gen.close() # Check that number of frames is correct nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == n # Check again by actually reading gen2 = imageio_ffmpeg.read_frames(test_file2) gen2.__next__() # == meta count = 0 for frame in gen2: count += 1 assert count == n
def processvideo(file): print('processing ' + file) cap = cv2.VideoCapture(file) # cap.get(3) width and cap.get(4) is heigh if int(args.framelimit) >= 0: if args.outputfile.startswith('udp') or args.outputfile.startswith( 'rtsp'): print("Exporting to stream: " + args.outputfile) writer = imageio.write_frames(args.outputfile, (int(cap.get(3)), int(cap.get(4))), output_params=['-f', 'mpegts']) else: print("Exporting to file " + args.outputfile) writer = imageio.write_frames(args.outputfile, (int(cap.get(3)), int(cap.get(4)))) writer.send(None) frame_counter = 0 while (cap.isOpened()): frame_counter = frame_counter + 1 # check limits of processed frames if int(args.framelimit) > 0 and frame_counter > int( args.framestart) + int(args.framelimit): print('Processed ' + args.framelimit + 'ending...') break # read input frame ret, frame = cap.read() # throttle control if frame_counter % int(args.fpsthrottle) != 0: continue # check frame start if int(args.framestart) > 0 and frame_counter < int(args.framestart): break if ret == True: if not frame is None: print('Detecting objects in frame ' + str(frame_counter), end="\r") image = detect(frame) #writer.send(frame) writer.send(image) else: print('Frame error in frame ' + str(frame_counter)) else: break cap.release() writer.close()
def _write_frames(pixfmt, bpp, tout): gen = imageio_ffmpeg.write_frames(test_file2, (2048, 2048), pix_fmt_in=pixfmt, ffmpeg_timeout=tout) gen.send(None) # seed for i in range(n): data = (255 * np.random.rand(2048 * 2048 * bpp)).astype(np.uint8) data = bytes(data) gen.send(data) gen.close()
def encode_video(video: np.ndarray, output_path: str, fps: float, bitrate: str = "0", crf: int = 20) -> str: """Encode a video with vp9 codec via imageio-ffmpeg Parameters ---------- video : np.ndarray Video to be encoded output_path : str Desired output path for encoded video fps : float Desired frame rate for encoded video pix_fmt_out : str Desired pixel format for output encoded video, by default "yuv420p" bitrate : str, optional Desired bitrate of output, by default "0". The default *MUST* be zero in order to encode in constant quality mode. Other values will result in constrained quality mode. crf : int, optional Desired perceptual quality of output, by default 20. Value can be from 0 - 63. Lower values mean better quality (but bigger video sizes). Returns ------- str Output path of the encoded video """ # ffmpeg expects video shape in terms of: (width, height) video_shape = (video[0].shape[1], video[0].shape[0]) writer = mpg.write_frames(output_path, video_shape, pix_fmt_in="gray8", pix_fmt_out="yuv420p", codec="libvpx-vp9", fps=fps, bitrate=bitrate, output_params=["-crf", str(crf)]) writer.send(None) # Seed ffmpeg-imageio writer generator for frame in video: writer.send(frame) writer.close() return output_path
def test_write_wmv(): # Switch to MS friendly codec when writing .wmv files for ext, codec in [("", "h264"), (".wmv", "msmpeg4")]: fname = test_file2 + ext gen = imageio_ffmpeg.write_frames(fname, (64, 64)) gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * 3) gen.send(data) gen.close() # meta = imageio_ffmpeg.read_frames(fname).__next__() assert meta["codec"].startswith(codec)
async def write_video_to_stream(path_to_video, path_to_socket): video_reader = imageio.get_reader(path_to_video) sock = await create_socket(path_to_socket) _, socket_writer = await asyncio.open_unix_connection(sock=sock) img_writer = None for frame in video_reader: if img_writer is None: img_writer = imageio_ffmpeg.write_frames(socket_writer, size=(frame.shape[1], frame.shape[0])) img_writer.send(None) img_writer.send(frame) await socket_writer.drain() img_writer.close() socket_writer.close()
def test_write_close(): for i in range(N): pids0 = get_ffmpeg_pids() w = imageio_ffmpeg.write_frames(test_file2, (64, 64)) pids1 = get_ffmpeg_pids().difference( pids0) # generator has not started w.send(None) w.send(b"x" * 64 * 64 * 3) pids2 = get_ffmpeg_pids().difference(pids0) # now ffmpeg is running w.close() pids3 = get_ffmpeg_pids().difference(pids0) # now its not assert len(pids1) == 0 assert len(pids2) == 1 assert len(pids3) == 0
def test_write_pix_fmt_in(): sizes = [] for pixfmt, bpp in [("gray", 1), ("rgb24", 3), ("rgba", 4)]: # Prepare for writing gen = imageio_ffmpeg.write_frames(test_file2, (64, 64), pix_fmt_in=pixfmt) gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * bpp) gen.send(data) gen.close() with open(test_file2, "rb") as f: sizes.append(len(f.read())) # Check nframes nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == 9
def test_write_quality(): sizes = [] for quality in [2, 5, 9]: # Prepare for writing gen = imageio_ffmpeg.write_frames(test_file2, (64, 64), quality=quality) gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * 3) gen.send(data) gen.close() with open(test_file2, "rb") as f: sizes.append(len(f.read())) # Check nframes nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == 9 assert sizes[0] < sizes[1] < sizes[2]
def test_write_pix_fmt_out(): sizes = [] for pixfmt in ["gray", "yuv420p"]: # Prepare for writing gen = imageio_ffmpeg.write_frames(test_file2, (64, 64), pix_fmt_out=pixfmt) gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * 3) gen.send(data) gen.close() with open(test_file2, "rb") as f: sizes.append(len(f.read())) # Check nframes nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == 9 assert sizes[0] < sizes[1]
def test_write_macro_block_size(): frame_sizes = [] for mbz in [None, 10]: # None is default == 16 # Prepare for writing gen = imageio_ffmpeg.write_frames(test_file2, (40, 50), macro_block_size=mbz) gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 40 * 50 * 3) gen.send(data) gen.close() # Check nframes nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == 9 # Check size meta = imageio_ffmpeg.read_frames(test_file2).__next__() frame_sizes.append(meta["size"]) assert frame_sizes[0] == (48, 64) assert frame_sizes[1] == (40, 50)
def test_write_bitrate(): # Mind that we send uniform images, so the difference is marginal sizes = [] for bitrate in ["1k", "10k", "100k"]: # Prepare for writing gen = imageio_ffmpeg.write_frames(test_file2, (64, 64), bitrate=bitrate) gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * 3) gen.send(data) gen.close() with open(test_file2, "rb") as f: sizes.append(len(f.read())) # Check nframes nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == 9 assert sizes[0] < sizes[1] < sizes[2]
def processvideo(file): cap = cv2.VideoCapture(file) writer = imageio.write_frames(args.outputfile, (int(cap.get(3)), int(cap.get(4)))) writer.send(None) frame_counter = 0 while(cap.isOpened()): frame_counter = frame_counter + 1 ret, frame = cap.read() print('Detecting objects in frame ' + str(frame_counter)) if ret==True: if not frame is None: image = detect(frame) writer.send(frame) else: print('Frame error in frame ' + str(frame_counter)) else: break cap.release() writer.close()
def test_write_audio_path(): # Provide an audio gen = imageio_ffmpeg.write_frames(test_file2, (64, 64), audio_path=test_file3, audio_codec="aac") gen.send(None) # seed for i in range(9): data = bytes([min(255, 100 + i * 10)] * 64 * 64 * 3) gen.send(data) gen.close() # Check nframes nframes, nsecs = imageio_ffmpeg.count_frames_and_secs(test_file2) assert nframes == 9 # Check size meta = imageio_ffmpeg.read_frames(test_file2).__next__() audio_codec = meta["audio_codec"] assert nframes == 9 assert audio_codec == "aac"
def write(self, frame: np.ndarray): assert len(frame.shape) == 3 assert frame.shape[2] == 3 assert frame.dtype == np.uint8 frame_size = (frame.shape[1], frame.shape[0]) if self.out is None: self.path.parent.mkdir(exist_ok=True, parents=True) self.frame_size = frame_size self.out = imageio_ffmpeg.write_frames( str(self.path), self.frame_size, fps=self.fps, codec=self.codec, macro_block_size=1, ffmpeg_log_level='error', ) self.out.send(None) else: assert self.frame_size == frame_size, f"Wrong frame size: should be {self.frame_size}, got {frame_size}" if self.bgr2rgb: frame = frame[:, :, (2, 1, 0)] self.out.send(np.ascontiguousarray(frame))
def write_video_to_stream(path_to_video, rtp_port): video_reader = imageio.get_reader(path_to_video) print(video_reader.get_meta_data()) rtp_url = f"rtp://localhost:{rtp_port}" output_params = ['-f', 'rtp_mpegts'] input_params = ['-re'] img_writer = imageio_ffmpeg.write_frames( rtp_url, size=video_reader.get_meta_data()['size'], output_params=output_params, input_params=input_params, fps=24, codec='mpeg4', bitrate='1024K') img_writer.send(None) for frame in video_reader: try: img_writer.send(frame) except KeyboardInterrupt as e: img_writer.close() break img_writer.close()
def write_video_to_stream(path_to_video, rtp_port): video_reader = imageio.get_reader(path_to_video) print(video_reader.get_meta_data()) rtp_url = f"http://localhost:{rtp_port}" output_params = [ '-f', 'webm', '-listen', "1", '-seekable', '0', '-headers', 'Access-Control-Allow-Origin: http://localhost:3000\r\n'] input_params = ['-re'] img_writer = imageio_ffmpeg.write_frames( rtp_url, size=video_reader.get_meta_data()['size'], output_params=output_params, input_params=input_params, codec='vp9', ) img_writer.send(None) for frame in video_reader: #try: img_writer.send(frame) #except KeyboardInterrupt as e: # print(e) # img_writer.close() # break img_writer.close()
break cap.release() writer.close() # Doing some Object Detection on a video classes = None with open(args.classes, 'r') as f: classes = [line.strip() for line in f.readlines()] COLORS = np.random.uniform(0, 255, size=(len(classes), 3)) if args.input.startswith('rtsp'): cap = cv2.VideoCapture(args.input) if int(args.framelimit) > 0: writer = imageio.write_frames(args.outputfile, (int(cap.get(3)), int(cap.get(4)))) writer.send(None) frame_counter = 0 while(True): if int(args.framelimit) > 0 and frame_counter > int(args.framestart) + int(args.framelimit): writer.close() break if frame_counter % int(args.fpsthrottle) ==0: ret, frame = cap.read() if ret and frame_counter >= int(args.framestart): print('Detecting objects in frame ' + str(frame_counter)) frame = detect(frame) if int(args.framelimit) > 0: writer.send(frame) else:
def style_video(video_path, style, output_name=None, alpha=1.0, preserve_color=False, custom_resolution=None): toTensor = ToTensor() reader = imageio.get_reader(video_path) resolution = None if custom_resolution is None: resolution = reader.get_meta_data()["source_size"] else: resolution = custom_resolution fps = reader.get_meta_data()["fps"] style = transform.resize(style, resolution) if not preserve_color: style_tensor = torch.tensor(np.expand_dims(toTensor(style), 0)).float().to(device) style_tensor = net.encode(style_tensor) frame_number = 0 driving_video = [] for frame_number in tqdm(range(reader.get_meta_data()['nframes'])): try: content = reader.get_next_data() except imageio.core.CannotReadFrameError: break except IndexError: break else: if custom_resolution is not None: content = transform.resize(content, list(resolution)[::-1]) else: content = content.astype(float) / 255 if preserve_color: style_tensor = preserve_color_stable(content, style) style_tensor = torch.tensor( np.expand_dims(toTensor(style_tensor), 0)).float().to(device) style_tensor = net.encode(style_tensor) x = torch.tensor(toTensor(content)).unsqueeze(0).float().to(device) x = net.encode(x) result = adaIN(x, style_tensor) result = (1 - alpha) * x + alpha * result result = test(content, result, alpha, False, False) result = cv2.convertScaleAbs(result * 255) driving_video.append(result) reader.close() if output_name == None: return driving_video writer = imageio_ffmpeg.write_frames(f'{output_name}.mp4', (resolution), fps=fps, macro_block_size=1) writer.send(None) # seed the generator for frame in driving_video: writer.send(frame) writer.close()
prediction = model(in_data) torch.cuda.synchronize() end = time.time() times.append(end - start) # store video frame(s) x = in_data[:, data.N_FRAMES, 0:3, :, :] if 'reproj' in args.type else in_data[:, 0:3, :, :] p = prediction[:, 0:3, :, :] y = target[:, 0:3, :, :] # postprocess frame = (torch.sqrt(torch.clamp(torch.cat((x, p, y), dim=-1), 0, 1)) * 255).to(torch.uint8) frame = frame.transpose_(-3, -1).transpose_(-3, -2) video[args.batch_size*idx:args.batch_size*idx+p.size(0)] = frame.cpu().numpy() # write images to disk? if args.images: img = torch.cat((x, p, y), dim=-1) if args.cmp else p data.write([f'{name}/{name}_{args.batch_size*idx+j:06}.{args.format}' for j in range(frame.size(0))], img.cpu().numpy()) tq.update(args.batch_size) tq.close() print(f'avg inference time (in s):', np.array(times).mean(), 'std:', np.array(times).std()) # write video ffmpeg = imageio_ffmpeg.write_frames(f'{name}/{name}.mp4', (3*sample.shape[-1], sample.shape[-2]), fps=args.fps, quality=5) ffmpeg.send(None) # seed ffmpeg.send(video) ffmpeg.close() print(f'{name}/{name}.mp4 written.') # make sure images were written data.pool.close() data.pool.join()
# %% Write a series of large frames # In earlier versions of imageio-ffmpeg, the ffmpeg process was given a timeout # to complete, but this timeout must be longer for longer movies. The default # is now to wait for ffmpeg. import os import numpy as np import imageio_ffmpeg ims = [ np.random.uniform(0, 255, size=(1000, 1000, 3)).astype(np.uint8) for i in range(10) ] filename = os.path.expanduser("~/Desktop/foo.mp4") w = imageio_ffmpeg.write_frames(filename, (1000, 1000), ffmpeg_timeout=0) w.send(None) for i in range(200): w.send(ims[i % 10]) print(i) w.close() # %% Behavior of KeyboardInterrupt / Ctrl+C import os import imageio_ffmpeg filename = os.path.expanduser("~/.imageio/images/cockatoo.mp4") reader = imageio_ffmpeg.read_frames(filename)
def OpenWriter(cam_params): n_cam = cam_params["n_cam"] folder_name = os.path.join(cam_params["videoFolder"], cam_params["cameraName"]) if cam_params["cameraMake"] == "emu": cam_params["videoFilename"] = "emu" + cam_params["videoFilename"] full_file_name = os.path.join(folder_name, cam_params["videoFilename"]) if not os.path.isdir(folder_name): os.makedirs(folder_name) print('Made directory {}.'.format(folder_name)) # Load defaults pix_fmt_out = cam_params["pixelFormatOutput"] codec = cam_params["codec"] gpu_params = [] # CPU compression if cam_params["gpuID"] == -1: print('Opened: {} using CPU to compress the stream.'.format( full_file_name)) if pix_fmt_out == 'rgb0': pix_fmt_out = 'yuv420p' if cam_params["codec"] == 'h264': codec = 'libx264' elif cam_params["codec"] == 'h265': codec = 'libx265' gpu_params = [ '-r:v', str(cam_params["frameRate"]), '-preset', 'fast', '-tune', 'fastdecode', '-crf', cam_params["quality"], '-bufsize', '20M', '-maxrate', '10M', '-bf:v', '4', '-vsync', '0', ] # GPU compression else: print('Opened: {} using GPU {} to compress the stream.'.format( full_file_name, cam_params["gpuID"])) if cam_params["gpuMake"] == 'nvidia': if cam_params["codec"] == 'h264': codec = 'h264_nvenc' elif cam_params["codec"] == 'h265': codec = 'hevc_nvenc' gpu_params = [ '-r:v', str(cam_params["frameRate"] ), # important to play nice with vsync '0' '-preset', 'fast', # set to 'fast', 'llhp', or 'llhq' for h264 or hevc '-qp', cam_params["quality"], '-bf:v', '0', '-vsync', '0', '-2pass', '0', '-gpu', str(cam_params["gpuID"]), ] elif cam_params["gpuMake"] == 'amd': if pix_fmt_out == 'rgb0': pix_fmt_out = 'yuv420p' if cam_params["codec"] == 'h264': codec = 'h264_amf' elif cam_params["codec"] == 'h265': codec = 'hevc_amf' gpu_params = [ '-r:v', str(cam_params["frameRate"]), '-usage', 'lowlatency', '-rc', 'cqp', # constant quantization parameter '-qp_i', cam_params["quality"], '-qp_p', cam_params["quality"], '-qp_b', cam_params["quality"], '-bf:v', '0', '-hwaccel', 'auto', '-hwaccel_device', str(cam_params["gpuID"]), ] elif cam_params["gpuMake"] == 'intel': if pix_fmt_out == 'rgb0': pix_fmt_out = 'nv12' if cam_params["codec"] == 'h264': codec = 'h264_qsv' elif cam_params["codec"] == 'h265': codec = 'hevc_qsv' gpu_params = [ '-r:v', str(cam_params["frameRate"]), '-bf:v', '0', ] # Initialize writer object (imageio-ffmpeg) while (True): try: try: writer = write_frames( full_file_name, [cam_params["frameWidth"], cam_params["frameHeight"] ], # size [W,H] fps=cam_params["frameRate"], quality=None, codec=codec, pix_fmt_in=cam_params[ "pixelFormatInput"], # 'bayer_bggr8', 'gray', 'rgb24', 'bgr0', 'yuv420p' pix_fmt_out=pix_fmt_out, bitrate=None, ffmpeg_log_level=cam_params[ "ffmpegLogLevel"], # 'warning', 'quiet', 'info' input_params=['-an'], # '-an' no audio output_params=gpu_params, ) writer.send(None) # Initialize the generator break except Exception as e: logging.error('Caught exception: {}'.format(e)) time.sleep(0.1) except KeyboardInterrupt: break return writer
def make_movie(simulation, dynstate, name: str, pfilm=5, fps=24, callback=None): """ Makes a .mp4 movie of simulation from the history saved in dynstate. Parameters ---------- simulation : Simulation The simulation object dynstate : DynState The dynState object containing the position history file. name : str The path (and name) of the file to be writen. pfilm : int To make the movie, takes every pfilm position. fps : int The number of frame per seconds of the film. callback : function An optional callback function that is called every time a frame is made and is passed the current iteration number. Returns ------- """ # ouverture pour la lecture npas = simulation.current_iter YlimB = simulation.model.y_lim_inf YlimH = simulation.model.y_lim_sup XlimG = simulation.model.x_lim_inf XlimD = simulation.model.x_lim_sup if not name.endswith(".mp4"): name += ".mp4" with dynstate.open(dynstate.POS_H, 'r') as fix: # liste de k ou tracer le graph klist = set(range(0, npas, pfilm)) # boucle pour creer le film figure_size = (1920, 1088) try: os.remove(name) except FileNotFoundError: pass gen = write_frames(name, figure_size, fps=fps, quality=9) gen.send(None) fig = plt.figure(figsize=(figure_size[0] / (72 * 2), figure_size[1] / (72 * 2))) plt.clf() # definition du domaine de dessin plt.ioff() # pour ne pas afficher les graphs) plt.axis('scaled') plt.ylim(YlimB, YlimH) plt.xlim(XlimG, XlimD) pos = fix.load() line1, = plt.plot(*pos[:simulation.model.n_a, :].T, 'ro', markersize=0.5) line2, = plt.plot(*pos[simulation.model.n_a:, :].T, 'bo', markersize=0.5) plt.xlabel("0") temp = io.BytesIO() for k in range(1, npas): # dessin a chaque pas (ne s'ouvre pas: est sauvegarde de maniere incrementale) if k in klist: temp.seek(0) plt.xlabel("Iteration : {}".format(k)) plt.title(f"T = {simulation.state_fct['T'][k]:.2f} K") line1.set_data(*pos[:simulation.model.n_a, :].T) line2.set_data(*pos[simulation.model.n_a:, :].T) fig.savefig(temp, format='raw', dpi=72 * 2) # sauvegarde incrementale temp.seek(0) if callback: callback(k) gen.send( Image.frombytes('RGBA', figure_size, temp.read()).convert('RGB').tobytes()) pos = fix.load() # on charge a chaque pas de temps gen.close() plt.close(fig)
def main(): home = '/Users/Max/' filename = home + '/Documents/videos/daily_dose1.mp4' out_filename = home + '/Documents/videos/daily_dose1_edited.mp4' reader = get_reader(filename) meta = reader.get_meta_data() writer = write_frames(out_filename, meta['size'], fps=meta['fps']) writer.send(None) start_offset = 4170 looper = Looper(meta) looper.set_video(reader) looper.set_offset(start_offset) looper.set_loop_length(500) looper.set_frame_interval(5) looper1 = Looper(meta) looper1.set_video(reader) looper1.set_offset(start_offset - 1000) looper1.set_loop_length(500) looper1.set_frame_interval(5) reducer = BitReducer(meta) sharpen = Sharpener(meta) delay = VideoDelay(meta) net = ModuleNet() net.add_module(looper) # net.add_module(looper1) net.add_module(sharpen, input_names=(['Looper0'], [], [])) net.add_module(reducer, input_names=(['Sharpener0'], [], []), output_names=(['OUT'], [], [])) # net.add_module(reducer, input_names=(['Looper0'], [], []), output_names=(['OUT'], [], [])) # net.add_module(delay, input_names=(['Looper0'], [], []), output_names=(['OUT'], [], [])) # net.add_connection('Looper1', 'VideoDelay0', 'video') net.nodes['BitReducer0'].module.set_num_bits(4) # net.nodes['Sharpener0'].module.set_strength(7) # net.nodes['VideoDelay0'].module.disable() sharpen_max_strength = 20 sharpen_min_strength = 10 sharpen_delta = (sharpen_max_strength - sharpen_min_strength) / 2 angular_period = 2 * np.pi / 30 instrument = Instrument() instrument.net = net total_frames = 300 start_time = timer() for i in range(total_frames): net.nodes['Sharpener0'].module.set_strength(sharpen_min_strength + sharpen_delta + sharpen_delta * np.sin(angular_period * i)) # # net.nodes['Looper0'].tick() # net.nodes['Sharpener0'].tick() # net.nodes['BitReducer0'].tick() # instrument.reset_flags() # print(net.nodes['Looper0'].module.current_frame) # if i == 100: # net.nodes['VideoDelay0'].module.disable() # print(i) video_frame = instrument.tick() # frame = net.nodes['BitReducer0'].module.video_output_buffer writer.send(video_frame) # # print(looper.current_frame) end_time = timer() reader.close() writer.close() print( f"Time taken to run {total_frames} frames: {end_time - start_time}s.")
def style_video(video_path, style, output_name=None, alpha=1.0, preserve_color=False, custom_resolution=None, frame_skip=1): toTensor = ToTensor() reader = imageio.get_reader(video_path) resolution = None if custom_resolution is None: resolution = reader.get_meta_data()["source_size"] else: resolution = custom_resolution fps = reader.get_meta_data()["fps"] / frame_skip style = transform.resize(style, (resolution[1], resolution[0])) if not preserve_color: style_tensor = torch.tensor(np.expand_dims(toTensor(style), 0)).float().to(device) style_tensor = net.encode(style_tensor) driving_video = [] for frame_number in tqdm(range(reader.count_frames())): try: content = reader.get_next_data() except IndexError: break else: if frame_number % frame_skip == 0: if custom_resolution is not None: content = transform.resize(content, (resolution[1], resolution[0])) else: content = content.astype(float) / 255 if preserve_color: style_tensor = preserve_color_stable(content, style) style_tensor = torch.tensor( np.expand_dims(toTensor(style_tensor), 0)).float().to(device) style_tensor = net.encode(style_tensor) x = torch.tensor( toTensor(content)).unsqueeze(0).float().to(device) x = net.encode(x) result = adaIN(x, style_tensor) result = (1 - alpha) * x + alpha * result result = test(content, result, alpha, False, False) result = cv2.convertScaleAbs(result * 255) driving_video.append(result) reader.close() if output_name == None: return driving_video writer = imageio_ffmpeg.write_frames(f'{output_name}_plain.mp4', (resolution), fps=fps, macro_block_size=1) writer.send(None) for frame in driving_video: writer.send(frame) writer.close() os.system( f'ffmpeg -y -i "{video_path}" -vn -acodec copy "output_audio.aac"') os.system( f'ffmpeg -y -i "{output_name}_plain.mp4" -i "output_audio.aac" -c:v copy -c:a aac "{output_name}.mp4"' ) if os.path.isfile("output_audio.aac"): os.remove("output_audio.aac") if not os.path.isfile(f"{output_name}.mp4"): os.rename(f"{output_name}_plain.mp4", f"{output_name}.mp4") else: os.remove(f'{output_name}_plain.mp4')