ラベル画像の切り出し

Date:

Share post:

前回の投稿(Mac+PostgreSQLでベクトル検索)で、「次回は、先に生成したモデルを使用してラベル部分をトリミングし、ベクトル化した上での検索精度に関して確認してみる」と風呂敷を広げましたが、いざ作業を始めてみると、ラベル部分を切り出すだけで結構な作業が発生しましたので、一旦そこまでの内容をまとめたいと思います。

なお、昨今の関連投稿においては「ラベル画像のトリミング」と銘打ってシリーズ化してきましたが、「トリミング」とは「画像全体の構図を整えるために、不要な部分を切り落とすことを指すのが一般的」(Gemini先生談)とのことで、ラベル部分を切り出すという本作業の趣旨とは若干異なる意味になるようなので、今後はトリミングという表現は使わないことにします。

関連パッケージインストール

新たにPythonのvenv環境を構築し、必要なパッケージをインストールします。
現時点の作業に必要なパッケージは以下の3つです。

# pip install tensorflow tensorflow-metal opencv-python

モデルのロード

当然ながら、ラベル部分の特定には今まで苦労して育ててきたセグメンテーション用モデルを使用します。
学習時にも既存モデルの読み込み処理はありましたが、改めて今回は以下のような関数を用意しました。

import tensorflow as tf

from train_model import dice_loss
from train_model import dice_coef

def get_model(model_path):
    try:
        model = tf.keras.models.load_model(
            model_path,
            custom_objects={'dice_loss': dice_loss, 'dice_coef': dice_coef}
        )
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        exit()

損失関数、評価指標関数は学習時に利用したtrain_model.py内の関数を指定します。
これら関数はあくまで学習用に必要なもので、今回のように当該モデルで予測を実施する段階では使わないのですが、同関数を使って学習を行なったモデルのロードに際しては、それら関数への参照が有効になっている必要があるようです。

マスク画像の生成

マスク画像の生成に関しては以下の通りです。

import cv2
import numpy as np
import tensorflow as tf
import os

IMAGE_SIZE = (256, 256)

def preprocess_image(image_path, image_size):
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, image_size)
    image = image / 255.0
    return image.astype(np.float32)

def infer_mask(model, image):
    image = np.expand_dims(image, axis=0)
    prediction = model.predict(image)
    mask = prediction.squeeze()
    return mask

