真っ黒のpng画像の特定の画素値を透明化する。

Question

前提

https://teratail.com/questions/0vl3yd7kmwsl47
こちらの質問の続きです。動画から得られた画像に対して学習モデルを用いた検出はできるようになりました。その続きになります。

import segmentation_models as sm
sm.set_framework('tf.keras')
sm.framework()
import glob
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import cv2
import os
import shutil
import re

from tensorflow.python.keras.models import load_model
from PIL import Image

#データセット
x_image_size = 320 #インプット画像のサイズ
y_image_size = 224 #インプット画像のサイズ

#画像表示の関数
def visualize(**images):
    n = len(images)
    plt.figure(figsize=(32, 16))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(' '.join(name.split('_')).title())
        plt.imshow(image)
    plt.show()
    
def atoi(text):
    return int(text) if text.isdigit() else text

def natural_keys(text):
    return [ atoi(c) for c in re.split(r'(\d+)', text) ]

#判別したい画像群用配列
TestImgArray = []

#TestImageディレクトリ内の画像をTestIArrayに格納していく
dir = 'C:/Users/XXXX/Desktop/VideoImage/VideoImages'
dir_NumSum = sorted(glob.glob(dir+"/*.png"), key = natural_keys)
for i, file in enumerate(dir_NumSum): 
    image = Image.open(file) #画像の読み込み
    image = image.convert("RGB") #RGBに変換
    image = image.resize((x_image_size, y_image_size)) #サイズ変更
    image = np.asarray(image) #画像データを配列へ変更
    image = np.expand_dims(image, axis=0)
    pr_mask = model.predict(image)

    #抽出画像の描画
    visualize(
        image = image.squeeze(),
        A_Prediction = pr_mask[..., 1].squeeze(),
        B_Prediction = pr_mask[..., 2].squeeze(),
    )
    
    print(type(A_Prediction))
    
    #顔部分のマスクを取得
    A_Prediction = pr_mask[..., 1].squeeze()
    cv2.imwrite("C:/Users/XXXX/Desktop/VideoImage/MaskImages/FaceMasks/FaceMask_%d.png" % i, A_Prediction)
    
    #口部部分のマスクを取得
    B_Prediction = pr_mask[..., 2].squeeze()
    cv2.imwrite("C:/Users/XXXX/Desktop/VideoImage/MaskImages/MouthMasks/MouthMask_%d.png" % i, B_Prediction)

できたこと

皆様のご助力のおかげで、リンク先で質問させていただいたことは解決いたしました。
次は、取得できた画像(学習モデルが判別した顔パーツ)から、マスク画像の生成を行おうとしています。
取得できた画像は以下のように真っ黒なのですが、学習モデルがパーツを判別した部分の画素値は[1, 1, 1]となっています。

これより、画素値が[0, 0, 0]となっている個所を透過させるプログラムを作成し、以下の画像が得られました。

#ライブラリの読み込み
import cv2
import numpy as np
import glob
import re
import matplotlib.pyplot as plt

from IPython.display import Image, display

MaskPath = "C:/Users/XXXX/Desktop/VideoImage/MaskImages/FaceMasks/"
ClearPath = "C:/Users/XXXX/Desktop/VideoImage/MaskImages/ClearMask/"

i = 0

for i in range(300):
    img = cv2.imread(MaskPath + "FaceMask_" + str(i)+".png")
    
    #マスク画像を生成
    mask = np.all(img[:] == [0], axis=-1)

    # 元画像をBGR形式からBGRA形式に変換
    dst = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    #マスク画像をもとに透明化
    dst[mask,3] = 0
    
    #画像保存
    cv2.imwrite(ClearPath+str(i)+"result.png", dst)
    
    i += 1

やりたいこと

動画から画像を取得しているので、30fps, 10秒の動画の場合、300枚の画像が取得されます。
現状の流れとしましては

1. 動画から取得された300枚の画像1枚1枚にパーツの判別を行う。
2. 「できたこと」であげたような真っ黒な画像が300枚(600枚)生成されるので、これを専用のディレクトリに保存。(最初のプログラム63行目, 67行目)
3. 専用ディレクトリに格納された300枚の画像一枚一枚に対して、透過の処理を行う。(２番目のプログラム)

