YOLO学習データを自動生成する｜射影変換・ラベル生成・データ拡張をまとめて自動化

import io
import os
import cv2
import random
import numpy as np
from PIL import Image, ImageEnhance

# 設定
BASE_DIR = "./sample/base_dir"
CROP_DIR = "./sample/crop_dir"
OUT_ROOT = "./sample/out_dir"

PAD = 5
EDGE_KERNEL_SIZE = 5
EDGE_BLUR_SIZE = 5
NUM = 10
MAX_DUPLICATE = 4
TRAIN_RATIO = 0.8
VAL_RATIO = 0.1
TEST_RATIO = 0.1


for split in ["train", "val", "test"]:
    os.makedirs(os.path.join(OUT_ROOT, "images", split), exist_ok=True)
    os.makedirs(os.path.join(OUT_ROOT, "labels", split), exist_ok=True)


# augmentation設定
AUGMENT_CONFIG = {
    "brightness": (0.7, 1.3),
    "contrast": (0.7, 1.3),
    "color": (0.7, 1.3),
    "sharpness": (0.5, 2.0),
    "gamma": (0.8, 1.2),
    "noise_sigma": (0, 10),
    "jpeg_quality": (50, 100),
    "rotate_deg": (-5, 5),
    "shift_x": (-0.05, 0.05),
    "shift_y": (-0.05, 0.05),
}

def clip_box(x1, y1, x2, y2, w, h):
    x1 = max(0, min(x1, w - 1))
    y1 = max(0, min(y1, h - 1))
    x2 = max(0, min(x2, w - 1))
    y2 = max(0, min(y2, h - 1))
    return x1, y1, x2, y2

def xyxy_to_yolo(x1, y1, x2, y2, w, h):
    xc = ((x1 + x2) / 2) / w
    yc = ((y1 + y2) / 2) / h
    bw = abs(x2 - x1) / w
    bh = abs(y2 - y1) / h
    return xc, yc, bw, bh

def apply_gamma(img, gamma):
    arr = np.array(img).astype(np.float32) / 255.0
    arr = np.power(arr, gamma)
    arr = np.clip(arr * 255, 0, 255).astype(np.uint8)
    return Image.fromarray(arr)

def apply_noise(img, sigma):
    arr = np.array(img).astype(np.float32)
    noise = np.random.normal(0, sigma, arr.shape)
    arr += noise
    arr = np.clip(arr, 0, 255).astype(np.uint8)
    return Image.fromarray(arr)

def apply_jpeg(img, quality):
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=int(quality))
    buf.seek(0)
    return Image.open(buf).convert("RGB")

def augment_image(img):
    img = ImageEnhance.Brightness(img).enhance(
        random.uniform(*AUGMENT_CONFIG["brightness"])
    )
    img = ImageEnhance.Contrast(img).enhance(
        random.uniform(*AUGMENT_CONFIG["contrast"])
    )
    img = ImageEnhance.Color(img).enhance(
        random.uniform(*AUGMENT_CONFIG["color"])
    )
    img = ImageEnhance.Sharpness(img).enhance(
        random.uniform(*AUGMENT_CONFIG["sharpness"])
    )

    img = apply_gamma(
        img,
        random.uniform(*AUGMENT_CONFIG["gamma"])
    )

    img = apply_noise(
        img,
        random.uniform(*AUGMENT_CONFIG["noise_sigma"])
    )

    img = apply_jpeg(
        img,
        random.randint(*AUGMENT_CONFIG["jpeg_quality"])
    )

    return img

def shift_image_and_boxes(img, labels):
    w, h = img.size
    dx = int(random.uniform(*AUGMENT_CONFIG["shift_x"]) * w)
    dy = int(random.uniform(*AUGMENT_CONFIG["shift_y"]) * h)

    shifted = Image.new("RGB", (w, h), (0, 0, 0))
    shifted.paste(img, (dx, dy))

    out = []

    for line in labels:
        cid, xc, yc, bw, bh = map(float, line.split())
        x1 = (xc - bw/2) * w + dx
        y1 = (yc - bh/2) * h + dy
        x2 = (xc + bw/2) * w + dx
        y2 = (yc + bh/2) * h + dy
        x1,y1,x2,y2 = clip_box(x1,y1,x2,y2,w,h)
        if x2 <= x1 or y2 <= y1:
            continue
        xc2,yc2,bw2,bh2 = xyxy_to_yolo(x1,y1,x2,y2,w,h)
        out.append(f"{int(cid)} {xc2:.6f} {yc2:.6f} {bw2:.6f} {bh2:.6f}")

    return shifted, out


