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

# 입력 / 출력 폴더
INPUT_DIR = "test2"
OUTPUT_DIR = "test3"

os.makedirs(OUTPUT_DIR, exist_ok=True)

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) 배경 제거 + 사람 가운데 정렬 + 사람 최대 확대 ----------

MARGIN_RATIO = 0.7      # 여백 비율
SHOULDER_RATIO = 0.6    # 사람 전체 높이 중 위에서부터 어느 정도까지 남길지 (0.5~0.6 정도 조절)

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)

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

            alpha = result.getchannel("A")
            bbox = alpha.getbbox()  # (left, upper, right, lower) 또는 None

            if bbox is None:
                # 사람이 안 잡힌 경우: 전체 이미지 기준으로 상단 일부만 사용
                left, upper, right, lower = 0, 0, result.width, result.height
            else:
                left, upper, right, lower = bbox

            # ----- 🔥 어깨까지만 남기도록 아래쪽 잘라내기 -----
            box_height = lower - upper
            # 위에서부터 SHOULDER_RATIO 비율 만큼만 남김 (예: 55%)
            new_lower = upper + int(box_height * SHOULDER_RATIO)
            # 혹시라도 이상하게 계산되어 높이가 0이 되지 않도록 최소 보정
            if new_lower <= upper:
                new_lower = lower

            cropped = result.crop((left, upper, right, new_lower))
            cw, ch = cropped.size

            # ----- 캔버스에 "사람이 꽉 차게" 배치 -----
            scale_w = (target_width * MARGIN_RATIO) / cw
            scale_h = (target_height * MARGIN_RATIO) / ch
            scale = min(scale_w, scale_h)

            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)

            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))