という操作をしています。

これを、最初のプログラムだけで完結させることはできないか
(3.を、最初のプログラムの専用ディレクトに保存する段階でできないか)

    #顔部分のマスクを取得
    A_Prediction = pr_mask[..., 1].squeeze()
    cv2.imwrite("C:/Users/XXXX/Desktop/VideoImage/MaskImages/FaceMasks/FaceMask_%d.png" % i, A_Prediction)
    
    #口部部分のマスクを取得
    B_Prediction = pr_mask[..., 2].squeeze()
    cv2.imwrite("C:/Users/XXXX/Desktop/VideoImage/MaskImages/MouthMasks/MouthMask_%d.png" % i, B_Prediction)

が一番やりたいことで、
それが無理ならば3.の処理を行いたいのですが、エラーが発生してしまい詰まっています。

発生している問題・エラーメッセージ

2番目のプログラムを実行すると以下のエラーが発生します。

TypeError                                 Traceback (most recent call last)
\AppData\Local\Temp\ipykernel_19140\2419611442.py in <module>
      9 
     10     #マスク画像を生成
---> 11     mask = np.all(img[:] == [0], axis=-1)
     12 
     13     # 元画像をBGR形式からBGRA形式に変換

TypeError: 'NoneType' object is not subscriptable

試したこと

試したことと言いますか、2番目のプログラムは最初の数枚程度であれば透過画像を出力してくれるので、途中で止まる理由がわかりません。
2番目のプログラムの15行目、0～300という指定を行っているのにディレクトリ内に100枚しか画像がなく、実行の途中で存在しない画像を参照しようとすることが原因のエラーでした。大変失礼しました。また、2番目のプログラムは改変を行い、真っ黒画像から透過画像を作成するプログラムを以下のように改変しました。

#ライブラリの読み込み
import cv2
import numpy as np
import glob
import re
import matplotlib.pyplot as plt

from IPython.display import Image, display

def atoi(text):
    return int(text) if text.isdigit() else text

def natural_keys(text):
    return [ atoi(c) for c in re.split(r'(\d+)', text) ]

dir = 'C:/Users/XXXX/Desktop/VideoImage/MaskImages/FaceMasks'
dir_NumSum = sorted(glob.glob(dir+"/*.png"), key = natural_keys)
for i, file in enumerate(dir_NumSum): 
    img = cv2.imread(file)
    
    #透過処理
    mask = np.all(img[:] == [0], axis=-1)
    dst = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
    dst[mask,3] = 0
    
    #画像保存
    cv2.imwrite("C:/Users/XXXX/Desktop/VideoImage/MaskImages/ClearMask/ClearMask_%d.png" % i, dst)

先ほども申し上げた通り　変更後のプログラムにせよ、最初のプログラムで真っ黒のpng画像を取得して、2番目のプログラムで透過作業を行うことは変わっていません。素人ながらに冗長だと思いますので、最初のプログラムで透過までできる方法があればアドバイス頂けたらと思います。何卒よろしくお願い申し上げます。

補足情報（FW/ツールのバージョンなど）

実行環境
・CPU：windows11
・GPU：RTX3060
・メモリ：16GB

構築環境
・CUDA 11.0
・cuDNN 8.0
・tensorflow-gpu == 2.4.0
・python 3.7

備考
・利用モデルはSegmentation models
・Jupyter notebook上での動作を想定

Comments

[meg_]

> 試したことと言いますか、2番目のプログラムは最初の数枚程度であれば透過画像を出力してくれるので、途中で止まる理由がわかりません。
特定の画像で止まるのでしょうか？ランダムでしょうか？

[osumosan]

コメント頂き、ありがとうございます。
自分でもう一度確認したところ、単純に2番目のプログラムでfor文の指定を300にしているにもかかわらず、ディレクトリ内に画像が300未満しか入っていないことによるエラーでした…
ディレクトリに100枚しか画像がないのに、101枚目を参照したら、そりゃエラーになりますよね…お騒がせしました。