"""
YOLOv8 추론 스크립트 (Ultralytics)
=====================================
COCO 사전학습 YOLOv8로 이미지에서 객체를 탐지하고 결과를 저장한다.
앞서 만든 밑바닥 Faster R-CNN(infer.py)과 같은 이미지로 비교해보기 좋다.

설치:
  pip install ultralytics

실행 예:
  # 이미지 한 장
  python yolo_infer.py --image test.jpg

  # 폴더 전체
  python yolo_infer.py --image_dir ./samples --out_dir ./yolo_results

  # 모델 크기 변경 (n < s < m < l < x, 뒤로 갈수록 정확·무거움)
  python yolo_infer.py --image test.jpg --model yolov8s.pt

  # 점수 임계값 조정 (기본 0.25)
  python yolo_infer.py --image test.jpg --conf 0.4

메모:
  - 처음 실행 시 모델 가중치(yolov8n.pt 등)가 자동 다운로드된다.
  - Faster R-CNN(20 VOC 클래스)과 달리 YOLOv8은 COCO 80 클래스를 탐지한다.
  - 라이선스: Ultralytics YOLO는 AGPL-3.0. 상업적 폐쇄소스 사용 시 상용 라이선스 필요.
"""

import os
import argparse
from ultralytics import YOLO


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="yolov8n.pt",
                    help="yolov8n/s/m/l/x.pt (없으면 자동 다운로드)")
    ap.add_argument("--image", help="단일 이미지 경로")
    ap.add_argument("--image_dir", help="이미지 폴더 경로")
    ap.add_argument("--out_dir", default="./yolo_results", help="결과 저장 폴더")
    ap.add_argument("--conf", type=float, default=0.25,
                    help="이 신뢰도 이상만 표시 (기본 0.25)")
    args = ap.parse_args()

    # 모델 로드 (COCO 사전학습, 첫 실행 시 자동 다운로드)
    print(f"모델 로드: {args.model}")
    model = YOLO(args.model)

    # 처리할 이미지 목록 구성
    targets = []
    if args.image:
        targets.append(args.image)
    if args.image_dir:
        for fn in os.listdir(args.image_dir):
            if fn.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
                targets.append(os.path.join(args.image_dir, fn))

    if not targets:
        print("이미지를 지정하세요: --image 또는 --image_dir")
        return

    os.makedirs(args.out_dir, exist_ok=True)

    for path in targets:
        # 추론 실행
        results = model(path, conf=args.conf, verbose=False)
        r = results[0]

        # 탐지 개수와 클래스별 요약
        n = len(r.boxes)
        names = r.names  # 클래스 인덱스 → 이름
        detected = {}
        for cls_id in r.boxes.cls.tolist():
            name = names[int(cls_id)]
            detected[name] = detected.get(name, 0) + 1

        # 박스 그려서 저장
        out_path = os.path.join(args.out_dir, "yolo_" + os.path.basename(path))
        r.save(filename=out_path)

        summary = ", ".join(f"{k}×{v}" for k, v in detected.items()) or "없음"
        print(f"  {os.path.basename(path)}: {n}개 탐지 ({summary}) → {out_path}")

    print(f"완료. 결과는 {args.out_dir} 폴더에 저장됨.")


if __name__ == "__main__":
    main()