dataset_aug.py

"""
画像の拡張処理を行うための関数群を提供します。

1. 画像の平均色を計算する関数
2. 画像を指定された角度で回転させ、平均色で余白を埋める関数
3. 回転した画像から最大の長方形を切り出す関数
4. ランダムな正方形を切り出す関数
"""

import os
from PIL import Image, ImageStat, ImageOps
from collections import Counter
import random
import math
from tqdm import tqdm
import argparse

def get_average_color(image):
    """画像の平均色を取得"""
    # RGBAの場合はRGBに変換
    if image.mode == 'RGBA':
        image = image.convert('RGB')
    
    stat = ImageStat.Stat(image)
    r, g, b = map(int, stat.mean)
    return (r, g, b)

def get_edge_mode_color(img, edge_width=10):
    """画像の外周の最頻値（mode）を取得する"""
    # 外周の10ピクセル領域を取得
    left = img.crop((0, 0, edge_width, img.height))  # 左端
    right = img.crop((img.width - edge_width, 0, img.width, img.height))  # 右端
    top = img.crop((0, 0, img.width, edge_width))  # 上端
    bottom = img.crop((0, img.height - edge_width, img.width, img.height))  # 下端

    # 各領域のピクセルデータを取得して結合
    colors = list(left.getdata()) + list(right.getdata()) + list(top.getdata()) + list(bottom.getdata())
    # 最頻値（mode）を計算
    mode_color = Counter(colors).most_common(1)[0][0]  # 最も頻繁に出現する色を取得

    return mode_color

def rotate_image(image, angle, fill_color=(255, 255, 255)):
    """画像を指定された角度で回転させ、指定された色で余白を埋める"""
    return image.rotate(angle, expand=True, fillcolor=fill_color)

def crop_square(cropped_rect_image, left, top, crop_size):
    """ランダムな正方形を切り出す"""
    return cropped_rect_image.crop((left, top, left + crop_size, top + crop_size))

def apply_random_flip(image, is_horizontal):
    """画像にランダムなフリップ（水平または垂直）を適用する"""
    if is_horizontal:
        return ImageOps.mirror(image)  # 水平フリップ
    return image

def process_image_pair(
    source_image,
    target_image,
    output_size=(1024, 1024),
    is_flip=False,
    rotation_range=40,
    min_scale=0.6,
    max_scale=1.2,
    source_is_avg_color_fill=False,
    source_is_edge_mode_fill=False,
    target_is_avg_color_fill=True,
    target_is_edge_mode_fill=False,
    expand_to_long_side=False
    ):
    """1組の画像に対して拡張処理を行う"""
    orig_source_width, orig_source_height = source_image.size
    orig_target_width, orig_target_height = target_image.size
    
    # ソース画像の余白の色を決定
    if source_is_edge_mode_fill:
        source_fill_color = get_edge_mode_color(source_image, edge_width=10)
    elif source_is_avg_color_fill:
        source_fill_color = get_average_color(source_image)
    else:
        source_fill_color = (255, 255, 255)
    
    # ターゲット画像の余白の色を決定
    if target_is_edge_mode_fill:
        target_fill_color = get_edge_mode_color(target_image, edge_width=10)
    elif target_is_avg_color_fill:
        target_fill_color = get_average_color(target_image)
    else:
        target_fill_color = (255, 255, 255)

    base_source = source_image
    base_target = target_image

    # 長辺を基準にする場合の処理を追加
    if expand_to_long_side:
        # sourceの長辺を取得して正方形のキャンバスを作成
        source_long_side = max(base_source.width, base_source.height)
        source_canvas = Image.new("RGB", (source_long_side, source_long_side), source_fill_color)
        # 中央に配置
        source_paste_x = (source_long_side - base_source.width) // 2
        source_paste_y = (source_long_side - base_source.height) // 2
        source_canvas.paste(base_source, (source_paste_x, source_paste_y))
        base_source = source_canvas

        # targetも同様に処理
        target_long_side = max(base_target.width, base_target.height)
        target_canvas = Image.new("RGB", (target_long_side, target_long_side), target_fill_color)
        target_paste_x = (target_long_side - base_target.width) // 2
        target_paste_y = (target_long_side - base_target.height) // 2
        target_canvas.paste(base_target, (target_paste_x, target_paste_y))
        base_target = target_canvas

    if rotation_range > 0:
        angle = random.uniform(-rotation_range, rotation_range)
        rotated_source = rotate_image(source_image, angle, source_fill_color)
        rotated_target = rotate_image(target_image, angle, target_fill_color)
        base_source = rotated_source
        base_target = rotated_target

    if is_flip:
        is_horizontal = random.choice([True, False])
        flipped_source = apply_random_flip(base_source, is_horizontal)
        flipped_target = apply_random_flip(base_target, is_horizontal)
        base_source = flipped_source
        base_target = flipped_target

    scale = random.uniform(min_scale, max_scale)
    canvas_scale = 1/scale

    if canvas_scale > 1.0:
        # 新規画像(canvas)を作成し中心に画像を配置
        scaled_source = Image.new("RGB", (int(base_source.width*canvas_scale), int(base_source.height*canvas_scale)), source_fill_color)
        scaled_target = Image.new("RGB", (int(base_target.width*canvas_scale), int(base_target.height*canvas_scale)), target_fill_color)
        scaled_source.paste(base_source, (int((scaled_source.width-base_source.width)/2), int((scaled_source.height-base_source.height)/2)))
        scaled_target.paste(base_target, (int((scaled_target.width-base_target.width)/2), int((scaled_target.height-base_target.height)/2)))
    else:
        scaled_source = base_source
        scaled_target = base_target

    base_source_width, base_source_height = base_source.size
    base_source_max_square_size = min(base_source_height, base_source_width)
    crop_source_size = int(base_source_max_square_size * canvas_scale)

    base_target_width, base_target_height = base_target.size
    base_target_max_square_size = min(base_target_height, base_target_width)
    crop_target_size = int(base_target_max_square_size * canvas_scale)

    scaled_source_width, scaled_source_height = scaled_source.size
    left_source = random.randint(0, scaled_source_width - crop_source_size)
    top_source = random.randint(0, scaled_source_height - crop_source_size)

    # sourceとtargetの位置合わせ. この場合、sourceとtargetのアスペクト比は同じと仮定
    left_target = left_source * orig_target_width // orig_source_width
    top_target = top_source * orig_target_height // orig_source_height

    final_source = crop_square(scaled_source, left_source, top_source, crop_source_size).resize(output_size)
    final_target = crop_square(scaled_target, left_target, top_target, crop_target_size).resize(output_size)

    return final_source, final_target

