import os
import json
import cv2
from tqdm import tqdm  # 진행률 표시를 위한 라이브러리

# JSON 파일 경로와 이미지 폴더 경로 설정
# json_file_path = '/DATA2/ltb/mmyolo/it1_final_fusion/val.json'
# images_folder_path = '/DATA2/ltb/mmyolo/it1_final_fusion/val/images'
json_file_path = '/DATA2/ltb/mmyolo/it1_conequip_for_filtering/it1_중장비_v7_construction_equipment_8class.json'
images_folder_path = '/DATA2/ltb/mmyolo/it1_conequip_for_filtering/images'

# 라벨링된 이미지를 저장할 폴더 생성
output_folder = '/DATA2/ltb/mmyolo/it1_conequip_for_filtering/labeled_images_con_960'
os.makedirs(output_folder, exist_ok=True)

# JSON 파일 로드
with open(json_file_path, 'r') as f:
    data = json.load(f)

# 카테고리 ID와 이름을 매핑하는 딕셔너리 생성
category_id_to_name = {category['id']: category['name'] for category in data['categories']}

# 클래스별 색상 설정 (category_id에 맞게 수정)
class_colors = {
    1: (0, 255, 0),   # 클래스 ID 1은 초록색
    2: (255, 0, 0)    # 클래스 ID 2는 파란색
}

# 이미지 ID와 파일명을 매핑하는 딕셔너리 생성
image_id_to_file = {image['id']: image['file_name'] for image in data['images']}

# 이미지별 모든 바운딩 박스를 적용하여 저장
total_images = len(data['images'])  # 전체 이미지 개수
with tqdm(total=total_images, desc="Processing Images") as pbar:  # 진행률 표시
    for image in data['images']:
        image_id = image['id']
        image_file_name = image['file_name']
        image_path = os.path.join(images_folder_path, image_file_name)

        # 이미지 로드
        img = cv2.imread(image_path)
        if img is None:
            print(f"Failed to load image: {image_path}")
            pbar.update(1)
            continue

        # 해당 이미지의 모든 어노테이션 가져오기
        annotations = [ann for ann in data['annotations'] if ann['image_id'] == image_id]

        # 바운딩 박스 그리기
        for annotation in annotations:
            bbox = annotation['bbox']  # [x, y, width, height] 형식
            category_id = annotation['category_id']
            category_name = category_id_to_name.get(category_id, 'Unknown')
            color = class_colors.get(category_id, (0, 0, 255))  # 기본값은 빨간색

            x, y, w, h = map(int, bbox)
            cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)

            # 클래스 이름을 바운딩 박스 위에 표시
            label = category_name
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            label_x = x
            label_y = y - 10 if y - 10 > 10 else y + 10
            cv2.rectangle(img, (label_x, label_y - label_size[1]), (label_x + label_size[0], label_y + 5), color, -1)
            cv2.putText(img, label, (label_x, label_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        # 출력 경로 생성 및 이미지 저장
        output_path = os.path.join(output_folder, image_file_name.replace("/", "_"))
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        cv2.imwrite(output_path, img)

        # 진행률 업데이트
        pbar.update(1)

print("라벨링된 이미지를 저장했습니다.")