본문 바로가기
AI/YOLO

CycleGAN을 이용한 데이터 증강

by wenect 2024. 7. 8.

공개된 데이터셋 및 기술

자동차 표면 결함 감지를 위한 AI 데이터 증강과 관련하여, 다음과 같은 공개 데이터셋과 기술을 활용할 수 있습니다.

1. CarDD (Car Damage Detection)

CarDD는 자동차 외부 손상을 감지하고 평가하기 위한 고해상도 데이터셋입니다. 이 데이터셋은 4,000개의 이미지와 9,000개 이상의 손상 인스턴스로 구성되어 있으며, 덴트(dent), 긁힘(scratch), 균열(crack), 유리 파손(glass shatter), 램프 손상(lamp broken), 타이어 펑크(tire flat) 등 다양한 손상 유형을 포함합니다. 이 데이터셋은 분류(classification), 객체 탐지(object detection), 인스턴스 분할(instance segmentation) 및 주목 객체 탐지(salient object detection) 작업에 사용될 수 있습니다​ (ar5iv)​.

2. KolektorSDD

Kolektor Surface-Defect Dataset(KolektorSDD)은 전기 코뮨테이터의 표면 결함을 포함하는 데이터셋입니다. 이 데이터셋은 결함이 있는 52개의 이미지와 결함이 없는 347개의 이미지로 구성되어 있으며, 결함 위치와 유형이 명확히 라벨링되어 있습니다. 이 데이터셋은 결함 감지 및 분류 작업에 유용하게 사용될 수 있습니다​ (GitHub)​.

3. 공개된 다른 데이터셋

다양한 표면 결함 감지를 위한 다른 데이터셋들도 공개되어 있습니다. 예를 들어, NEU-CLS는 강철 표면의 전형적인 결함을 포함하는 데이터셋으로, 여러 결함 유형에 대한 위치와 범주를 제공하는 라벨을 포함합니다. 또한 Kaggle의 Severstal: Steel Defect Detection 데이터셋은 철강 제조 과정에서 발생하는 다양한 결함을 포함하고 있습니다​ (GitHub)​.

CycleGAN을 이용한 데이터 증강

CycleGAN은 두 도메인 간의 이미지를 변환하는 모델로, 결함이 없는 자동차 표면 이미지를 결함이 있는 이미지로 변환하는 데 사용할 수 있습니다. CycleGAN을 통해 생성된 결함 이미지는 실제 결함 이미지를 보강하여 학습 데이터셋을 확장하고, 모델의 성능을 향상시킬 수 있습니다.

결론

공개된 데이터셋과 CycleGAN과 같은 생성 모델을 활용하여 자동차 표면 결함 감지를 위한 고품질 데이터를 증강할 수 있습니다. CarDD와 KolektorSDD와 같은 데이터셋을 사용하여 초기 모델을 학습시키고, CycleGAN을 통해 생성된 데이터를 추가하여 모델의 일반화 성능을 높일 수 있습니다. 이러한 접근 방식은 다양한 결함 사례를 포함한 풍부한 데이터셋을 확보하는 데 도움이 될 것입니다.


데이터 증강을 위해 생성 모델을 사용하는 방법은 매우 효과적입니다. 특히, 자동차 표면의 결함(찍힘, 긁힘 등)을 생성하기 위해서는 다양한 생성 모델을 사용할 수 있습니다. 여기서는 주로 GAN(Generative Adversarial Networks)과 같은 기술을 활용한 방법을 설명하겠습니다.

1. GAN (Generative Adversarial Networks)

GAN은 두 개의 신경망(생성자와 판별자)이 서로 경쟁하며 학습하는 모델입니다. 생성자는 새로운 이미지를 생성하고, 판별자는 이 이미지가 진짜인지 가짜인지 구별합니다. 이 과정을 통해 생성자는 점점 더 현실적인 이미지를 생성하게 됩니다.

2. 결함 이미지 생성 방법

결함 이미지 생성에 사용할 수 있는 구체적인 GAN 모델로는 DCGAN, CycleGAN, StyleGAN 등이 있습니다. 이 모델들은 특정 결함(찍힘, 긁힘 등)을 생성하는 데 매우 유용합니다.

3. CycleGAN을 이용한 결함 생성

CycleGAN은 두 도메인 간의 이미지를 서로 변환하는 모델입니다. 예를 들어, 결함이 없는 자동차 표면 이미지를 결함이 있는 이미지로 변환할 수 있습니다.

