import os
from rembg import remove, new_session
from PIL import Image

# 입력 / 출력 폴더
INPUT_DIR = "frames"
OUTPUT_DIR = "rembg"

# 출력 폴더 없으면 생성
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 사람 전용 세그멘테이션 모델 세션 (u2net_human_seg)
session = new_session("u2net_human_seg")

# 처리할 이미지 확장자들
VALID_EXT = (".png", ".jpg", ".jpeg", ".bmp", ".webp")

# ---------- 1) 타겟 사이즈 결정 (첫 이미지 기준) ----------
target_width = None
target_height = None

for filename in os.listdir(INPUT_DIR):
    if not filename.lower().endswith(VALID_EXT):
        continue
    first_path = os.path.join(INPUT_DIR, filename)
    with Image.open(first_path) as img:
        target_width, target_height = img.size
    break

if target_width is None:
    print("입력 폴더에 유효한 이미지가 없습니다.")
    exit(1)

print(f"타겟 캔버스 크기: {target_width} x {target_height}")

# ---------- 2) 배경 제거 + 사람 가운데 정렬 + 사람 최대 확대 ----------

# 사람 주변에 아주 살짝 여백을 남기고 싶다면 (1.0 이면 여백 없음)
MARGIN_RATIO = 0.95  # 0.9 ~ 0.98 정도에서 조절

for filename in sorted(os.listdir(INPUT_DIR)):
    if not filename.lower().endswith(VALID_EXT):
        continue  # 이미지가 아니면 스킵

    input_path = os.path.join(INPUT_DIR, filename)

    try:
        with Image.open(input_path) as img:
            # rembg로 배경 제거 (사람 위주로 남김)
            result = remove(img, session=session)

            # RGBA 보장
            if result.mode != "RGBA":
                result = result.convert("RGBA")

            # ----- 알파 채널 기준으로 사람 영역(bbox) 구하기 -----
            alpha = result.getchannel("A")
            bbox = alpha.getbbox()  # (left, upper, right, lower) 또는 None

            if bbox is None:
                # 사람이 안 잡힌 경우: 그냥 전체 이미지를 사용
                cropped = result
            else:
                # 사람 영역만 crop
                cropped = result.crop(bbox)

            cw, ch = cropped.size

            # ----- ✨ 고정 크기 캔버스에 "사람이 꽉 차게" 배치 -----
            # 캔버스 내부에 사람이 최대한 크게 들어가도록 확대/축소 (비율 유지)
            # MARGIN_RATIO 로 아주 약간의 여백을 남김
            scale_w = (target_width * MARGIN_RATIO) / cw
            scale_h = (target_height * MARGIN_RATIO) / ch
            scale = min(scale_w, scale_h)

            # (원하면 너무 큰 확대 방지용으로 최대 배율 제한 가능)
            # 예: scale = min(scale, 3.0)

            new_cw = int(cw * scale)
            new_ch = int(ch * scale)

            # 고품질 리사이즈
            resized = cropped.resize((new_cw, new_ch), Image.LANCZOS)

            # 최종 캔버스 (투명 배경)
            canvas = Image.new("RGBA", (target_width, target_height), (0, 0, 0, 0))

            # 가운데 위치 계산
            offset_x = (target_width  - new_cw) // 2
            offset_y = (target_height - new_ch) // 2

            # 캔버스에 붙여넣기
            canvas.paste(resized, (offset_x, offset_y), resized)

            # 배경 제거 후이므로 png 저장 추천
            base_name, _ = os.path.splitext(filename)
            output_filename = base_name + ".png"
            output_path = os.path.join(OUTPUT_DIR, output_filename)

            canvas.save(output_path)
            print(f"[OK] {input_path} -> {output_path}")

    except Exception as e:
        print(f"[ERROR] {input_path}: {e}")

print("✅ 배경 제거 + 사람 확대 + 가운데 정렬 완료! 결과 폴더:", os.path.abspath(OUTPUT_DIR))