def rotate_image_and_boxes(img, labels):
    w, h = img.size
    angle = random.uniform(*AUGMENT_CONFIG["rotate_deg"])
    rotated = img.rotate(
        angle,
        expand=False,
        resample=Image.BICUBIC
    )
    rad = np.deg2rad(-angle)
    cx = w / 2.0
    cy = h / 2.0
    cos_r = np.cos(rad)
    sin_r = np.sin(rad)
    out = []

    for line in labels:
        cid, xc, yc, bw, bh = map(float, line.split())
        # YOLO -> xyxy
        box_w = bw * w
        box_h = bh * h

        x_center = xc * w
        y_center = yc * h

        x1 = x_center - box_w / 2
        y1 = y_center - box_h / 2
        x2 = x_center + box_w / 2
        y2 = y_center + box_h / 2

        # 4頂点
        corners = np.array([
            [x1, y1],
            [x2, y1],
            [x2, y2],
            [x1, y2]
        ])

        # 中心基準へ移動
        corners[:, 0] -= cx
        corners[:, 1] -= cy

        # 回転
        rotated_corners = np.zeros_like(corners)

        rotated_corners[:, 0] = (
            corners[:, 0] * cos_r -
            corners[:, 1] * sin_r
        )

        rotated_corners[:, 1] = (
            corners[:, 0] * sin_r +
            corners[:, 1] * cos_r
        )

        # 元の座標系へ戻す
        rotated_corners[:, 0] += cx
        rotated_corners[:, 1] += cy

        # 回転後の外接矩形
        new_x1 = rotated_corners[:, 0].min()
        new_y1 = rotated_corners[:, 1].min()
        new_x2 = rotated_corners[:, 0].max()
        new_y2 = rotated_corners[:, 1].max()

        # 画像外クリップ
        new_x1, new_y1, new_x2, new_y2 = clip_box(
            new_x1,
            new_y1,
            new_x2,
            new_y2,
            w,
            h
        )

        if new_x2 <= new_x1 or new_y2 <= new_y1:
            continue

        xc2, yc2, bw2, bh2 = xyxy_to_yolo(
            new_x1,
            new_y1,
            new_x2,
            new_y2,
            w,
            h
        )

        out.append(f"{int(cid)} {xc2:.6f} {yc2:.6f} {bw2:.6f} {bh2:.6f}")

    return rotated, out



