이제 IRFS구현을 위한 코드를 짜준후에 모델을 훈련시킵니다. 가벼운 coco128 데이터셋을 활용하여 실험을 진행하였습니다. (128개의 이미지, 941개의 바운딩박스, 79개의 클래스)
coco128 클래스 분포
논문의 수식을 다시 한번 살펴보고 가겠습니다. 적은 클래스가 포함된 이미지일수록 image-level repeat factor가 커지게 되고, 해당이미지가 훈련에 적용될 수 있는 기회를 더 많이 부여하는 작업이 필요합니다.
Weighted Random Sampler로 각 이미지들이 뽑힐 확률을 조정해주는 방법으로도 IRFS를 구현할 수 있지만, image-level repeat factor만큼 이미지를 복사한 데이터셋을 만들어준다면, 기존 모델의 코드를 전혀 바꾸지 않고도 바뀐 데이터셋을 활용하여 샘플링 효과가 적용된 훈련을 할 수 있습니다.
import os
import glob
from collections import defaultdict
import numpy as np
import shutil
import math
# 클래스별 바운딩 박스 개수 및 총 바운딩 박스 개수 계산 함수
def count_labels(directory):
class_counts = defaultdict(int)
total_bboxes = 0
images_with_class = defaultdict(int)
total_images = 0
image_class_map = {}
# 지정된 디렉토리 내 모든 .txt 파일 찾기
for filename in glob.glob(os.path.join(directory, '*.txt')):
total_images += 1
classes_in_image = set() # 각 이미지에 대해 새로운 빈 집합 생성
with open(filename, 'r') as file:
for line in file:
if line.strip(): # 공백 줄이 아닌 경우에만 처리
class_id = int(line.split()[0])
class_counts[class_id] += 1
classes_in_image.add(class_id) # 클래스 ID를 집합에 추가
total_bboxes += 1
for class_id in classes_in_image:
images_with_class[class_id] += 1
# 이미지 파일 경로와 해당 이미지에 포함된 클래스 매핑 저장
image_class_map[filename] = list(classes_in_image)
return class_counts, total_bboxes, images_with_class, total_images, image_class_map
# 라벨 파일들이 위치한 디렉토리 경로 설정
labels_dir = '/content/drive/MyDrive/Colab/ultralytics/ultralytics/cfg/datasets/coco128/labels/train2017'
images_dir = '/content/drive/MyDrive/Colab/ultralytics/ultralytics/cfg/datasets/coco128/images/train2017'
# 클래스별 바운딩 박스 개수 및 총 바운딩 박스 개수 계산
class_counts, total_bboxes, images_with_class, total_images, image_class_map = count_labels(labels_dir)
# f_(i,c)와 f_(b,c) 계산 함수
def calculate_f_ic_and_f_bc(class_counts, images_with_class, total_images, total_bboxes):
f_ic = {class_id: images_with_class[class_id] / total_images for class_id in class_counts}
f_bc = {class_id: class_counts[class_id] / total_bboxes for class_id in class_counts}
return f_ic, f_bc
# f_(i,c)와 f_(b,c) 계산
f_ic, f_bc = calculate_f_ic_and_f_bc(class_counts, images_with_class, total_images, total_bboxes)
# r_c 계산 함수
def calculate_r_c(f_ic, f_bc, t):
r_c_values = {}
for class_id in f_ic:
value = np.sqrt(t / np.sqrt(f_ic[class_id] * f_bc[class_id]))
r_c = max(1, value)
r_c_values[class_id] = r_c
return r_c_values
# 예시로 사용할 t 값
t = 0.5
# r_c 계산
r_c_values = calculate_r_c(f_ic, f_bc, t)
# r_i 계산 함수
def calculate_r_i(image_class_map, r_c_values):
r_i_values = {}
for image_path, class_ids in image_class_map.items():
max_r_c = max(r_c_values[class_id] for class_id in class_ids)
r_i_values[image_path] = max_r_c
return r_i_values
# r_i 계산
r_i_values = calculate_r_i(image_class_map, r_c_values)
# 새 데이터셋 생성 함수
def create_augmented_dataset(r_i_values, images_dir, labels_dir, output_dir):
if not os.path.exists(output_dir):
os.makedirs(os.path.join(output_dir, 'images/train2017'))
os.makedirs(os.path.join(output_dir, 'labels/train2017'))
for image_path, r_i in r_i_values.items():
ceil_r_i = math.ceil(r_i)
base_image_name = os.path.basename(image_path).replace('.txt', '.jpg')
image_file = os.path.join(images_dir, base_image_name)
label_file = image_path
if not os.path.exists(image_file):
print(f"Warning: Image file {image_file} does not exist. Skipping...")
continue
for i in range(ceil_r_i):
new_image_name = f"{os.path.splitext(base_image_name)[0]}_{i}.jpg"
new_label_name = f"{os.path.splitext(os.path.basename(label_file))[0]}_{i}.txt"
new_image_path = os.path.join(output_dir, 'images/train2017', new_image_name)
new_label_path = os.path.join(output_dir, 'labels/train2017', new_label_name)
# 이미지 파일 복사
shutil.copy(image_file, new_image_path)
# 라벨 파일 복사
shutil.copy(label_file, new_label_path)
# 출력 디렉토리 설정
output_dir = '/content/drive/MyDrive/Colab/ultralytics/ultralytics/cfg/datasets/coco128_IRFS_05'
# 새 데이터셋 생성
create_augmented_dataset(r_i_values, images_dir, labels_dir, output_dir)
print("Augmented dataset created successfully.")
위의 코드로 각 이미지들의 image-level repeat factor들을 구할 수 있고, 그 수만큼의 이미지를 복제합니다.(해당 코드에서는 소수점의 repeat factor는 올림처리하였다)
t=0.5로 진행하였습니다. 논문에서의 최적값은 0.001이었으나 coco128은 LVIS 데이터셋보다는 훨씬 크기가 작고 클래스 수도 적기 때문에 f_ic와 f_bc 가 훨씬 크게 산출됩니다.(LVIS는 1202개의 클래스를 가진 데이터셋)
따라서 t=0.001로 설정하고 만든 데이터셋에는 복제된 사진의 양이 거의 없었고 t를 늘릴수록 복제되는 데이터량이 많은것을 확인 후에 가장 적절하다고 생각되는 t=0.5로 채택하여 진행하였습니다.
이제 복제한 데이터셋을 활용할 수 있게 하기 위한 yaml 파일 같은 디렉토리에 생성해줍니다.
기존 coco128 의 yaml 파일에서 path 만 .../coco128_IRFS_05 (새로 생성한 데이터셋) 로 변경하여 생성해주면 됩니다. (yaml 파일명 > coco128_IRFS_05.yaml)
yaml 파일 생성까지 마쳤다면 실행코드를 돌려주면 됩니다!
from ultralytics import YOLO
# pretrained model
model = YOLO('yolov8n.pt')
# 사용자 지정 데이터셋 구성 파일 경로
data = '/content/drive/MyDrive/Colab/ultralytics/ultralytics/cfg/datasets/coco128_IRFS_05.yaml'
img_size = 640
epoch_num = 50
batch_size = 8
task = 'detect'
mode = 'train'
# 훈련 결과를 저장할 폴더의 절대 경로
save_dir = '/content/drive/MyDrive/Colab/training_result/coco128_IRFS_05'
#모델 훈련
results = model.train(data=data,
imgsz=img_size,
epochs=epoch_num,
batch=batch_size,
task=task,
mode=mode,
project=save_dir)
save_dir 을 설정해주어 훈련결과를 확인할 폴더를 지정해주면 훈련이 끝난 후에 결과파일들(성능지표 등)을 받아볼 수 있습니다.
