はじめに

FFStreamはPython上で動画を高速に入出力する機能をまとめたクラスです。
主に動画から画像を取得して AI等で処理した結果を再び動画にする目的で使います。
最初はOpenCVやPyAV等を使っていましたが、動画の処理がボトルネックになる（AIの処理より動画の処理に時間がかかる）ケースが出てきたのでffmpeg-pythonを使う方法を試みたところ大幅な改善がみられたので、クラス化しました。

特徴

処理が速い
動画の音声をコピーできる
画質やファイルサイズの調整ができる
PyTorchで扱いやすい Tensor[C,H,W]形式のRGBフレームに対応

主な使用例

ダウンロード

使い方

ffstream.zipを解凍してできる'ffstream'フォルダをプロジェクトへコピー

ffstream_sample.py

# ffstream.pyの使用例
# 入力動画からRGB画像を受け取り 簡単な処理を加えて 音声付きの動画を作成します
import torch
from tqdm import tqdm
from ffstream.ffstream import FFStream

input_path  = 'input.mp4'                   # 入力する 動画ファイル
output_path = 'output.mp4'                  # 作成する 動画ファイル

device = torch.device('cuda')               # ビデオメモリを使って CUDAで処理する(早い)
#device = torch.device('cpu')               # メインメモリを使って CPU で処理する(遅い)

# 入力パス, 出力パス, 画質20, ピクセルフォーマットyuv420p
# 音声がある場合はコピーする, フレーム処理用のデバイスを指定する
ff = FFStream(input_path, output_path, crf=20, pix_fmt='yuv420p',
              copy_audio_stream=True, device=device)

w = ff.width        # 入力動画の 画像サイズ
h = ff.height       # 入力動画の 画像サイズ
frames = ff.frames  # 入力動画の 総フレーム数

bar = tqdm(total=frames, dynamic_ncols=True)
while True:
    # FFmpegからRGBフレームを受け取る
    frame_rgb = ff.get_rgb()
    if (frame_rgb == None) :
        break                               # 動画の最後に達した

    # RGBの処理
    frame_rgb[2]=0                          # 青成分を0にする  [C,H,W]

    # FFmpegの出力ストリームへRGBフレームを送る
    ff.put_rgb(frame_rgb)

    bar.update(1)

bar.close()
ff.close()

リファレンス

class FFStream() v1.3.0

コンストラクタ

項目	説明
書式	ff = FFStream(’input.mp4’, ’output.mp4’...)
引数	下記
戻り値	なし
使用例	下記

引数

    input_path:str=None		# 入力パス、省略可
    
    output_path:str=None	# 出力パス、省略可
    						# 入力パスと出力パスの両方または一方を指定して使う
          
	mbps:int=None			# 画質指定(ビットレート) 10なら10Mbpsでエンコード
 
	crf:float=None			# 画質指定(品質) 20が標準
 							# 数値が低いほど 高画質でファイルサイズが大きくなる
        
	pix_fmt:str='yuv420p'	# ピクセルフォーマット
 							# 主に 動画なら'yuv420p', jpgなら'rgb24', pngなら'rgba'
        
	copy_audio_stream:bool=False
 							# 入力動画の音声を出力動画へコピーする場合True
        					# 音声がない場合はコピーしません
             
	ss:str=None				# 開始位置
 							# '60'				60秒
        					# f'{1000/60}'		60fpsの動画の場合に1000フレーム当たり
             				# '1:30:00'			1時間30分0秒
                 
	ow:int=None				# 出力画像サイズを指定する
 							# 入力と異なる画像サイスにする場合や
        					# 入力を使わず出力だけする場合に指定する
             
	oh:int=None				# 出力画像サイズを指定する
 
	ofps:float=None			# 出力フレームレートを指定する
 							# 入力と異なる場合や 入力を使わず出力だけする場合に指定する
        					# 29.97の場合は 30000/1001
             
	image_scale:float=1		# 出力画像サイズの倍率
 
	fps_scale:float=1		# 出力フレームレートの倍率
 
	no_probe:bool=False		# プローブを使わない
 							# 画像サイズ等の情報を取得する処理は遅い
        					# 連番動画などを大量に読み込む場合Trueにする
             				# 代わりに画像サイズをow,ohで指定する
                 
	get_info:bool=False		# 入力ファイルの情報を取得するだけで ffmpegを実行しない
 
	device:str=None			# フレーム処理をビデオメモリにするかメインメモリにするか指定する
 							# 省略した場合は メインメモリ
        
	input_options:str=None	# ffmpegの入力ストリームへ設定するオプション文字列
 							# ダブルクォーテーションは任意
        					# ハードウエアデコーダーを使う例 "-hwaccel cuda"
             
    output_options:str=None	# ffmpegの出力ストリームへ設定するオプション文字列
 							# ダブルクォーテーションは任意
        					# プロファイルを指定する例 "-profile:v baseline -level:v 4.1"
             
    mt_buffer:int=0			# 1以上の場合 転送処理を別スレッドで行う
    						# その時使うバッファ数を指定する
          					# フレームの入出力は get_rgb(), put_rgb()を使う
          					# 最後は必ず def close(self):を使う               
          					# 0の場合はメインスレッドで転送する