CycleGAN을 이용한 결함 이미지 생성 예시

먼저 CycleGAN을 구현하려면 다음과 같은 단계를 따릅니다.

  1. 데이터 준비:
    • 결함이 없는 자동차 표면 이미지와 결함이 있는 자동차 표면 이미지 데이터셋을 준비합니다.
  2. 모델 훈련:
    • CycleGAN 모델을 훈련시켜 결함이 없는 이미지를 결함이 있는 이미지로 변환할 수 있도록 합니다.
  3. 이미지 생성:
    • 훈련된 CycleGAN 모델을 사용하여 새로운 결함 이미지를 생성합니다.

CycleGAN 구현 예시 (Python, TensorFlow 사용)

import tensorflow as tf
from tensorflow_examples.models.pix2pix import pix2pix
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt

# 데이터 로드
dataset, metadata = tfds.load('cycle_gan/monet2photo', with_info=True, as_supervised=True)
train_monet, train_photo = dataset['trainA'], dataset['trainB']

# 이미지를 전처리하는 함수
def preprocess_image(image, label):
    image = tf.image.resize(image, [256, 256])
    image = (image / 127.5) - 1
    return image

train_monet = train_monet.map(preprocess_image)
train_photo = train_photo.map(preprocess_image)

# CycleGAN 모델 정의
generator_g = pix2pix.unet_generator(output_channels=3, norm_type='instancenorm')
generator_f = pix2pix.unet_generator(output_channels=3, norm_type='instancenorm')
discriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False)
discriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False)

# 모델 컴파일
generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

# CycleGAN 훈련 함수 정의
@tf.function
def train_step(real_x, real_y):
    with tf.GradientTape(persistent=True) as tape:
        fake_y = generator_g(real_x, training=True)
        cycled_x = generator_f(fake_y, training=True)

        fake_x = generator_f(real_y, training=True)
        cycled_y = generator_g(fake_x, training=True)

        same_x = generator_f(real_x, training=True)
        same_y = generator_g(real_y, training=True)

        disc_real_x = discriminator_x(real_x, training=True)
        disc_real_y = discriminator_y(real_y, training=True)

        disc_fake_x = discriminator_x(fake_x, training=True)
        disc_fake_y = discriminator_y(fake_y, training=True)

        gen_g_loss = generator_loss(disc_fake_y)
        gen_f_loss = generator_loss(disc_fake_x)

        total_cycle_loss = cycle_loss(real_x, cycled_x) + cycle_loss(real_y, cycled_y)

        total_gen_g_loss = gen_g_loss + total_cycle_loss + identity_loss(real_y, same_y)
        total_gen_f_loss = gen_f_loss + total_cycle_loss + identity_loss(real_x, same_x)

        disc_x_loss = discriminator_loss(disc_real_x, disc_fake_x)
        disc_y_loss = discriminator_loss(disc_real_y, disc_fake_y)

    generator_g_gradients = tape.gradient(total_gen_g_loss, generator_g.trainable_variables)
    generator_f_gradients = tape.gradient(total_gen_f_loss, generator_f.trainable_variables)

    discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables)
    discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables)

    generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables))
    generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables))

    discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables))
    discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables))

# CycleGAN 학습
EPOCHS = 40
for epoch in range(EPOCHS):
    for image_x, image_y in tf.data.Dataset.zip((train_photo, train_monet)):
        train_step(image_x, image_y)

# 이미지 변환 및 저장
def generate_images(model, test_input):
    prediction = model(test_input)
    plt.figure(figsize=(12, 12))

    display_list = [test_input[0], prediction[0]]
    title = ['Input Image', 'Generated Image']

    for i in range(2):
        plt.subplot(1, 2, i+1)
        plt.title(title[i])
        plt.imshow((display_list[i] * 0.5 + 0.5))
        plt.axis('off')
    plt.show()

sample_photo = next(iter(train_photo.take(1)))
generate_images(generator_g, sample_photo)

위 코드는 TensorFlow와 CycleGAN을 사용하여 결함 없는 이미지를 결함이 있는 이미지로 변환하는 예시입니다. 이를 통해 현실적인 결함 이미지를 생성할 수 있습니다.

4. 모델 평가 및 검증

