증강 20 → 100

"""
연구실 구역(A/B/C) 이미지 증강 스크립트
20장씩 → 100장씩으로 늘림 (원본 20장 + 증강 80장)

설치:
    pip install albumentations opencv-python

폴더 구조 (입력):
    images_airlab/
        A/  (20장)
        B/  (20장)
        C/  (20장)

폴더 구조 (출력):
    images_augmented/
        A/  (100장: orig_*.jpg 20장 + aug_*.jpg 80장)
        B/  (100장)
        C/  (100장)

증강 후에는 이 images_augmented/ 폴더를 기존 train.json 생성 스크립트의
IMAGE_DIR로 넣어서 그대로 쓰면 됨.

"""

import os
import math
import shutil
import cv2
import albumentations as A

INPUT_DIR = "./images_airlab"
OUTPUT_DIR = "./images_augmented" 
TARGET_COUNT = 100 
ZONE_FOLDERS = ["A", "B", "C"]

# 실제로 사람이 다른 각도/조명/거리에서 찍었을 때 생길 수 있는 변화 범위로 제한함.
#
# - RandomBrightnessContrast : 조명 차이 시뮬레이션
# - Affine   : 약간의 회전 + 이동 + 스케일 + 기울임(shear)
#  → 카메라를 살짝 틀어서 찍은 효과

# - Perspective : 모서리를 살짝 늘리는 원근 왜곡
#   → 비스듬히 본 효과

# - RandomResizedCrop : 확대/크롭 → 카메라 거리 차이 시뮬레이션
#   (scale을 0.8~1.0으로 제한해 과도한 크롭 방지)

# - GaussianBlur (약하게) : 살짝 흔들린 사진 대비용, 너무 세게 주지 않음
# - HueSaturationValue (약하게): 색온도 차이 정도만 미세 조정

transform = A.Compose([
    A.RandomBrightnessContrast(
        brightness_limit=0.2, contrast_limit=0.2, p=0.7
    ),
    A.Affine(
        scale=(0.92, 1.08),
        translate_percent=(-0.04, 0.04),
        rotate=(-8, 8),
        shear=(-5, 5),
        p=0.7,
    ),
    A.Perspective(
        scale=(0.02, 0.06), p=0.4
    ),
    A.RandomResizedCrop(
        size=(720, 1280),
        scale=(0.85, 1.0),
        ratio=(0.9, 1.1),
        p=0.6,
    ),
    A.GaussianBlur(blur_limit=(3, 3), p=0.15),
    A.HueSaturationValue(
        hue_shift_limit=5, sat_shift_limit=10, val_shift_limit=10, p=0.3
    ),
])

def list_images(folder):
    exts = (".jpg", ".jpeg", ".png")
    return sorted(f for f in os.listdir(folder) if f.lower().endswith(exts))

def augment_zone(zone_name):
    src_dir = os.path.join(INPUT_DIR, zone_name)
    dst_dir = os.path.join(OUTPUT_DIR, zone_name)
    os.makedirs(dst_dir, exist_ok=True)

    image_files = list_images(src_dir)
    n_original = len(image_files)
    if n_original == 0:
        print(f"[{zone_name}] 이미지가 없음. 건너뜀.")
        return

    # 원본 그대로 복사
    for idx, fname in enumerate(image_files, start=1):
        src_path = os.path.join(src_dir, fname)
        dst_name = f"orig_{idx:03d}{os.path.splitext(fname)[1]}"
        shutil.copy2(src_path, os.path.join(dst_dir, dst_name))

    # 부족한 만큼 증강 생성 
    need = max(TARGET_COUNT - n_original, 0)
    if need == 0:
        print(f"[{zone_name}] 이미 {n_original}장으로 목표 충족, 증강 생략")
        return

    # 원본 이미지마다 몇 장씩 증강을 만들지 균등 분배
    per_image = math.ceil(need / n_original)

    aug_count = 0
    for idx, fname in enumerate(image_files, start=1):
        src_path = os.path.join(src_dir, fname)
        img = cv2.imread(src_path)
        if img is None:
            print(f"  ! 읽기 실패: {src_path}")
            continue
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        for j in range(per_image):
            if aug_count >= need:
                break
            augmented = transform(image=img_rgb)["image"]
            out_bgr = cv2.cvtColor(augmented, cv2.COLOR_RGB2BGR)

            aug_count += 1
            out_name = f"aug_{idx:03d}_{j+1:02d}.jpg"
            cv2.imwrite(os.path.join(dst_dir, out_name), out_bgr)

        if aug_count >= need:
            break

    total = n_original + aug_count
    print(f"[{zone_name}] 원본 {n_original}장 + 증강 {aug_count}장 = 총 {total}장 → {dst_dir}")

def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    for zone in ZONE_FOLDERS:
        augment_zone(zone)
    print("\n증강 완료")

if __name__ == "__main__":
    main()