최종 학습 및 평가 (3) : Segformer

조훈·2025년 6월 21일

1. Segformer 학습 mask

Labeltxt_2_Mask.py

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)
  • Label.txt에 라벨링 된 Class_Index가 Mask.png의 픽셀값이 됨
    -> Driving Area는 1의 픽셀값, Parking Area는 2의 픽셀값을 가짐 (나머지는 Background = 0 )


(자세히 봐야 알 수 있음...)


2. Hyperparameter

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(),
)
  • 학습시간이 비교적 길기 때문에 save_strategy를 steps(100)로 설정
  • 각 Epoch 학습이 종료될 때 evaluation

3. 결과

0개의 댓글