画像分類・EfficientNetの学習データを作る｜画像分割とデータ拡張

from datetime import datetime


# =========================
# サンプル設定
# =========================
TARGET_SAMPLE = [
    "./image_classification/01_data/20260627_123025",
    "./image_classification/01_data/20260710_120118",
    "./image_classification/01_data/20260714_214916",
    "./image_classification/01_data/20260722_194307",
    "./image_classification/01_data/20260725_133610",
    "./image_classification/01_data/20260729_214106",
    "./image_classification/01_data/20260729_215809",
]


# クラス一覧
CLASSES = [
    "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
    "10", "11", "12", "13", "14", "15", "16", "17", "18", "19",
    "20", "21", "22", "23", "24", "25", "26", "27", "28", "29",
    "30", "31", "32", "33", "34", "35", "36"
]

# 出力先
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
DATASET_PATH = f"./image_classification/02_augmentation/{timestamp}"

# 分割比率
DATASET_RATIO = [0.8, 0.1, 0.1]
# クラス毎の最終サンプル数
TARGET_SAMPLES_PER_CLASS = 1000
# 画像サイズ
IMAGE_SIZE = 224

import os
import random
from pathlib import Path
from PIL import Image
import torchvision.transforms as transforms

import train_utils.config as conf

TRAIN_NUM = int(conf.TARGET_SAMPLES_PER_CLASS * conf.DATASET_RATIO[0])
VAL_NUM = int(conf.TARGET_SAMPLES_PER_CLASS * conf.DATASET_RATIO[1])
TEST_NUM = conf.TARGET_SAMPLES_PER_CLASS - TRAIN_NUM - VAL_NUM

class RandomRotate180:
    def __call__(self, img):
        if random.random() < 0.5:
            return img.rotate(180)
        return img

augment_transform = transforms.Compose([
    transforms.Resize((conf.IMAGE_SIZE, conf.IMAGE_SIZE)),
    RandomRotate180(),
    transforms.RandomRotation(5),
    transforms.RandomAffine(
        degrees=0,
        translate=(0.03, 0.03),
        scale=(0.95, 1.05)
    ),
    transforms.ColorJitter(
        brightness=0.25,
        contrast=0.25,
        saturation=0.25,
        hue=0.02
    ),
])


# ======================================
# 元画像を分割
# ======================================
def split_original_images(files):
    random.shuffle(files)
    total = len(files)
    train_end = int(total * conf.DATASET_RATIO[0])
    val_end = train_end + int(total * conf.DATASET_RATIO[1])

    return (
        files[:train_end],
        files[train_end:val_end],
        files[val_end:]
    )


# ======================================
# splitデータ生成
# ======================================
def generate_split_dataset(miss, correct):
    def generate(t_samples, target_count):
        generated = []
        originals = []

        for img_path in t_samples:
            img = Image.open(str(img_path)).convert("RGB")

            img = img.resize((conf.IMAGE_SIZE, conf.IMAGE_SIZE))
            originals.append((f"{img_path.stem}_orig.jpg",img))

        generated.extend(originals)

        # --------------------------
        # 元画像が目標以上
        # --------------------------
        if len(generated) >= target_count:
            return random.sample(generated, target_count)

        # --------------------------
        # 不足分Aug
        # --------------------------
        need = target_count - len(generated)
        for i in range(need):
            img_path = random.choice(t_samples)
            img = Image.open(str(img_path)).convert("RGB")
            aug_img = augment_transform(img)
            generated.append((f"{img_path.stem}_aug_{i}.jpg", aug_img))

        return generated



    samples = miss + correct

    if len(samples) >= conf.TARGET_SAMPLES_PER_CLASS:
        if len(miss) >= conf.TARGET_SAMPLES_PER_CLASS:
            samples = random.sample(miss, conf.TARGET_SAMPLES_PER_CLASS)
        else:
            samples = miss + random.sample(correct, conf.TARGET_SAMPLES_PER_CLASS - len(miss))
    
    
    random.shuffle(samples)
    total = len(samples)
    train_end = int(total * conf.DATASET_RATIO[0])
    val_end = train_end + int(total * conf.DATASET_RATIO[1])
    
    train_samples = samples[:train_end]
    val_samples = samples[train_end:val_end]
    test_samples = samples[val_end:]
    
    train_samples = generate(train_samples, TRAIN_NUM)
    val_samples = generate(val_samples, VAL_NUM)
    test_samples = generate(test_samples, TEST_NUM)
    
    return train_samples, val_samples, test_samples

    


def augment():
    for cls_name in conf.CLASSES:
        ###### ディレクトリ作成
        for split in ["train", "val", "test"]:
            path = os.path.join(conf.DATASET_PATH, split, cls_name)
            os.makedirs(path, exist_ok=True)

        ###### データセット分割

        # 1. 対象サンプルを走査し、数を取得
        miss_img_paths = []
        correct_img_paths = []
        for target_dir in conf.TARGET_SAMPLE:
            target_dir = Path(target_dir)
            miss_class = target_dir / "miss" / cls_name
            correct_class = target_dir / "correct" / cls_name

            miss_img_paths.extend(sorted(miss_class.rglob("*.jpg")))
            correct_img_paths.extend(sorted(correct_class.rglob("*.jpg")))
 
        print(f"{cls_name}_count: [miss: {len(miss_img_paths)}, correct: {len(correct_img_paths)}]")
        
        # 2. miss, correctの数に応じて処理を行う
        train, val, test = generate_split_dataset(miss_img_paths, correct_img_paths)

        
        # 3. 保存
        save_split_data(train,"train",cls_name)
        save_split_data(val,"val",cls_name)
        save_split_data(test,"test",cls_name)



# ======================================
# 保存
# ======================================
def save_split_data(data,split_name,cls_name):
    output_dir = os.path.join(conf.DATASET_PATH, split_name, cls_name)
    for file_name, img in data:
        save_path = os.path.join(output_dir,file_name)
        img.save(save_path, quality=95)