# %% [셀 1: 초기 설정 및 라이브러리 임포트]
import os
import json
import glob

# (!!!) 자주 바꿀 수 있는 값: 기본 경로 및 파일명 설정
# 이미지들이 저장된 루트 경로 (하위 폴더 포함)
COMMON_DIR = "/ltb"
IMAGES_ROOT = COMMON_DIR + "/media/ltb/90887173887158A4/Use/si/it1_worker_helmet/images"
# 개별 annotation 파일들이 있는 경로
ANNOTATIONS_DIR = COMMON_DIR + "/media/ltb/90887173887158A4/Use/si/it1_worker_helmet/releases/it1-worker-helmet-250326-4class-background/annotations"
# 최종 COCO 형식의 통합 JSON 파일이 저장될 경로 및 파일명
OUTPUT_JSON_PATH = COMMON_DIR + "/media/ltb/90887173887158A4/Use/si/it1_worker_helmet/coco_annotations.json"

print("이미지 경로:", IMAGES_ROOT)
print("Annotation 폴더:", ANNOTATIONS_DIR)
print("출력 JSON 파일:", OUTPUT_JSON_PATH)

# %% [셀 2: Annotation 파일 불러오기 및 COCO 형식 데이터 구성]
# COCO 포맷용 기본 리스트 생성
coco_images = []       # 이미지 정보 목록
coco_annotations = []  # annotation 정보 목록
category_mapping = {}  # 카테고리 이름 -> id 매핑 (동적으로 생성)
next_category_id = 1   # 카테고리 id는 1부터 할당
next_image_id = 1      # 이미지 id (각 annotation 파일마다 1개씩)
next_anno_id = 1       # annotation id (전체에 대해 순차적 할당)

# ANNOTATIONS_DIR 내의 모든 json 파일 목록
annotation_files = glob.glob(os.path.join(ANNOTATIONS_DIR, "*.json"))
print(f"총 {len(annotation_files)}개의 annotation 파일을 찾았습니다.")

for anno_file in annotation_files:
    try:
        with open(anno_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except Exception as e:
        print(f"파일 {anno_file} 로딩 중 오류 발생: {e}")
        continue  # 오류 발생 시 다음 파일로 넘어감

    # 이미지 정보 추출: data["item"] 에서 이미지 관련 정보가 있음.
    item = data.get("item", {})
    # 이미지 이름 (예: "00014.jpg")
    image_name = item.get("name", None)
    if image_name is None:
        print(f"{anno_file}에 이미지 이름 정보 누락")
        continue

    # 이미지 경로: 보통 slot 안에 있는 source_files에서 local_path 사용 (절대경로)
    # 여러 slot이 있을 수 있으므로 'image' 타입의 slot을 우선 탐색
    image_slot = None
    for slot in item.get("slots", []):
        if slot.get("type", "") == "image":
            image_slot = slot
            break
    if image_slot is None:
        print(f"{anno_file}에서 image slot을 찾을 수 없음")
        continue

    # 이미지의 width, height 값
    width = image_slot.get("width", None)
    height = image_slot.get("height", None)
    if width is None or height is None:
        print(f"{anno_file}에서 width/height 정보 누락")
        continue

    # source_files 내 첫번째 파일의 local_path 사용 (이미지의 절대 경로)
    source_files = image_slot.get("source_files", [])
    if len(source_files) == 0:
        print(f"{anno_file}의 source_files가 비어 있음")
        continue
    local_path = source_files[0].get("local_path", "")
    # (선택사항) 이미지 경로가 IMAGES_ROOT를 기준으로 상대 경로가 필요하다면 아래처럼 처리할 수 있음:
    # relative_path = os.path.relpath(local_path, IMAGES_ROOT)

    # COCO 이미지 정보 생성 (id, file_name, width, height 등)
    coco_img = {
        "id": next_image_id,
        "file_name": local_path,  # 또는 relative_path 사용 가능
        "width": width,
        "height": height
    }
    coco_images.append(coco_img)

    # 각 annotation 파일 내의 annotation 항목 처리
    for anno in data.get("annotations", []):
        # 카테고리 이름 (예: "Worker", "helmet", 등)
        cat_name = anno.get("name", "undefined")
        # 동적으로 category id 할당: 카테고리 매핑에 없으면 추가
        if cat_name not in category_mapping:
            category_mapping[cat_name] = next_category_id
            next_category_id += 1

        # bounding_box 정보 추출: COCO의 bbox는 [x, y, w, h]
        bb = anno.get("bounding_box", {})
        x = bb.get("x", 0)
        y = bb.get("y", 0)
        w = bb.get("w", 0)
        h = bb.get("h", 0)
        # area 계산 (w*h)
        area = w * h

        # COCO annotation 객체 생성
        coco_anno = {
            "id": next_anno_id,
            "image_id": next_image_id,  # 현재 이미지에 연결
            "category_id": category_mapping[cat_name],
            "bbox": [x, y, w, h],
            "area": area,
            "iscrowd": 0,
            "segmentation": []  # bbox만 있을 경우 빈 리스트로 처리
        }
        coco_annotations.append(coco_anno)
        next_anno_id += 1

    next_image_id += 1

print(f"총 {len(coco_images)}개의 이미지, {len(coco_annotations)}개의 annotation이 처리되었습니다.")

# %% [셀 3: COCO 포맷의 카테고리 목록 구성]
# category_mapping 딕셔너리를 기반으로 COCO 카테고리 목록 생성
coco_categories = []
for cat_name, cat_id in category_mapping.items():
    # supercategory는 단순히 카테고리 이름과 동일하게 지정하거나 필요에 따라 변경 가능
    coco_categories.append({
        "id": cat_id,
        "name": cat_name,
        "supercategory": cat_name
    })

print("발견된 카테고리 목록:")
for cat in coco_categories:
    print(cat)

# %% [셀 4: 최종 COCO JSON 생성 및 저장]
# COCO 포맷 최종 JSON 구조 구성
coco_output = {
    "images": coco_images,
    "annotations": coco_annotations,
    "categories": coco_categories
}

# 출력 파일이 저장될 폴더가 없다면 생성 (덮어쓰지 않도록 폴더명에 번호 추가 등도 고려 가능)
output_dir = os.path.dirname(OUTPUT_JSON_PATH)
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
    print(f"폴더 생성: {output_dir}")

# JSON 파일 저장 (indent=2로 예쁘게 저장)
with open(OUTPUT_JSON_PATH, 'w', encoding='utf-8') as f:
    json.dump(coco_output, f, ensure_ascii=False, indent=2)

print("COCO 형식 JSON 파일이 저장되었습니다:")
print(OUTPUT_JSON_PATH)