LabelClass

import os
from PIL import Image

def create_yolo_txt(image_path, class_id):
    # 이미지 크기 얻기
    with Image.open(image_path) as img:
        width, height = img.size
    
    # YOLO 포맷으로 변환 (전체 이미지를 바운딩박스로)
    # YOLO 포맷: <class> <x_center> <y_center> <width> <height>
    x_center = 0.5  # 중심점 x (정규화된 좌표)
    y_center = 0.5  # 중심점 y (정규화된 좌표)
    w = 1.0        # 너비 (정규화된 값)
    h = 1.0        # 높이 (정규화된 값)
    
    # txt 파일 경로 생성 (이미지와 동일한 이름, 확장자만 .txt로)
    txt_path = os.path.splitext(image_path)[0] + '.txt'
    
    # txt 파일 작성
    with open(txt_path, 'w') as f:
        f.write(f"{class_id} {x_center} {y_center} {w} {h}")

def process_dataset(base_dir):
    # 클래스 이름과 ID 매핑 (알파벳 순)
    class_mapping = {
        'ants': 0,
        'bees': 1
    }
    
    # 각 클래스 폴더 처리
    for class_name in class_mapping:
        class_dir = os.path.join(base_dir, class_name)
        if not os.path.exists(class_dir):
            continue
            
        class_id = class_mapping[class_name]
        
        # 폴더 내의 모든 이미지 파일 처리
        for filename in os.listdir(class_dir):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_path = os.path.join(class_dir, filename)
                create_yolo_txt(image_path, class_id)
                print(f"Created annotation for: {image_path}")

# 실행
dataset_path = "/Users/m1_mini/Desktop/antbee"  # 데이터셋 기본 경로
process_dataset(dataset_path)