import os
import cv2
import numpy as np
from tqdm import tqdm
# 경로 설정
img_dir = "YOUR_IMG_DIR" # 원본 이미지 폴더 (크기 참고용)
label_dir = "YOUR_LABEL_DIR" # YOLO 형식 label 폴더
mask_dir = "YOUR_MASK_DIR" # 출력 마스크 저장 폴더
os.makedirs(mask_dir, exist_ok=True)
for label_file in tqdm(os.listdir(label_dir)):
if not label_file.endswith(".txt"):
continue
# 이미지와 마스크 크기 추출
img_name = os.path.splitext(label_file)[0] + ".png"
img_path = os.path.join(img_dir, img_name)
image = cv2.imread(img_path)
if image is None:
print(f"이미지 없음: {img_path}")
continue
h, w = image.shape[:2]
mask = np.zeros((h, w), dtype=np.uint8)
# 레이블 파싱
with open(os.path.join(label_dir, label_file), "r") as f:
for line in f:
parts = line.strip().split()
cls = int(parts[0]) +1
points = np.array(parts[1:], dtype=np.float32).reshape(-1, 2)
points *= [w, h] # 정규화된 좌표를 원래 크기로
points = points.astype(np.int32)
cv2.fillPoly(mask, [points], color=cls)
# 마스크 저장
mask_path = os.path.join(mask_dir, os.path.splitext(label_file)[0] + ".png")
cv2.imwrite(mask_path, mask)

(자세히 봐야 알 수 있음...)
training_args = TrainingArguments(
output_dir=output_dir,
per_device_train_batch_size=16,
num_train_epochs=10,
logging_dir=logging_dir,
save_strategy="steps",
save_steps=100,
logging_steps=50,
evaluation_strategy="epoch",
learning_rate=5e-5,
save_total_limit=3,
fp16=torch.cuda.is_available(),
)
