import json
import os
from collections import defaultdict

def convert_coco_to_yolo(root_dir):
    """
    COCO 형식의 JSON 파일을 YOLO 형식의 txt 파일로 변환하고,
    원본 이미지와 동일한 디렉토리 구조로 라벨 파일을 생성합니다.

    Args:
        root_dir (str): 데이터셋의 루트 디렉토리 경로.
                        (예: 'ForestPersons_v3')
    """
    annotations_path = os.path.join(root_dir, 'annotations.json')
    images_dir = os.path.join(root_dir, 'images')
    labels_dir = os.path.join(root_dir, 'labels')

    # 1. labels 디렉토리 생성
    print(f"'{labels_dir}' 디렉토리를 생성합니다.")
    os.makedirs(labels_dir, exist_ok=True)

    # 2. JSON 파일 로드
    print(f"'{annotations_path}' 파일을 로드합니다.")
    with open(annotations_path, 'r') as f:
        coco_data = json.load(f)

    # 3. 데이터 처리를 용이하게 하기 위해 dictionary로 변환
    images_info = {img['id']: img for img in coco_data['images']}
    
    # image_id를 키로 하여 해당 이미지의 모든 어노테이션을 리스트로 저장
    annotations_by_image_id = defaultdict(list)
    for ann in coco_data['annotations']:
        annotations_by_image_id[ann['image_id']].append(ann)

    # 4. 각 이미지에 대해 YOLO 형식의 라벨 파일 생성
    total_images = len(images_info)
    processed_count = 0
    print(f"총 {total_images}개의 이미지에 대한 라벨 파일 생성을 시작합니다.")

    for image_id, image in images_info.items():
        # 파일 경로 및 이름 정보 추출
        file_name = image['file_name']
        image_width = image['width']
        image_height = image['height']

        # 출력 파일 경로 설정
        # 예: '001.../abc.jpg' -> '001.../abc.txt'
        label_relative_path = os.path.splitext(file_name)[0] + '.txt'
        label_abs_path = os.path.join(labels_dir, label_relative_path)

        # 라벨 파일을 저장할 하위 디렉토리 생성
        os.makedirs(os.path.dirname(label_abs_path), exist_ok=True)
        
        yolo_annotations = []
        
        # 해당 이미지에 대한 모든 어노테이션을 순회
        if image_id in annotations_by_image_id:
            for ann in annotations_by_image_id[image_id]:
                # COCO bbox: [x_min, y_min, width, height]
                bbox = ann['bbox']
                category_id = ann['category_id']

                x_min, y_min, w, h = bbox
                
                # YOLO 형식으로 변환 (정규화)
                # YOLO category_id는 0부터 시작하는 경우가 많으므로, 원본 id에서 1을 빼주는 것이 일반적입니다.
                # 만약 카테고리 ID가 0부터 시작한다면 이 부분을 수정해야 합니다.
                yolo_class_id = category_id - 1
                
                x_center = (x_min + w / 2) / image_width
                y_center = (y_min + h / 2) / image_height
                width_norm = w / image_width
                height_norm = h / image_height
                
                # 소수점 6자리까지 표현
                yolo_line = f"{yolo_class_id} {x_center:.6f} {y_center:.6f} {width_norm:.6f} {height_norm:.6f}"
                yolo_annotations.append(yolo_line)

        # 변환된 어노테이션을 .txt 파일에 저장
        with open(label_abs_path, 'w') as f:
            f.write('\n'.join(yolo_annotations))
        
        processed_count += 1
        if processed_count % 100 == 0: # 100개마다 진행 상황 출력
            print(f"진행 상황: {processed_count}/{total_images}")

    print("-" * 30)
    print(f"✅ 변환 완료! 총 {processed_count}개의 라벨 파일이 '{labels_dir}'에 생성되었습니다.")


if __name__ == '__main__':
    # 데이터셋의 루트 디렉토리 경로를 지정해주세요.
    # 이 스크립트가 'ForestPersons_v3' 폴더와 같은 위치에 있다면 그대로 사용하면 됩니다.
    # ROOT_DIR = 'ForestPersons_v3' 
    # ROOT_DIR = 'ForestPersons_v3_0924' 
    ROOT_DIR = 'data_forestpersons_v3' 
    
    if not os.path.exists(ROOT_DIR):
        print(f"오류: '{ROOT_DIR}' 디렉토리를 찾을 수 없습니다. 경로를 확인해주세요.")
    else:
        convert_coco_to_yolo(ROOT_DIR)