def load_quads(txt_path):
    quads = []

    if not os.path.exists(txt_path):
        return quads

    with open(txt_path, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            vals = list(map(float, line.split(",")))

            if len(vals) != 8:
                continue

            quad = np.float32([
                [vals[0], vals[1]],
                [vals[2], vals[3]],
                [vals[4], vals[5]],
                [vals[6], vals[7]],
            ])

            quads.append(quad)
    return quads


def get_surrounding_color(base_img, dst_pts):
    h, w = base_img.shape[:2]

    x_min = int(np.min(dst_pts[:, 0]))
    x_max = int(np.max(dst_pts[:, 0]))

    y_min = int(np.min(dst_pts[:, 1]))
    y_max = int(np.max(dst_pts[:, 1]))

    margin = 20

    x1 = max(0, x_min - margin)
    x2 = min(w, x_max + margin)

    y1 = max(0, y_min - margin)
    y2 = min(h, y_max + margin)

    region = base_img[y1:y2, x1:x2]

    if region.size == 0:
        return (127, 127, 127)

    mean_color = tuple(
        int(v)
        for v in region.mean(axis=(0, 1))
    )

    return mean_color


def paste_perspective(base_img, crop_img, dst_pts, pad=100,):
    base_h, base_w = base_img.shape[:2]
    crop_h, crop_w = crop_img.shape[:2]

    mean_color = get_surrounding_color(
        base_img,
        dst_pts
    )

    crop_pad = cv2.copyMakeBorder(
        crop_img,
        pad,
        pad,
        pad,
        pad,
        cv2.BORDER_CONSTANT,
        value=mean_color
    )

    pad_h, pad_w = crop_pad.shape[:2]

    mask = np.zeros((pad_h, pad_w), dtype=np.uint8)

    mask[pad:pad + crop_h, pad:pad + crop_w] = 255


    src_pts = np.float32([
        [pad, pad],
        [pad + crop_w - 1, pad],
        [pad + crop_w - 1, pad + crop_h - 1],
        [pad, pad + crop_h - 1]
    ])

    H = cv2.getPerspectiveTransform(
        src_pts,
        dst_pts
    )

    warped_img = cv2.warpPerspective(
        crop_pad,
        H,
        (base_w, base_h),
        flags=cv2.INTER_LINEAR
    )

    warped_mask = cv2.warpPerspective(
        mask,
        H,
        (base_w, base_h),
        flags=cv2.INTER_NEAREST
    )

    result = base_img.copy()

    result[warped_mask > 0] = warped_img[warped_mask > 0]

    kernel = np.ones(
        (
            EDGE_KERNEL_SIZE,
            EDGE_KERNEL_SIZE
        ),
        np.uint8
    )

    edge_mask = cv2.morphologyEx(
        warped_mask,
        cv2.MORPH_GRADIENT,
        kernel
    )

    blurred = cv2.GaussianBlur(result, (EDGE_BLUR_SIZE,EDGE_BLUR_SIZE), 0)

    result[edge_mask > 0] = blurred[
        edge_mask > 0
    ]

    return result


def quad_to_yolo(quad, img_w, img_h):
    x_min = np.min(quad[:, 0]) - 2
    x_max = np.max(quad[:, 0]) + 2

    y_min = np.min(quad[:, 1]) - 10
    y_max = np.max(quad[:, 1]) + 10

    x_center = (x_min + x_max) / 2.0
    y_center = (y_min + y_max) / 2.0

    width = x_max - x_min
    height = y_max - y_min

    x_center /= img_w
    y_center /= img_h

    width /= img_w
    height /= img_h

    return (
        x_center,
        y_center,
        width,
        height
    )


def get_random_split():
    r = random.random()
    if r < TRAIN_RATIO:
        return "train"
    elif r < TRAIN_RATIO + VAL_RATIO:
        return "val"
    else:
        return "test"



crop_files = sorted([
    os.path.join(CROP_DIR, f)
    for f in os.listdir(CROP_DIR)
    if f.lower().endswith(".jpg")
])

if len(crop_files) == 0:
    raise RuntimeError(
        "crop jpg が見つかりません"
    )


base_files = sorted([
    f
    for f in os.listdir(BASE_DIR)
    if f.lower().endswith(".jpg")
])

for idx, file_name in enumerate(base_files):
    print(f"{idx} / {len(base_files)}")

    base_path = os.path.join(BASE_DIR,file_name)
    txt_path = os.path.splitext(base_path)[0] + ".txt"

    if not os.path.exists(txt_path):
        print("skip (txtなし):", file_name)
        continue

    print("processing:", file_name)

    base_img = cv2.imread(base_path)

    if base_img is None:
        print("skip (読込失敗):", file_name)
        continue

    quads = load_quads(txt_path)
    stem = os.path.splitext(file_name)[0]

    for aug_idx in range(NUM):
        result_img = base_img.copy()
        # crop使用回数管理
        crop_use_count = {
            path: 0
            for path in crop_files
        }

        label_lines = []
        for quad in quads:
            available_crops = [
                path
                for path in crop_files
                if crop_use_count[path] < MAX_DUPLICATE
            ]

            if len(available_crops) == 0:
                break

            crop_path = random.choice(available_crops)

            crop_use_count[crop_path] += 1

            crop_img = cv2.imread(crop_path)

            if crop_img is None:
                continue


            if random.random() < 0.5:

                crop_img = cv2.rotate(
                    crop_img,
                    cv2.ROTATE_180
                )

            result_img = paste_perspective(
                result_img,
                crop_img,
                quad,
                pad=PAD
            )
            
            img_h, img_w = result_img.shape[:2]

            x_center, y_center, width, height = quad_to_yolo(
                quad,
                img_w,
                img_h
            )

            label_lines.append(
                f"0 "
                f"{x_center:.6f} "
                f"{y_center:.6f} "
                f"{width:.6f} "
                f"{height:.6f}"
            )
        
        
        # 合成後に画像処理で変換
        result_img
        img = Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))
        img = augment_image(img)
        img, new_labels = shift_image_and_boxes(img, label_lines)
        img, new_labels_ = rotate_image_and_boxes(img, new_labels)
         
        split = get_random_split() 

        img.save(os.path.join(OUT_ROOT, "images", split, f"{stem}_{aug_idx:03d}.jpg"))

        with open(os.path.join(OUT_ROOT, "labels", split, f"{stem}_{aug_idx:03d}.txt"), "w") as f:
            f.write("\n".join(new_labels_))

print("finished")