def process_images(
    source_img,
    target_img,
    num_copies,
    output_size,
    is_flip,
    rotation_range,
    min_scale,
    max_scale,
    source_is_avg_color_fill,
    source_is_edge_mode_fill,
    target_is_avg_color_fill,
    target_is_edge_mode_fill,
    expand_to_long_side
):
    aug_sources = []
    aug_targets = []
    
    for i in range(num_copies):
        # 拡張処理を実行
        aug_source, aug_target = process_image_pair(
            source_img,
            target_img,
            output_size,
            is_flip,
            rotation_range,
            min_scale,
            max_scale,
            source_is_avg_color_fill,
            source_is_edge_mode_fill,
            target_is_avg_color_fill,
            target_is_edge_mode_fill,
            expand_to_long_side
        )
        
        aug_sources.append(aug_source)
        aug_targets.append(aug_target)
    
    return aug_sources, aug_targets

if __name__ == '__main__':
    args = argparse.ArgumentParser()
    args.add_argument('source_folder', type=str, help='source画像フォルダのパス / Path to source image folder')
    args.add_argument('target_folder', type=str, help='target画像フォルダのパス / Path to target image folder')
    args.add_argument('output_folder', type=str, help='出力先フォルダのパス / Path to output folder')
    args.add_argument('--output_size', '-s', type=int, default=1024, 
                     help='出力画像の一辺のサイズ / Output image size')
    args.add_argument('--num_copies', '-n', type=int, default=1, 
                     help='出力画像の枚数 / Number of augmented copies')
    args.add_argument('--is_flip', '-f', type=bool, default=True, 
                     help='フリップを適用するかどうか / Whether to apply random flip')
    args.add_argument('--rotation_range', '-r', type=int, default=0, 
                     help='回転角度の範囲 / Range of rotation angle')
    args.add_argument('--min_scale', '--ms', type=float, default=1.0, 
                     help='最小の画像サイズ / Minimum scale of the image')
    args.add_argument('--max_scale', '--xs', type=float, default=1.0, 
                     help='最大の画像サイズ / Maximum scale of the image')
    args.add_argument('--source_is_avg_color_fill', '--sa', type=bool, default=True, 
                     help='source画像を平均色で余白を埋めるかどうか / Whether to fill source image padding with average color')
    args.add_argument('--source_is_edge_mode_fill', '--se', type=bool, default=False, 
                     help='source画像を外周の最頻値で余白を埋めるかどうか / Whether to fill source image padding with edge mode color')
    args.add_argument('--target_is_avg_color_fill', '--ta', type=bool, default=False, 
                     help='target画像を平均色で余白を埋めるかどうか / Whether to fill target image padding with average color')
    args.add_argument('--target_is_edge_mode_fill', '--te', type=bool, default=False, 
                     help='target画像を外周の最頻値で余白を埋めるかどうか / Whether to fill target image padding with edge mode color')
    args.add_argument('--expand_to_long_side', '--el', type=bool, default=False,
                     help='長辺まで拡張して正方形にするかどうか / Whether to expand the image to a square using the long side')
    args = args.parse_args()
    
    output_path_source = os.path.join(args.output_folder, 'aug_source')
    output_path_target = os.path.join(args.output_folder, 'aug_target')

    # 画像処理
    for image_name in tqdm(os.listdir(args.source_folder)):
        if image_name.endswith('.jpg'):
            source_path = os.path.join(args.source_folder, image_name)
            target_path = os.path.join(args.target_folder, image_name)
            
            process_images(
                source_path,
                target_path,
                args.num_copies,
                (args.output_size, args.output_size),
                args.is_flip,
                args.rotation_range,
                args.min_scale,
                args.max_scale,
                args.source_is_avg_color_fill,
                args.source_is_edge_mode_fill,
                args.target_is_avg_color_fill,
                args.target_is_edge_mode_fill,
                args.expand_to_long_side
            )