def extract_largest_mask(mask_image: np.ndarray, binary_threshold: int = 127) -> np.ndarray:
    temp_mask = (np.clip(mask_image, 0, 1) * 255).astype(np.uint8)
    _, binary_mask = cv2.threshold(temp_mask, binary_threshold, 255, cv2.THRESH_BINARY)

    contours, _ = cv2.findContours(binary_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    max_area = 0
    max_contour = None
    for contour in contours:
        area = cv2.contourArea(contour)
        if area > max_area:
            max_area = area
            max_contour = contour

    largest_object_mask = np.zeros(binary_mask.shape, dtype=np.uint8)
    if max_contour is not None:
        cv2.drawContours(largest_object_mask, [max_contour], -1, 255, cv2.FILLED)
    
    return largest_object_mask

def get_mask(model, image_path,
             binary_threshold: int = 127):
    image = preprocess_image(image_path, IMAGE_SIZE)

    raw_mask_prediction = infer_mask(model, image)

    final_mask = extract_largest_mask(raw_mask_prediction, binary_threshold)

    return final_mask

全体的には今までも実施してきた処理ですが、特筆すべきは今回追加した関数extract_largest_maskの内容です。

今までのセグメンテーションにおいて、生成されたマスク画像にはラベル以外の部分を対象と誤認した箇所や、逆にラベル内の一部を対象外と誤認した箇所が相当数見受けられました。
extract_largest_maskはそれらへの対処を行う関数です。
具体的には、複数のマスク箇所(マスク画像内の白く塗られた箇所)のそれぞれに関して輪郭を抽出し、その中から最大の面積を持つものを抽出して、その内部を全て白く塗り潰したマスク画像を新たに生成しています。
これにより、ラベル部分以外の離れ小島のようなマスク箇所は新マスク画像生成時に無視されますし、ラベル部分にシミのように生じた黒い箇所も白く塗り潰されるので、本来のラベル部分を切り抜くためのマスクとして、より適切なものになります。

ラベル画像の生成

先に生成されたマスクを元画像に適用することで、ラベル(と思われる)部分をそのまま残し、それ以外を黒く塗り潰した画像が生成できます。
ただ、これだけではラベルの位置もサイズも様々な画像ができてしまうので、切り出したラベル部分が所定のサイズ(その後のベクトル化を考慮して224×224を想定)で大写しされた画像に加工したいと思います。

具体的な内容は以下の通り(例によってGemini先生作)。

import cv2
import numpy as np

def apply_mask_and_crop(org_image_path, mask, target_crop_size: int = 224) -> np.ndarray:
    org_image = cv2.imread(org_image_path)
    org_image = cv2.cvtColor(org_image, cv2.COLOR_BGR2RGB)
    org_height, org_width, _ = org_image.shape

    resized_mask = cv2.resize(mask, (org_width, org_height), interpolation=cv2.INTER_NEAREST)
    resized_mask_expanded = np.expand_dims(resized_mask, axis=-1)
    masked_image = cv2.bitwise_and(org_image, org_image, mask=resized_mask)

    contours, _ = cv2.findContours(resized_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    x, y, w, h = cv2.boundingRect(contours[0])
    min_x, min_y = x, y
    max_x, max_y = x + w, y + h
    min_x, min_y = max(0, min_x), max(0, min_y)
    max_x, max_y = min(org_width, max_x), min(org_height, max_y)
    
    bbox_width = max_x - min_x
    bbox_height = max_y - min_y

    crop_width_orig = max(bbox_width, bbox_height)
    crop_height_orig = max(bbox_width, bbox_height)

    center_x = min_x + bbox_width // 2
    center_y = min_y + bbox_height // 2

    crop_x1 = center_x - crop_width_orig // 2
    crop_y1 = center_y - crop_height_orig // 2
    crop_x2 = crop_x1 + crop_width_orig
    crop_y2 = crop_y1 + crop_height_orig

    pad_left = max(0, -crop_x1)
    pad_top = max(0, -crop_y1)
    pad_right = max(0, crop_x2 - org_width)
    pad_bottom = max(0, crop_y2 - org_height)
    padded_masked_image = cv2.copyMakeBorder(masked_image, 
                                               pad_top, pad_bottom, 
                                               pad_left, pad_right, 
                                               cv2.BORDER_CONSTANT, value=[0, 0, 0])

    crop_x1_padded = crop_x1 + pad_left
    crop_y1_padded = crop_y1 + pad_top
    crop_x2_padded = crop_x2 + pad_left
    crop_y2_padded = crop_y2 + pad_top
    cropped_image = padded_masked_image[crop_y1_padded:crop_y2_padded, crop_x1_padded:crop_x2_padded]
    
    final_cropped_image = cv2.resize(cropped_image,
                                     (target_crop_size, target_crop_size),
                                     interpolation=cv2.INTER_AREA)
    
    return final_cropped_image

ポイントになりそうな部分を個別に見ていきます。

    resized_mask = cv2.resize(mask, (org_width, org_height), interpolation=cv2.INTER_NEAREST)
    resized_mask_expanded = np.expand_dims(resized_mask, axis=-1)
    masked_image = cv2.bitwise_and(org_image, org_image, mask=resized_mask)

モデルの学習に際しては元画像(org_image)を256×256にリサイズしていたので、同モデルが生成したマスク画像も同じサイズ(256×256)になっています。
このマスクをリサイズ後の画像に適用すると、ラベル部分はかなり小さい(情報量の少ない)ものになってしまいます。
よって、リサイズ前の元画像に対してマスクが適用できるよう、マスク画像の方を元画像に合わせて拡大し、その結果を元画像に適用をしています。
これにより、サイズは元画像と同じで、ラベル部分以外が黒く塗りつぶされた画像(masked_image)が生成されます。

    contours, _ = cv2.findContours(resized_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    x, y, w, h = cv2.boundingRect(contours[0])
    min_x, min_y = x, y
    max_x, max_y = x + w, y + h
    min_x, min_y = max(0, min_x), max(0, min_y)
    max_x, max_y = min(org_width, max_x), min(org_height, max_y)
    
    bbox_width = max_x - min_x
    bbox_height = max_y - min_y

    crop_width_orig = max(bbox_width, bbox_height)
    crop_height_orig = max(bbox_width, bbox_height)

masked_imageにはラベル以外の余計な部分が大量に含まれますので、それらを極力除外し、ラベル部分がギリギリ納まる正方形の画像として切り出せるよう、位置やサイズを計算します。
まず、cv2.findContoursでマスク画像からマスクの輪郭を抽出し、cv2.boundingRectでそれをギリギリ囲む矩形の座標とサイズを算出しています。
bbox_width, bbox_heightは前述の矩形の幅と高さになり、crop_width_orig, crop_height_origはそれを包含する正方形の幅と高さになります。

なお、若干冗長な処理を行なっているように見える部分もありますが、Gemini先生なりのこだわりがあるようです(単に計算的に合っているかどうかではなく、「コードを書く人の”意図”と”思考のプロセス”を示すため、あるいは汎用的なパターンの適用のため」とのこと)。

    center_x = min_x + bbox_width // 2
    center_y = min_y + bbox_height // 2

    crop_x1 = center_x - crop_width_orig // 2
    crop_y1 = center_y - crop_height_orig // 2
    crop_x2 = crop_x1 + crop_width_orig
    crop_y2 = crop_y1 + crop_height_orig

先の計算結果を元に、ラベル部分が中心に位置するように考慮しつつ、切り出す正方形の座標を算出しています。

    pad_left = max(0, -crop_x1)
    pad_top = max(0, -crop_y1)
    pad_right = max(0, crop_x2 - org_width)
    pad_bottom = max(0, crop_y2 - org_height)
    padded_masked_image = cv2.copyMakeBorder(masked_image, 
                                               pad_top, pad_bottom, 
                                               pad_left, pad_right, 
                                               cv2.BORDER_CONSTANT, value=[0, 0, 0])

元画像からラベル部分を含む正方形を切り出す場合、ラベル部分が元画像の端の方に位置していた場合などは切り出す正方形が元画像内に納まらない可能性があります。
よって、正方形が元画像からはみ出す場合は、その分元画像を拡張し、追加した部分は黒く塗りつぶすという処理をしています。

蛇足ながら、上記可能性については、個人的には盲点でした。
さすがGemini先生。

    crop_x1_padded = crop_x1 + pad_left
    crop_y1_padded = crop_y1 + pad_top
    crop_x2_padded = crop_x2 + pad_left
    crop_y2_padded = crop_y2 + pad_top
    cropped_image = padded_masked_image[crop_y1_padded:crop_y2_padded, crop_x1_padded:crop_x2_padded]

先の処理で元画像の形が変わっている可能性があり、その場合は切り抜く範囲の座標も変わっているので、その点を考慮して座標を補正しつつ、ラベル部分(cropped_image)を切り抜きます。

    final_cropped_image = cv2.resize(cropped_image,
                                     (target_crop_size, target_crop_size),
                                     interpolation=cv2.INTER_AREA)

切り抜いたラベル部分(正方形)のサイズは様々なので、これを所定のサイズ(デフォルトは224×224)にリサイズします。

上記処理により、元画像(org_image)においてラベル部分がどのようなサイズでどのような位置に写っていても、ラベル部分が中央に、かつ最大サイズで写された同じサイズの画像(final_cropped_image)として加工されることになります。

ラベル画像の保存

前述した処理を実行しつつ、生成されたラベル画像を所定のディレクトリ配下に保存します。

import os
import cv2

from get_model import get_model
from get_mask import get_mask
from apply_mask_and_crop import apply_mask_and_crop

MODEL_PATH = "./my_model/unet.keras"
INPUTS = "./inputs"
OUTPUTS = "./outputs"
CROP_SIZE = 224

def save_img(img, save_path):
    if img.ndim == 3 and img.shape[2] == 3:
        img_to_save = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    else:
        img_to_save = img

    cv2.imwrite(save_path, img_to_save)

# MAIN
os.makedirs(OUTPUTS, exist_ok=True)

model = get_model(MODEL_PATH)

image_files = os.listdir(INPUTS)
for image_file in image_files:
    print(f"Processing {image_file}...")
    image_path = os.path.join(INPUTS, image_file)

    mask = get_mask(model, image_path)

    save_path = os.path.join(OUTPUTS, os.path.splitext(image_file)[0] + '_mask.png')
    save_img(mask, save_path)

    cropped_masked_image = apply_mask_and_crop(image_path, mask, CROP_SIZE)

    base_name, original_ext = os.path.splitext(image_file)
    save_path_cropped = os.path.join(OUTPUTS, base_name + '_cropped' + original_ext)
    save_img(cropped_masked_image, save_path_cropped)

セグメンテーションに使用するモデルは「./my_model/unet.keras」として配置しておきます。
今回は「ラベル画像のトリミング(6)」でval_lossの値が最も小さかった(セグメンテーション結果も総合的に良さそうな)「学習率減衰」モデルを使用することにします。

実行結果

ラベル画像の切り出し結果を以下に示します。
上が元画像で下が切り出し結果です。

セット1

セット2

セット3

セット4

まとめ

率直に言って、かなり良いのではないでしょうか!
一部ラベル以外の部分を含んでしまっているケースもありますが、ベクトル化を想定して画像内のラベル部分を同等のサイズ・構図で画像化すると言う目的に照らせば、十分満足できるレベルではないかと思っています。
特にセット4などはベストに近い結果かと自負してます。

かつて「画像検索の精度確認(5)」においてrembgを使用したラベルの切り抜きに挫折し、独自に「セマンティックセグメンテーション」用モデルを構築すると決意してから幾星霜(実は4ヶ月程度)、やっとここまで来たかと感慨深いです。

次回はいよいよ「ベクトル化した上での検索精度に関して確認してみる」にチャレンジしたいと思います。

Related articles

フロント担当のTIL 2025/06

「Today I Learned(今日学んだこと)」を記録するTIL、6月は新マシンが納品され、GPU仮想化からローカルLLM環境の構築まで、いろいろ試行錯誤しました。

Mac+PostgreSQLでベクトル検索

前回の投稿で予告したように、今回から...

ラベル画像のトリミング(6)

「ラベル画像のトリミング(5)」では...

EC-CUBE + Vagrant + Xdebu...

EC-CUBEのデバッグで今まで使用し...