使用例

動画の入出力

from ffstream.ffstream import FFStream
input_path  = 'input.mp4'                   # 入力する 動画ファイル
output_path = 'output.mp4'                  # 作成する 動画ファイル
ff = FFStream(input_path, output_path)
for i in range(ff.frames):
    frame_rgb = ff.get_rgb()
    ff.put_rgb(frame_rgb)
    print(f'{i+1} / {ff.frames}')
ff.close()

動画の入力のみ

from ffstream.ffstream import FFStream
input_path  = 'input.mp4'                   # 入力する 動画ファイル
ff = FFStream(input_path, None)
frame_rgb = ff.get_rgb()
ff.close()

動画の出力のみ

from ffstream.ffstream import FFStream
import torch
w:int       = 1920
h:int       = 1080
fps:float   = 30000/1001
frame_rgb   = torch.ones(3,h,w)
output_path = 'output.mp4'                  # 作成する 動画ファイル
ff = FFStream(None, output_path, ow=w, oh=h, ofps=fps)
ff.put_rgb(frame_rgb)
ff.close()

動画や画像の情報だけ取得する

from ffstream.ffstream import FFStream
input_path  = 'input.mp4'                   # 入力する 動画ファイル
ff = FFStream(input_path, None, get_info=True)
w = ff.width        # 入力動画の 画像サイズ
h = ff.height       # 入力動画の 画像サイズ
frames = ff.frames  # 入力動画の 総フレーム数
print(f'{w}x{h}, {frames}frames')

画像を読み込む

from ffstream.ffstream import FFStream
input_path  = 'input.jpg'                   # 入力する 画像ファイル
#input_path = 'input.png'                   # 入力する 画像ファイル
ff = FFStream(input_path, None, pix_fmt='rgb24')
frame_rgb = ff.get_rgb()
ff.close()

アルファチャンネル付きpng画像を読み込む

from ffstream.ffstream import FFStream
import torch
input_path  = 'input.png'               # 入力する 画像ファイル
ff = FFStream(input_path, None, pix_fmt='rgba')
w, h = ff.width, ff.height              # 画像サイズ
frame_bytes = ff.recv_frame()           # RGBARGBARGBARGBA
img_src     = torch.frombuffer(frame_bytes, dtype=torch.uint8)
img_src     = img_src.cuda().reshape(h,w,4)
img_src     = img_src.permute(2, 0, 1)  # 配列を整形 [C,H,W] RRRRGGGGBBBBAAAA
frame_rgb   = img_src[0:3]              # RGBフレーム
frame_alpha = img_src[3:4]              # アルファチャンネル
ff.close()

recv_frame()

ローレベル入力。
FFmpegの入力ストリームから受け取ったフレームを返す。
フレームが送られてくるまで処理は中断する。

項目	説明
書式	frame_bytes = recv_frame()
引数	なし
戻り値	受け取ったフレーム（1次元のByte配列）動画の最後まで達した場合はNoneを返す

send_frame(bytes)

ローレベル出力。
FFmpegの出力ストリームへフレームを送る。

項目	説明
書式	send_frame(bytes)
引数	- bytes 送るフレーム（1次元のByte配列）
戻り値	なし

close()

起動しているFFmpegのプロセスを終了する。

項目	説明
書式	close()
引数	なし
戻り値	なし

get_fps()->float:

入力ストリームのフレームレートを float型で返す。

項目	説明
書式	fps = get_fps()
引数	なし
戻り値	フレームレート（float型）

get_rgb()->Tensor

入力ストリームからRGBフレームを取得する。
yuv420pまたは rgb24を受け取って RGBフレーム（Tensor[C,H,W]）へ変換する。
動画の最後に達したら Noneを返す。

項目	説明
書式	frame_rgb = get_rgb()
引数	なし
戻り値	RGBフレーム（Tensor[C,H,W]）動画の最後まで達した場合はNoneを返す

put_rgb(frame_rgb:Tensor)

出力ストリームへRGBフレームを送る。
RGBフレーム（Tensor[C,H,W]）をyuv420pまたは rgb24へ変換する。

項目	説明
書式	put_rgb(frame_rgb)
引数	- frame_rgb:Tensor RGBフレーム（Tensor[C,H,W]）
戻り値	なし

そのほか

yuv420pのストリーム用入出力関数
- get_yuv420_to_rgb()->Tensor
- put_rgb_to_yuv420(frame_rgb:Tensor)
rgb24のストリーム用入出力関数
- get_rgb24_to_rgb()->Tensor
- put_rgb_to_rgb24(frame_rgb:Tensor)