생성된 이미지를 실제 데이터와 비교하여 모델의 성능을 평가하고, 필요에 따라 모델을 조정합니다.

5. 데이터 증강

생성된 결함 이미지를 기존의 결함 데이터와 함께 사용하여 데이터셋을 확장하고, 모델 학습에 활용합니다.

이와 같은 방식으로 생성 모델을 활용하여 자동차 표면의 결함 이미지를 증강할 수 있습니다. 이를 통해 다양한 결함 사례를 포함한 데이터를 확보하여 모델의 성능을 향상시킬 수 있습니다.


CycleGAN 학습 과정

CycleGAN은 두 가지 이미지를 서로 변환하는 방법을 학습합니다. 예를 들어, 깨끗한 자동차 이미지와 긁힌 자동차 이미지를 학습한다고 생각해 보세요.

  1. 깨끗한 자동차 이미지를 CycleGAN에 입력합니다.
  2. CycleGAN 생성자는 이를 긁힌 자동차 이미지로 변환합니다.
  3. 변환된 이미지를 다시 깨끗한 자동차 이미지로 되돌립니다.
  4. 원래 이미지와 되돌린 이미지의 차이를 확인하며 학습합니다.

이 과정에서 생성자는 점점 더 진짜 같은 긁힌 이미지를 만들고, 판별자는 그 이미지를 진짜와 가짜로 구별하는 능력이 향상됩니다.

간단한 예시 코드

이제 CycleGAN을 실제로 사용하여 이미지를 변환하는 간단한 예시를 보여드릴게요.

import tensorflow as tf
from tensorflow_examples.models.pix2pix import pix2pix

# 데이터 준비
clean_car_images = ... # 깨끗한 자동차 이미지 데이터셋
scratched_car_images = ... # 긁힌 자동차 이미지 데이터셋

# 모델 준비
generator_g = pix2pix.unet_generator(output_channels=3)
generator_f = pix2pix.unet_generator(output_channels=3)
discriminator_x = pix2pix.discriminator()
discriminator_y = pix2pix.discriminator()

# 손실 함수 및 최적화 도구 설정
loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)
generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

# 학습 함수 정의
@tf.function
def train_step(real_x, real_y):
    with tf.GradientTape(persistent=True) as tape:
        fake_y = generator_g(real_x, training=True)
        cycled_x = generator_f(fake_y, training=True)
        
        fake_x = generator_f(real_y, training=True)
        cycled_y = generator_g(fake_x, training=True)

        disc_real_x = discriminator_x(real_x, training=True)
        disc_real_y = discriminator_y(real_y, training=True)

        disc_fake_x = discriminator_x(fake_x, training=True)
        disc_fake_y = discriminator_y(fake_y, training=True)

        gen_g_loss = loss_obj(disc_fake_y, tf.ones_like(disc_fake_y))
        gen_f_loss = loss_obj(disc_fake_x, tf.ones_like(disc_fake_x))
        
        disc_x_loss = loss_obj(disc_real_x, tf.ones_like(disc_real_x)) + loss_obj(disc_fake_x, tf.zeros_like(disc_fake_x))
        disc_y_loss = loss_obj(disc_real_y, tf.ones_like(disc_real_y)) + loss_obj(disc_fake_y, tf.zeros_like(disc_fake_y))

    generator_g_gradients = tape.gradient(gen_g_loss, generator_g.trainable_variables)
    generator_f_gradients = tape.gradient(gen_f_loss, generator_f.trainable_variables)

    discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables)
    discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables)

    generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables))
    generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables))

    discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables))
    discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables))

# 학습 실행
EPOCHS = 50
for epoch in range(EPOCHS):
    for image_x, image_y in zip(clean_car_images, scratched_car_images):
        train_step(image_x, image_y)

# 변환된 이미지 생성 및 확인
sample_image = clean_car_images[0]
fake_image = generator_g(sample_image)

CycleGAN을 사용하여 깨끗한 자동차 이미지를 긁힌 이미지로 변환하는 과정을 보여줍니다. CycleGAN은 이런 방식으로 이미지 데이터를 증강하여 결함 감지 모델의 성능을 향상시키는 데 사용할 수 있습니다.

샘플2

import tensorflow as tf
from tensorflow_examples.models.pix2pix import pix2pix
import tensorflow_datasets as tfds

