"""단독 프로세스로 ImageSegmenter(selfie_multiclass_256x256)를 실행해
category_mask를 .npy로 저장한다. (세그폴트 회피를 위해 FaceLandmarker와 분리.)

사용: python _run_segmentation.py <이미지경로> <출력npy경로>
"""
from __future__ import annotations

import sys
from pathlib import Path

import mediapipe as mp
import numpy as np
from mediapipe.tasks import python as mp_python
from mediapipe.tasks.python import vision
from PIL import Image

MODEL_SEG = Path(__file__).resolve().parent / "_models" / "selfie_multiclass_256x256.tflite"


def main() -> int:
    image_path = Path(sys.argv[1])
    out_path = Path(sys.argv[2])

    base_options = mp_python.BaseOptions(model_asset_path=str(MODEL_SEG))
    options = vision.ImageSegmenterOptions(
        base_options=base_options, output_confidence_masks=False, output_category_mask=True
    )
    im = Image.open(image_path).convert("RGB")
    arr = np.array(im)
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=arr)
    with vision.ImageSegmenter.create_from_options(options) as seg:
        result = seg.segment(mp_image)

    if result.category_mask is None:
        print("SEGMENTATION_FAILED")
        return 1

    category_mask = result.category_mask.numpy_view()[:, :, 0].copy()
    np.save(out_path, category_mask)
    print(f"SEGMENTATION_OK shape={category_mask.shape}")
    return 0


if __name__ == "__main__":
    sys.exit(main())