background_extraction.py

import cv2
import numpy as np

#動画ファイルパス

ESC_KEY = 27     # Escキー
INTERVAL= 33     # インターバル
FRAME_RATE = 30  # fps

#ウインドウ名の定義
WINDOW_ORG = "Orginal_frame"
WINDOW_BACK = "Back_frame"
WINDOW_DIFF = "Difference_frame"
WINDOW_DIFF_BINARY = "Difference_frame_Binary"

#ファイルリストのインデックスを入力
# print("Input number of index of movie:")
# number = input()

#ファイル名(拡張子含む)を取得
# FILE_ORG =file_list[int(number)]
FILE_ORG = "114C0002.mov"
background_FILE_ORG = "office_background_training_1080.mov"
#処理後に生成される動画ファイル名のためにひとまず拡張子を取り除いた文字列を用意しておく
#FILE_NAME = FILE_ORG.replace(".AVI", "")
FILE_NAME = "output.mov"
#今から処理される動画ファイル名を表示
print(background_FILE_ORG)

# ウィンドウの準備
cv2.namedWindow(WINDOW_ORG)
cv2.namedWindow(WINDOW_BACK)
cv2.namedWindow(WINDOW_DIFF)
cv2.namedWindow(WINDOW_DIFF_BINARY)

#動画処理

# 元ビデオファイル読み込み
mov_org = cv2.VideoCapture(background_FILE_ORG)

# 最初のフレーム(int型)とその次(つまり２番目)のフレームの読み込み
has_next, i_frame = mov_org.read()
# 背景フレームの初期化(値は全て「0」),最初のフレームと同じサイズでフロート型
back_frame = np.zeros_like(i_frame, np.float32)

# コードブック法(リアルタイムで背景を更新するやり方)を用いるが、動画開始の数フレームに
#日付が書かれていたので性能が落ちた。
#そこで一度動画を再生しきって背景を作成し、そのあとに差分を取ることにした
while has_next == True:
    # 入力画像を浮動小数点型に変換
    f_frame = i_frame.astype(np.float32)
    # 背景の更新、最後の引数は重み
    cv2.accumulateWeighted(f_frame, back_frame, 0.01)
    # 次のフレーム読み込み
    has_next, i_frame = mov_org.read()
    # 処理を中断したければ、Escキーで終了
    key = cv2.waitKey(INTERVAL)
    if key == ESC_KEY:
        break
        
print("背景学習完了")

# 動画の再生
thresh = 50
max_pixel = 255

# 元ビデオファイル読み込み(初期化)
mov_org = cv2.VideoCapture(FILE_ORG)

# ビデオライターオブジェクトの定義
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out_BINARY = cv2.VideoWriter(FILE_NAME,int(fourcc), FRAME_RATE, \
                        (int(mov_org.get(cv2.CAP_PROP_FRAME_WIDTH)),int(mov_org.get(cv2.CAP_PROP_FRAME_HEIGHT))))

# 最初のフレーム読み込み
has_next, i_frame = mov_org.read()
    
# 変換処理ループ
while has_next == True:
    # 入力画像を浮動小数点型に変換
    f_frame = i_frame.astype(np.float32)

    # 差分(絶対値)計算
    diff_frame = cv2.absdiff(f_frame, back_frame)
        
    #2値化
    ret, img_dst = cv2.threshold(diff_frame, thresh,max_pixel,cv2.THRESH_BINARY)

    # フレーム表示
    #Original_frame
    cv2.imshow(WINDOW_ORG, i_frame)
    #学習した背景画像
    cv2.imshow(WINDOW_BACK, back_frame.astype(np.uint8))
    #差分画像
    cv2.imshow(WINDOW_DIFF, diff_frame.astype(np.uint8))
    #差分画像"(2値化)"
    cv2.imshow(WINDOW_DIFF_BINARY, img_dst.astype(np.uint8))

    #2値化したフレームを動画に保存
    out_BINARY.write(img_dst.astype(np.uint8))
        
    # Escキーで終了
    key = cv2.waitKey(INTERVAL)
    if key == ESC_KEY:
        break

    # 次のフレーム読み込み
    has_next, i_frame = mov_org.read()

# 終了処理
#ウインドウを閉じる
cv2.destroyAllWindows()
#構造体の解放
mov_org.release()
out_BINARY.release()