YOLOv5로 객체 탐지 모델을 학습해야 하는 상황인데, 일부 데이터는 Segmentation 형식으로 되어 있어 바로 사용할 수 없었다. YOLO는 Segmentation을 직접 처리하지 못하기 때문에, 해당 데이터를 Bounding Box 형태로 변환하는 전처리 작업이 필요했다.
Polygon이란?
polygon은 다각형이라는 뜻으로, 컴퓨터 비전에서는 보통 물체의 외곽을 점(좌표)들의 집합으로 표현한 형태를 말한다.
다음과 같은 사진을 예시로 들어 다음과 같은 좌표들을 찍었다고 가정해보자.
좌표 예시: [(120, 80), (130, 75), (140, 78), (145, 90), (135, 100), (125, 98)]이 좌표들을 점대로 이으면 폐곡선이 된다. 이게 바로 Polygon annotation(=Segmentation)이다.
Segmentation 라벨 수 탐지
먼저 아래의 코드로 각각 폴더(train, test, valid)안 에 세그멘테이션 라벨의 수가 얼마인지부터 찍어보았다.
세그멘테이션 라벨 수가 0개라면 굳이 코드를 돌릴 필요가 없기 때문이다 !
import os
def count_label_types(label_dir):
bbox_count = 0
seg_count = 0
file_count = 0
for filename in os.listdir(label_dir):
if filename.endswith(".txt"):
file_path = os.path.join(label_dir, filename)
with open(file_path, "r") as f:
lines = f.readlines()
for line in lines:
items = line.strip().split()
if len(items) == 5:
bbox_count += 1
elif len(items) > 5:
seg_count += 1
file_count += 1
return bbox_count, seg_count, file_count
label_dir = "경로"
bbox, seg, files = count_label_types(label_dir)
print(f"🔍 총 라벨 파일 수: {files}개")
print(f"🟩 바운딩박스 라벨 수: {bbox}개")
print(f"🟦 세그멘테이션 라벨 수: {seg}개")Segmentation to Bounding Box
Segmentation 형식의 라벨 데이터를 YOLOv5에서 사용할 수 있도록 Bounding Box 형식으로 변환하는 코드이다.
YOLO는 polygon segmentation 데이터를 직접 처리할 수 없기 때문에, 해당 데이터를 x_center, y_center, width, height 형태의 바운딩 박스로 변환해주는 전처리 작업이 필요하다.
import os
import cv2
def convert_seg_to_yolo(label_dir, image_dir):
updated_count = 0
label_files = [f for f in os.listdir(label_dir) if f.endswith(".txt")]
for filename in label_files:
label_path = os.path.join(label_dir, filename)
base_name = os.path.splitext(filename)[0]
# 이미지 찾기
for ext in [".jpg", ".jpeg", ".png"]:
image_path = os.path.join(image_dir, base_name + ext)
if os.path.exists(image_path):
break
else:
print(f"🚫 이미지 없음: {base_name}")
continue
# 이미지 크기
img = cv2.imread(image_path)
if img is None:
print(f"🚫 이미지 로딩 실패: {image_path}")
continue
img_h, img_w = img.shape[:2]
with open(label_path, "r") as f:
lines = f.readlines()
new_lines = []
for line in lines:
items = list(map(float, line.strip().split()))
if len(items) <= 5:
new_lines.append(" ".join(map(str, items))) # bbox는 그대로 유지
continue
class_id = int(items[0])
coords = items[1:]
xs = coords[::2]
ys = coords[1::2]
# 정규화된 polygon이면 되돌리기
if max(xs) <= 1.0 and max(ys) <= 1.0:
xs = [x * img_w for x in xs]
ys = [y * img_h for y in ys]
x_min, x_max = min(xs), max(xs)
y_min, y_max = min(ys), max(ys)
# YOLO 형식으로 변환
x_center = (x_min + x_max) / 2.0 / img_w
y_center = (y_min + y_max) / 2.0 / img_h
w = (x_max - x_min) / img_w
h = (y_max - y_min) / img_h
new_line = f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}"
new_lines.append(new_line)
updated_count += 1
# 덮어쓰기
with open(label_path, "w") as f:
for l in new_lines:
f.write(l + "\n")
print(f"\n✅ 변환 완료: {updated_count}개의 세그멘테이션 객체가 bbox로 변환됨")
convert_seg_to_yolo("경로/labels", "경로/images")
이미지/라벨 불러오기
label_files = [f for f in os.listdir(label_dir) if f.endswith(".txt")].txt 파일만 골라서 처리하도록 해준다.
for ext in [".jpg", ".jpeg", ".png"]:
image_path = os.path.join(image_dir, base_name + ext)
if os.path.exists(image_path):
break
그리고 라벨 파일 이름과 매칭되는 이미지 파일을 탐색한다.
세그멘테이션 to 바운딩 박스
class_id = int(items[0])
coords = items[1:]
xs = coords[::2]
ys = coords[1::2]items의 길이가 6이상이면 Polygon으로 간주하고 좌표 쌍을 x, y로 분리한다.
if max(xs) <= 1.0 and max(ys) <= 1.0:
xs = [x * img_w for x in xs]
ys = [y * img_h for y in ys]그리고 좌표가 정규화된 값이라면 실제 이미지 크기로 되돌린다.
이때 정규화된 값임을 확인하기 위해서
img = cv2.imread(image_path)
img_h, img_w = img.shape[:2]다음과 같은 코드를 사용한다.
x_min, x_max = min(xs), max(xs)
y_min, y_max = min(ys), max(ys)
x_center = (x_min + x_max) / 2.0 / img_w
y_center = (y_min + y_max) / 2.0 / img_h
w = (x_max - x_min) / img_w
h = (y_max - y_min) / img_h가장 작은 사각형으로 감싸는 bounding box를 계산한다.
polygon을 이루는 x, y 좌표들 중
- 가장 작은 x, y: 좌상단 (top-left)
- 가장 큰 x, y: 우하단 (bottom-right)
이 두 점을 연결하면 해당 polygon을 감싸는 최소 크기의 사각형(bounding box)이 된다.
그 후 YOLO 포맷으로 변환시켜 준다.
확인
import os
import cv2
import matplotlib.pyplot as plt
base_path = '경로'
sets = ['train', 'test', 'valid']
def draw_yolo_bbox(image_path, label_path):
img = cv2.imread(image_path)
if img is None:
print(f"이미지를 읽을 수 없음: {image_path}")
return None
img_h, img_w = img.shape[:2]
with open(label_path, 'r') as f:
lines = f.readlines()
for line in lines:
parts = list(map(float, line.strip().split()))
if len(parts) != 5:
continue # 잘못된 형식 무시
class_id, x_center, y_center, w, h = parts
x_center *= img_w
y_center *= img_h
w *= img_w
h *= img_h
x1 = int(x_center - w / 2)
y1 = int(y_center - h / 2)
x2 = int(x_center + w / 2)
y2 = int(y_center + h / 2)
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(img, str(int(class_id)), (x1, y1 - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 1, cv2.LINE_AA)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
for dataset in sets:
label_dir = os.path.join(base_path, dataset, 'labels')
image_dir = os.path.join(base_path, dataset, 'images')
label_files = [f for f in os.listdir(label_dir) if f.endswith('.txt')]
print(f"\n📂 {dataset.upper()} 데이터셋 시각화 ({len(label_files)}장)")
for label_file in label_files:
base_name = os.path.splitext(label_file)[0]
label_path = os.path.join(label_dir, label_file)
for ext in ['.jpg', '.jpeg', '.png']:
image_path = os.path.join(image_dir, base_name + ext)
if os.path.exists(image_path):
break
else:
print(f"이미지 없음: {base_name}")
continue
img = draw_yolo_bbox(image_path, label_path)
if img is not None:
plt.figure(figsize=(6, 6))
plt.imshow(img)
plt.title(f"{dataset} / {label_file}")
plt.axis('off')
plt.show()
다음과 같은 코드를 통해 각각의 이미지가 Bounding Box로 잘 변환되었는지 확인해볼 수 있다.
공부 기록 공간 '◡'