# 데이터 로드
dataset, metadata = tfds.load('cycle_gan/horse2zebra', with_info=True, as_supervised=True)
train_horses, train_zebras = dataset['trainA'], dataset['trainB']

# 이미지 전처리 함수
def preprocess_image(image, label):
    image = tf.image.resize(image, [256, 256])
    image = (image / 127.5) - 1
    return image

train_horses = train_horses.map(preprocess_image)
train_zebras = train_zebras.map(preprocess_image)

# CycleGAN 모델 정의
generator_g = pix2pix.unet_generator(output_channels=3, norm_type='instancenorm')
generator_f = pix2pix.unet_generator(output_channels=3, norm_type='instancenorm')
discriminator_x = pix2pix.discriminator(norm_type='instancenorm', target=False)
discriminator_y = pix2pix.discriminator(norm_type='instancenorm', target=False)

# 모델 컴파일 및 학습
loss_obj = tf.keras.losses.BinaryCrossentropy(from_logits=True)
generator_g_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
generator_f_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_x_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_y_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

# 학습 함수 정의
@tf.function
def train_step(real_x, real_y):
    with tf.GradientTape(persistent=True) as tape:
        fake_y = generator_g(real_x, training=True)
        cycled_x = generator_f(fake_y, training=True)
        
        fake_x = generator_f(real_y, training=True)
        cycled_y = generator_g(fake_x, training=True)

        same_x = generator_f(real_x, training=True)
        same_y = generator_g(real_y, training=True)

        disc_real_x = discriminator_x(real_x, training=True)
        disc_real_y = discriminator_y(real_y, training=True)

        disc_fake_x = discriminator_x(fake_x, training=True)
        disc_fake_y = discriminator_y(fake_y, training=True)

        gen_g_loss = loss_obj(disc_fake_y, tf.ones_like(disc_fake_y))
        gen_f_loss = loss_obj(disc_fake_x, tf.ones_like(disc_fake_x))
        
        cycle_loss = tf.reduce_mean(tf.abs(real_x - cycled_x)) + tf.reduce_mean(tf.abs(real_y - cycled_y))
        identity_loss = tf.reduce_mean(tf.abs(real_x - same_x)) + tf.reduce_mean(tf.abs(real_y - same_y))

        total_gen_g_loss = gen_g_loss + cycle_loss + identity_loss
        total_gen_f_loss = gen_f_loss + cycle_loss + identity_loss

        disc_x_loss = loss_obj(disc_real_x, tf.ones_like(disc_real_x)) + loss_obj(disc_fake_x, tf.zeros_like(disc_fake_x))
        disc_y_loss = loss_obj(disc_real_y, tf.ones_like(disc_real_y)) + loss_obj(disc_fake_y, tf.zeros_like(disc_fake_y))

    generator_g_gradients = tape.gradient(total_gen_g_loss, generator_g.trainable_variables)
    generator_f_gradients = tape.gradient(total_gen_f_loss, generator_f.trainable_variables)

    discriminator_x_gradients = tape.gradient(disc_x_loss, discriminator_x.trainable_variables)
    discriminator_y_gradients = tape.gradient(disc_y_loss, discriminator_y.trainable_variables)

    generator_g_optimizer.apply_gradients(zip(generator_g_gradients, generator_g.trainable_variables))
    generator_f_optimizer.apply_gradients(zip(generator_f_gradients, generator_f.trainable_variables))

    discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients, discriminator_x.trainable_variables))
    discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients, discriminator_y.trainable_variables))

# 학습 실행
EPOCHS = 50
for epoch in range(EPOCHS):
    for image_x, image_y in tf.data.Dataset.zip((train_horses, train_zebras)):
        train_step(image_x, image_y)

# 이미지 변환 및 확인
sample_image = next(iter(train_horses))
fake_image = generator_g(sample_image)

결론

CycleGAN은 이미지 변환 작업에서 매우 효과적이지만, 다른 방법들도 함께 고려해볼 수 있습니다. 다양한 데이터셋과 기술을 활용하여 최적의 방법을 찾아보는 것이 중요합니다. CycleGAN은 그 중 하나의 좋은 선택지입니다.

 

'AI > YOLO' 카테고리의 다른 글

YOLOv8  (0) 2024.05.21
YOLOv8-오탐률(False Positive Rate) 낮추기  (0) 2024.05.21

댓글