import os import cv2 import random from collections import defaultdict # ============================================================================= # [유틸리티 함수] # ============================================================================= def ensure_dir(directory): """ 지정한 디렉토리 경로가 없으면 생성합니다. 이미 존재하면 그대로 사용합니다. """ if not os.path.exists(directory): os.makedirs(directory) print(f"디렉토리 생성: {directory}") else: print(f"디렉토리 이미 존재: {directory}") def parse_label_file(label_file_path): """ YOLO 형식의 라벨 파일을 파싱하여 annotation 리스트를 반환합니다. 파일 형식: 각 줄이 "class x_center y_center width height" 형태이며, 모든 값은 이미지 크기에 대해 정규화된 값(0~1)입니다. Parameters: label_file_path: 라벨 파일의 전체 경로 Returns: annotations: 리스트, 각 annotation은 (class, x_center, y_center, width, height) 튜플 """ annotations = [] try: with open(label_file_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if not line: continue tokens = line.split() cls = int(tokens[0]) x_center = float(tokens[1]) y_center = float(tokens[2]) width = float(tokens[3]) height = float(tokens[4]) annotations.append((cls, x_center, y_center, width, height)) except Exception as e: print(f"라벨 파일 파싱 오류: {label_file_path} / {e}") return annotations def collect_samples(image_dir, label_dir): """ 이미지 디렉토리와 라벨 디렉토리에서 파일들을 매칭하여 샘플 리스트를 생성합니다. 각 샘플은 파일 베이스 이름(확장자 제외)을 기준으로 매칭되며, 라벨 파일은 YOLO 형식의 txt 파일이고, 이미지 파일은 확장자 (.jpg, .png 등)로 매칭됩니다. Parameters: image_dir: 이미지 파일들이 위치한 디렉토리 (예: /path/to/images) label_dir: 라벨 txt 파일들이 위치한 디렉토리 (예: /path/to/labels) Returns: samples: 리스트, 각 샘플은 딕셔너리: { "base": 베이스 파일명(확장자 제외), "img_path": 이미지 파일 전체 경로, "label_path": 라벨 파일 전체 경로, "annotations": parse_label_file()를 통해 읽은 annotation 리스트, "classes": 해당 파일에 등장하는 클래스 집합 } """ samples = [] # 이미지 파일들의 베이스 이름과 전체 경로를 저장한 딕셔너리 생성 image_dict = {} for file in os.listdir(image_dir): if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')): base, _ = os.path.splitext(file) image_dict[base] = os.path.join(image_dir, file) # 라벨 파일들을 순회하며 이미지와 매칭 for file in os.listdir(label_dir): if file.lower().endswith(".txt"): base, _ = os.path.splitext(file) label_path = os.path.join(label_dir, file) if base not in image_dict: print(f"매칭되지 않는 라벨 파일 (이미지 없음): {label_path}") continue img_path = image_dict[base] annotations = parse_label_file(label_path) classes = set() for annot in annotations: classes.add(annot[0]) sample = { "base": base, "img_path": img_path, "label_path": label_path, "annotations": annotations, "classes": classes } samples.append(sample) return samples def select_test_samples(samples, desired_count): """ 샘플 리스트 중 모든 클래스가 최소 1회 등장하도록 선택한 후, 원하는 테스트 이미지 수(desired_count)만큼 샘플을 추가로 선정합니다. 방법: 1. 전체 샘플에서 등장하는 모든 클래스 집합을 계산합니다. 2. 최소 커버리지를 위해, 아직 포함되지 않은 클래스를 갖는 샘플들을 우선 선택합니다. 3. 선택된 샘플 수가 desired_count보다 작으면 나머지 샘플들 중에서 랜덤하게 추가 선택합니다. 4. 만약 desired_count가 전체 unique 클래스 수보다 작으면 자동으로 조정합니다. Parameters: samples: 샘플 리스트 desired_count: 사용자가 원하는 테스트 이미지 수 Returns: selected_samples: 선택된 샘플 리스트 """ # 전체 unique 클래스 집합 계산 all_classes = set() for sample in samples: all_classes.update(sample["classes"]) min_required = len(all_classes) if desired_count < min_required: print(f"지정한 테스트 이미지 수({desired_count})가 전체 클래스 수({min_required})보다 작습니다. 자동으로 {min_required}로 조정합니다.") desired_count = min_required # 최소 커버리지를 위한 선택 selected = [] covered = set() for sample in samples: new_classes = sample["classes"] - covered if new_classes: selected.append(sample) covered.update(sample["classes"]) if covered == all_classes: break # 추가로 랜덤하게 선택하여 desired_count에 도달 remaining = [s for s in samples if s not in selected] random.shuffle(remaining) while len(selected) < desired_count and remaining: selected.append(remaining.pop()) return selected def draw_annotations_on_image(image, annotations, class_names, color_palette=None, text_offsets=None, thickness=2): """ 주어진 이미지에 annotation 정보를 시각화하여 그립니다. - 각 annotation은 정규화된 좌표로 주어지며, 이미지 크기에 맞게 변환됩니다. - 클래스 번호에 따라 미리 정의된 색상 팔레트에서 색상을 선택하여 bbox와 텍스트에 적용합니다. - 텍스트는 bbox의 좌측 상단 위치에 표시되며, 각 클래스별로 오프셋(text_offsets)을 적용하여 겹치지 않도록 할 수 있습니다. Parameters: image: cv2 이미지 (numpy array) annotations: 리스트, 각 요소는 (class, x_center, y_center, width, height) 튜플 class_names: dictionary, 클래스 번호를 이름으로 매핑 (예: {0:"worker", 1:"helmet", ...}) color_palette: 리스트, 각 요소는 (B, G, R) 색상 튜플. None이면 기본 팔레트 사용. text_offsets: dictionary, 클래스 번호를 (x_offset, y_offset) 튜플로 매핑. None이면 기본 (0, -10) 사용. thickness: 선 두께 (기본값: 2) Returns: image: annotation이 그려진 이미지 (numpy array) """ # 기본 색상 팔레트 (BGR 형식). 충분한 개수를 미리 정의함. if color_palette is None: color_palette = [ (0, 0, 255), # Red (0, 255, 0), # Green (255, 0, 0), # Blue (0, 255, 255), # Yellow (255, 0, 255), # Magenta (255, 255, 0), # Cyan (128, 128, 0), # Olive (128, 0, 128), # Purple (0, 128, 128), # Teal (0, 0, 128) # Navy ] # 기본 텍스트 오프셋: 클래스별 지정 없으면 기본 (0, -10) if text_offsets is None: text_offsets = {} img_h, img_w = image.shape[:2] for annot in annotations: cls, x_center, y_center, width, height = annot # bbox 색상: 클래스 번호에 따라 색상 팔레트에서 선택 (클래스번호 % len(color_palette)) color = color_palette[cls % len(color_palette)] # 정규화 좌표를 실제 픽셀 좌표로 변환 abs_x_center = x_center * img_w abs_y_center = y_center * img_h abs_width = width * img_w abs_height = height * img_h # 좌상단, 우하단 좌표 계산 x1 = int(abs_x_center - abs_width / 2) y1 = int(abs_y_center - abs_height / 2) x2 = int(abs_x_center + abs_width / 2) y2 = int(abs_y_center + abs_height / 2) # 바운딩 박스 그리기 cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness) # 클래스 이름: class_names 매핑이 있다면 사용, 없으면 숫자 그대로 label = class_names.get(cls, str(cls)) # 클래스별 텍스트 오프셋 적용 (없으면 기본 (0, -10)) offset = text_offsets.get(cls, (0, -10)) text_position = (x1 + offset[0], y1 + offset[1]) cv2.putText(image, label, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, thickness, cv2.LINE_AA) return image def visualize_annotations(image_dir, label_dir, output_vis_dir, test_image_count=5, class_names=None, random_seed=42, color_palette=None, text_offsets=None): """ 이미지 디렉토리와 라벨 디렉토리에서 매칭된 샘플을 읽어, 지정된 수 만큼의 테스트 이미지에 bbox와 클래스 이름을 시각화하여 저장합니다. 주요 단계: 1. 이미지와 라벨 파일들을 매칭해 샘플 리스트를 생성합니다. 2. 각 라벨 파일 내 등장하는 전체 클래스 집합을 계산합니다. 3. 모든 클래스가 최소 1회 나타나도록 선택한 후, 사용자가 지정한 테스트 이미지 수(test_image_count)만큼 샘플을 선정합니다. (만약 지정 수가 전체 클래스 수보다 작으면 자동으로 조정됩니다.) 4. 선택된 각 샘플에 대해 이미지를 읽고, 각 annotation의 정규화 좌표를 실제 픽셀 좌표로 변환한 후 bbox와 클래스 이름(색상 팔레트 및 텍스트 오프셋 적용)을 그립니다. 5. 시각화된 이미지를 output_vis_dir에 저장합니다. Parameters: image_dir: 원본 이미지 디렉토리 경로 label_dir: YOLO 형식 라벨(txt) 파일 디렉토리 경로 output_vis_dir: 시각화 결과 이미지를 저장할 경로 (없으면 생성) test_image_count: 시각화할 테스트 이미지 수 (기본: 5) class_names: 클래스 번호를 이름으로 매핑하는 dictionary (예: {0:"worker", 1:"helmet", ...}) 지정하지 않으면 클래스 번호 문자열 그대로 사용 random_seed: 랜덤 선택을 위한 시드값 (재현성 확보) color_palette: 각 클래스의 bbox 및 텍스트 색상을 위한 (B, G, R) 튜플 리스트. None이면 기본 팔레트 사용. text_offsets: 클래스별 텍스트 위치 오프셋 (dictionary: 클래스번호 -> (x_offset, y_offset)). None이면 기본 (0, -10) """ # 출력 폴더 생성 ensure_dir(output_vis_dir) # 랜덤 시드 설정 random.seed(random_seed) # (1) 샘플 수집 samples = collect_samples(image_dir, label_dir) if not samples: print("선택된 샘플이 없습니다.") return # (2) 전체 unique 클래스 집합 계산 overall_classes = set() for sample in samples: overall_classes.update(sample["classes"]) print(f"전체 샘플에서 등장하는 클래스: {overall_classes}") # (3) 테스트 이미지 수 자동 조정: 지정한 수가 전체 클래스 수보다 작으면 조정 if test_image_count < len(overall_classes): print(f"지정한 테스트 이미지 수 ({test_image_count})가 전체 클래스 수 ({len(overall_classes)})보다 작아서 자동으로 조정합니다.") test_image_count = len(overall_classes) # (4) 최소 커버리지를 적용해 테스트용 샘플 선택 selected_samples = select_test_samples(samples, test_image_count) print(f"선택된 테스트 샘플 수: {len(selected_samples)}") # (5) 선택된 각 샘플에 대해 이미지 시각화 및 저장 for sample in selected_samples: img = cv2.imread(sample["img_path"]) if img is None: print(f"이미지 로드 실패: {sample['img_path']}") continue # color_palette 및 text_offsets가 지정되어 있으면 draw_annotations_on_image에 전달 vis_img = draw_annotations_on_image( img.copy(), sample["annotations"], class_names if class_names is not None else {}, color_palette=color_palette, text_offsets=text_offsets, thickness=2 ) vis_file_name = "vis_" + os.path.basename(sample["img_path"]) output_path = os.path.join(output_vis_dir, vis_file_name) cv2.imwrite(output_path, vis_img) print(f"시각화 이미지 저장: {output_path}") print("시각화 테스트 이미지 생성 완료.") # ============================================================================= # [스크립트 실행] # ============================================================================= if __name__ == "__main__": # ------------------------------------------------------------------------- # 1. 인풋 디렉토리 설정 # ------------------------------------------------------------------------- # 이미지 파일 디렉토리 (예: YOLO 변환된 이미지들이 있는 폴더) image_directory = "/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/robo_head_only_1561_250117/images" # 라벨 txt 파일 디렉토리 (예: YOLO 형식 라벨 텍스트 파일들이 있는 폴더) label_directory = "/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/robo_head_only_1561_250117/labels" # ------------------------------------------------------------------------- # 2. 아웃풋 디렉토리 설정 (시각화 결과를 저장할 경로) # ------------------------------------------------------------------------- output_visualization_dir = "/ltb/media/ltb/90887173887158A4/Users/ltb/박스_데이터셋/박스_yolo11/robo_head_only_1561_250117/시각화테스트" # ------------------------------------------------------------------------- # 3. 테스트 이미지 수 설정 (시각화할 이미지 파일 수) # ------------------------------------------------------------------------- test_img_count = 10 # 원하는 이미지 수로 조정 # ------------------------------------------------------------------------- # 4. 클래스 이름 수정 및 매핑 (필요시 사용자 정의; 예시) # ------------------------------------------------------------------------- custom_class_names = {0: "head", # 1: "helmet", 2: "head", 3: "background" } # ------------------------------------------------------------------------- # 5. 색상 팔레트 및 텍스트 오프셋 설정 (필요시 사용자 정의) # ------------------------------------------------------------------------- # color_palette: BGR 튜플 리스트 (각 클래스에 순차적 적용) custom_color_palette = [ (0, 0, 255), # Red (0, 255, 0), # Green (255, 0, 0), # Blue (0, 255, 255), # Yellow (255, 0, 255), # Magenta (255, 255, 0), # Cyan (128, 128, 0), # Olive (128, 0, 128), # Purple (0, 128, 128), # Teal (0, 0, 128) # Navy ] # text_offsets: 클래스별 텍스트 위치 오프셋, (x, y) 값 (좌측 상단 기준) # 예시: {클래스 번호: (x_offset, y_offset)} custom_text_offsets = { 0: (0, 0), 1: (0, -10), 2: (0, -10), 3: (0, 0) } # ------------------------------------------------------------------------- # 6. 함수 호출: 이미지와 라벨 시각화 및 테스트 이미지 생성 # ------------------------------------------------------------------------- visualize_annotations( image_directory, label_directory, output_visualization_dir, test_image_count=test_img_count, class_names=custom_class_names, random_seed=42, color_palette=custom_color_palette, text_offsets=custom_text_offsets )