@@ -1,301 +0,0 @@
# Copyright 2022 Lunar Ring. All rights reserved.
# Written by Johannes Stelzer, email stelzer@lunar-ring.ai twitter @j_stelzer
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import subprocess
import os
import numpy as np
from tqdm import tqdm
import cv2
from typing import List
import ffmpeg # pip install ffmpeg-python. if error with broken pipe: conda update ffmpeg
class MovieSaver ( ) :
def __init__ (
self ,
fp_out : str ,
fps : int = 24 ,
shape_hw : List [ int ] = None ,
crf : int = 21 ,
codec : str = ' libx264 ' ,
preset : str = ' fast ' ,
pix_fmt : str = ' yuv420p ' ,
silent_ffmpeg : bool = True ) :
r """
Initializes movie saver class - a human friendly ffmpeg wrapper .
After you init the class , you can dump numpy arrays x into moviesaver . write_frame ( x ) .
Don ' t forget toi finalize movie file with moviesaver.finalize().
Args :
fp_out : str
Output file name . If it already exists , it will be deleted .
fps : int
Frames per second .
shape_hw : List [ int , int ]
Output shape , optional argument . Can be initialized automatically when first frame is written .
crf : int
ffmpeg doc : the range of the CRF scale is 0 – 51 , where 0 is lossless
( for 8 bit only , for 10 bit use - qp 0 ) , 23 is the default , and 51 is worst quality possible .
A lower value generally leads to higher quality , and a subjectively sane range is 17 – 28.
Consider 17 or 18 to be visually lossless or nearly so ;
it should look the same or nearly the same as the input but it isn ' t technically lossless.
The range is exponential , so increasing the CRF value + 6 results in
roughly half the bitrate / file size , while - 6 leads to roughly twice the bitrate .
codec : int
Number of diffusion steps . Larger values will take more compute time .
preset : str
Choose between ultrafast , superfast , veryfast , faster , fast , medium , slow , slower , veryslow .
ffmpeg doc : A preset is a collection of options that will provide a certain encoding speed
to compression ratio . A slower preset will provide better compression
( compression is quality per filesize ) .
This means that , for example , if you target a certain file size or constant bit rate ,
you will achieve better quality with a slower preset . Similarly , for constant quality encoding ,
you will simply save bitrate by choosing a slower preset .
pix_fmt : str
Pixel format . Run ' ffmpeg -pix_fmts ' in your shell to see all options .
silent_ffmpeg : bool
Surpress the output from ffmpeg .
"""
if len ( os . path . split ( fp_out ) [ 0 ] ) > 0 :
assert os . path . isdir ( os . path . split ( fp_out ) [ 0 ] ) , " Directory does not exist! "
self . fp_out = fp_out
self . fps = fps
self . crf = crf
self . pix_fmt = pix_fmt
self . codec = codec
self . preset = preset
self . silent_ffmpeg = silent_ffmpeg
if os . path . isfile ( fp_out ) :
os . remove ( fp_out )
self . init_done = False
self . nmb_frames = 0
if shape_hw is None :
self . shape_hw = [ - 1 , 1 ]
else :
if len ( shape_hw ) == 2 :
shape_hw . append ( 3 )
self . shape_hw = shape_hw
self . initialize ( )
print ( f " MovieSaver initialized. fps= { fps } crf= { crf } pix_fmt= { pix_fmt } codec= { codec } preset= { preset } " )
def initialize ( self ) :
args = (
ffmpeg
. input ( ' pipe: ' , format = ' rawvideo ' , pix_fmt = ' rgb24 ' , s = ' {} x {} ' . format ( self . shape_hw [ 1 ] , self . shape_hw [ 0 ] ) , framerate = self . fps )
. output ( self . fp_out , crf = self . crf , pix_fmt = self . pix_fmt , c = self . codec , preset = self . preset )
. overwrite_output ( )
. compile ( )
)
if self . silent_ffmpeg :
self . ffmpg_process = subprocess . Popen ( args , stdin = subprocess . PIPE , stderr = subprocess . DEVNULL , stdout = subprocess . DEVNULL )
else :
self . ffmpg_process = subprocess . Popen ( args , stdin = subprocess . PIPE )
self . init_done = True
self . shape_hw = tuple ( self . shape_hw )
print ( f " Initialization done. Movie shape: { self . shape_hw } " )
def write_frame ( self , out_frame : np . ndarray ) :
r """
Function to dump a numpy array as frame of a movie .
Args :
out_frame : np . ndarray
Numpy array , in np . uint8 format . Convert with np . astype ( x , np . uint8 ) .
Dim 0 : y
Dim 1 : x
Dim 2 : RGB
"""
assert out_frame . dtype == np . uint8 , " Convert to np.uint8 before "
assert len ( out_frame . shape ) == 3 , " out_frame needs to be three dimensional, Y X C "
assert out_frame . shape [ 2 ] == 3 , f " need three color channels, but you provided { out_frame . shape [ 2 ] } . "
if not self . init_done :
self . shape_hw = out_frame . shape
self . initialize ( )
assert self . shape_hw == out_frame . shape , f " You cannot change the image size after init. Initialized with { self . shape_hw } , out_frame { out_frame . shape } "
# write frame
self . ffmpg_process . stdin . write (
out_frame
. astype ( np . uint8 )
. tobytes ( )
)
self . nmb_frames + = 1
def finalize ( self ) :
r """
Call this function to finalize the movie . If you forget to call it your movie will be garbage .
"""
if self . nmb_frames == 0 :
print ( " You did not write any frames yet! nmb_frames = 0. Cannot save. " )
return
self . ffmpg_process . stdin . close ( )
self . ffmpg_process . wait ( )
duration = int ( self . nmb_frames / self . fps )
print ( f " Movie saved, { duration } s playtime, watch here: \n { self . fp_out } " )
def concatenate_movies ( fp_final : str , list_fp_movies : List [ str ] ) :
r """
Concatenate multiple movie segments into one long movie , using ffmpeg .
Parameters
- - - - - - - - - -
fp_final : str
Full path of the final movie file . Should end with . mp4
list_fp_movies : list [ str ]
List of full paths of movie segments .
"""
assert fp_final [ - 4 ] == " . " , " fp_final seems to miss file extension: {fp_final} "
for fp in list_fp_movies :
assert os . path . isfile ( fp ) , f " Input movie does not exist: { fp } "
assert os . path . getsize ( fp ) > 100 , f " Input movie seems empty: { fp } "
if os . path . isfile ( fp_final ) :
os . remove ( fp_final )
# make a list for ffmpeg
list_concat = [ ]
for fp_part in list_fp_movies :
list_concat . append ( f """ file ' { fp_part } ' """ )
# save this list
fp_list = " tmp_move.txt "
with open ( fp_list , " w " ) as fa :
for item in list_concat :
fa . write ( " %s \n " % item )
cmd = f ' ffmpeg -f concat -safe 0 -i { fp_list } -c copy { fp_final } '
subprocess . call ( cmd , shell = True )
os . remove ( fp_list )
if os . path . isfile ( fp_final ) :
print ( f " concatenate_movies: success! Watch here: { fp_final } " )
def add_sound ( fp_final , fp_silentmovie , fp_sound ) :
cmd = f ' ffmpeg -i { fp_silentmovie } -i { fp_sound } -c copy -map 0:v:0 -map 1:a:0 { fp_final } '
subprocess . call ( cmd , shell = True )
if os . path . isfile ( fp_final ) :
print ( f " add_sound: success! Watch here: { fp_final } " )
def add_subtitles_to_video (
fp_input : str ,
fp_output : str ,
subtitles : list ,
fontsize : int = 50 ,
font_name : str = " Arial " ,
color : str = ' yellow '
) :
from moviepy . editor import VideoFileClip , TextClip , CompositeVideoClip
r """
Function to add subtitles to a video .
Args :
fp_input ( str ) : File path of the input video .
fp_output ( str ) : File path of the output video with subtitles .
subtitles ( list ) : List of dictionaries containing subtitle information
( start , duration , text ) . Example :
subtitles = [
{ " start " : 1 , " duration " : 3 , " text " : " hello test " } ,
{ " start " : 4 , " duration " : 2 , " text " : " this works " } ,
]
fontsize ( int ) : Font size of the subtitles .
font_name ( str ) : Font name of the subtitles .
color ( str ) : Color of the subtitles .
"""
# Check if the input file exists
if not os . path . isfile ( fp_input ) :
raise FileNotFoundError ( f " Input file not found: { fp_input } " )
# Check the subtitles format and sort them by the start time
time_points = [ ]
for subtitle in subtitles :
if not isinstance ( subtitle , dict ) :
raise ValueError ( " Each subtitle must be a dictionary containing ' start ' , ' duration ' and ' text ' . " )
if not all ( key in subtitle for key in [ " start " , " duration " , " text " ] ) :
raise ValueError ( " Each subtitle dictionary must contain ' start ' , ' duration ' and ' text ' . " )
if subtitle [ ' start ' ] < 0 or subtitle [ ' duration ' ] < = 0 :
raise ValueError ( " ' start ' should be non-negative and ' duration ' should be positive. " )
time_points . append ( ( subtitle [ ' start ' ] , subtitle [ ' start ' ] + subtitle [ ' duration ' ] ) )
# Check for overlaps
time_points . sort ( )
for i in range ( 1 , len ( time_points ) ) :
if time_points [ i ] [ 0 ] < time_points [ i - 1 ] [ 1 ] :
raise ValueError ( " Subtitle time intervals should not overlap. " )
# Load the video clip
video = VideoFileClip ( fp_input )
# Create a list to store subtitle clips
subtitle_clips = [ ]
# Loop through the subtitle information and create TextClip for each
for subtitle in subtitles :
text_clip = TextClip ( subtitle [ " text " ] , fontsize = fontsize , color = color , font = font_name )
text_clip = text_clip . set_position ( ( ' center ' , ' bottom ' ) ) . set_start ( subtitle [ " start " ] ) . set_duration ( subtitle [ " duration " ] )
subtitle_clips . append ( text_clip )
# Overlay the subtitles on the video
video = CompositeVideoClip ( [ video ] + subtitle_clips )
# Write the final clip to a new file
video . write_videofile ( fp_output )
class MovieReader ( ) :
r """
Class to read in a movie .
"""
def __init__ ( self , fp_movie ) :
self . video_player_object = cv2 . VideoCapture ( fp_movie )
self . nmb_frames = int ( self . video_player_object . get ( cv2 . CAP_PROP_FRAME_COUNT ) )
self . fps_movie = int ( self . video_player_object . get ( cv2 . CAP_PROP_FPS ) )
self . shape = [ 100 , 100 , 3 ]
self . shape_is_set = False
def get_next_frame ( self ) :
success , image = self . video_player_object . read ( )
if success :
if not self . shape_is_set :
self . shape_is_set = True
self . shape = image . shape
return image
else :
return np . zeros ( self . shape )
if __name__ == " __main__ " :
fps = 2
list_fp_movies = [ ]
for k in range ( 4 ) :
fp_movie = f " /tmp/my_random_movie_ { k } .mp4 "
list_fp_movies . append ( fp_movie )
ms = MovieSaver ( fp_movie , fps = fps )
for fn in tqdm ( range ( 30 ) ) :
img = ( np . random . rand ( 512 , 1024 , 3 ) * 255 ) . astype ( np . uint8 )
ms . write_frame ( img )
ms . finalize ( )
fp_final = " /tmp/my_concatenated_movie.mp4 "
concatenate_movies ( fp_final , list_fp_movies )
DGX