본문 바로가기
AI/딥러닝 프로젝트

[01 - ImageNet Sketch 데이터] - Augmentation [2]

by AI-BT 2024. 9. 26.
728x90
반응형

스케치 이미지에 유효한 Augmentation

이번 포스팅에서는 스케치 이미지 데이터셋을 학습할 때 사용한 다양한 Augmentation 기법들 중에서 성능에 유효했던 몇 가지를 공유하고자 합니다. 특히, 스케치 이미지의 특성에 맞춰 선명도와 밝기 조절을 통해 성능 개선을 이뤄낸 사례를 소개하겠습니다.

 

먼저, 사용했던 Augmentation 기법을 간단히 소개하면 아래와 같습니다.

  • Unsharp Mask: 이미지를 선명하게 만드는 기법.
  • RandomBrightnessContrast: 이미지의 밝기를 조절하는 기법.

이 두 가지 Augmentation은 특히 스케치 이미지에서 효과적인 성능 향상을 보여줬습니다. 이제 각각의 방법을 좀 더 자세히 살펴보겠습니다. 아래 코드는 적용한 전체 코드 입니다.

class UnsharpMask(A.ImageOnlyTransform):
    def __init__(self, kernel_size=7, sigma=1.5, amount=1.5, threshold=0, always_apply=False, p=1.0):
        super(UnsharpMask, self).__init__(always_apply, p)
        self.kernel_size = kernel_size
        self.sigma = sigma
        self.amount = amount
        self.threshold = threshold

    def apply(self, image, **params):
        return self.unsharp_mask(image)

    def unsharp_mask(self, image):
        blurred = cv2.GaussianBlur(image, (self.kernel_size, self.kernel_size), self.sigma)
        sharpened = cv2.addWeighted(image, 1.0 + self.amount, blurred, -self.amount, 0)
        if self.threshold > 0:
            low_contrast_mask = np.absolute(image - blurred) < self.threshold
            np.copyto(sharpened, ima
            
 UnsharpMask(kernel_size=7, sigma=1.5, amount=1.5, threshold=0, p=1.0),  # UnsharpMask 마스크 적용
 A.RandomBrightnessContrast(brightness_limit=(-0.2, -0.2), contrast_limit=0, p=1.0), # 20% 어둡게

 


1. UnsharpMask 

__init__ 메서드

  • kernel_size: 가우시안 블러를 적용할 때 사용하는 커널의 크기이다.
    이 값이 클수록 이미지가 더 많이 흐려진다. (기본값 7)
  • sigma: 가우시안 블러의 표준편차로, 블러의 강도를 조절한다. (기본값은1.5)
  • amount: 샤프닝의 강도를 결정한다. 값이 클수록 이미지가 더 날카롭게 된다. (기본값 1.5)
  • threshold: 이미지의 낮은 대비 부분을 샤프닝에서 제외하기 위한 임계값이다. (기본값은 0)
    샤프닝이 모든 픽셀에 적용된다.
  • always_apply: 변환을 항상 적용할지 여부를 결정한다. (기본값 False)
  • p: 변환이 적용될 확률이다. (기본값 1.0)

unsharp_mask 메서드

1. blurred: 입력 이미지에 가우시안 블러를 적용하여 흐릿한 이미지를 생성 한다.

blurred = cv2.GaussianBlur(image, (self.kernel_size, self.kernel_size), self.sigma)

 

2. sharpened: 원본 이미지와 블러 처리된 이미지를 가중치를 두고 더하여 샤프닝된 이미지를 만든다. addWeighted 함수는 두 이미지의 픽셀을 가중 평균한 결과를 반환 한다.

sharpened = cv2.addWeighted(image, 1.0 + self.amount, blurred, -self.amount, 0)

 

3. 임계값을 통한 샤프닝 조정: 만약 threshold 값이 0보다 크다면, 원본 이미지와 블러된 이미지 간의 차이가 임계값보다 작은 경우에는 샤프닝을 적용하지 않도록 조정한다.

if self.threshold > 0:
    low_contrast_mask = np.absolute(image - blurred) < self.threshold
    np.copyto(sharpened, image, where=low_contrast_mask)

 

 


2. RandomBrightnessContrast

Albumentations에서 제공하는 변환으로, 이미지의 밝기와 대비를 무작위로 조정한다.

이 경우 brightness_limit=(-0.2, -0.2)로 설정하여 이미지가 항상 20% 어두워지도록 했다.
contrast_limit=0으로 설정되어 있어 대비는 조정되지 않는다.

 A.RandomBrightnessContrast(brightness_limit=(-0.2, -0.2), contrast_limit=0, p=1.0)

3. 실제 이미지 비교

2번째 이미지는  RandomBrightnessContrast 20% 에 unsharpMask 를 적용한 이미지이고,

3번째 이미지는 RandomBrightnessContrast 20% 만 넣은 이미지 이다,

 

 

 

위에서 소개한 Unsharp MaskRandomBrightnessContrast는 특히 스케치 이미지에서 좋은 성능 향상을 보여주었습니다.

  1. Unsharp Mask는 이미지의 경계선을 강화하여 스케치 이미지의 세부 정보를 잘 학습할 수 있도록 했습니다.
  2. RandomBrightnessContrast는 이미지의 밝기를 조정하여 모델이 더 다양한 조명 환경에서도 정확하게 예측할 수 있도록 도왔습니다.

이 두 가지 Augmentation 기법은 실제로 스케치 이미지 데이터셋에서
학습 성능을 약 3~5% 이상 향상시키는 데 중요한 역할을 했습니다.


 

 

전체 코드 

# %%

import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import random
import os
import cv2


class UnsharpMask(A.ImageOnlyTransform):
    def __init__(self, kernel_size=7, sigma=1.5, amount=1.5, threshold=0, always_apply=False, p=1.0):
        super(UnsharpMask, self).__init__(always_apply, p)
        self.kernel_size = kernel_size
        self.sigma = sigma
        self.amount = amount
        self.threshold = threshold

    def apply(self, image, **params):
        return self.unsharp_mask(image)

    def unsharp_mask(self, image):
        blurred = cv2.GaussianBlur(image, (self.kernel_size, self.kernel_size), self.sigma)
        sharpened = cv2.addWeighted(image, 1.0 + self.amount, blurred, -self.amount, 0)
        if self.threshold > 0:
            low_contrast_mask = np.absolute(image - blurred) < self.threshold
            np.copyto(sharpened, image, where=low_contrast_mask)
        return sharpened
 
class AlbumentationsTransform_01:
    def __init__(self, is_train: bool = True):
        common_transforms = [
            A.Resize(224, 224),  # 이미지를 224x224 크기로 리사이즈
            UnsharpMask(kernel_size=7, sigma=1.5, amount=1.5, threshold=0, p=1.0),  # 언샤프 마스크 적용
            A.RandomBrightnessContrast(brightness_limit=(-0.2, -0.2), contrast_limit=0, p=1.0), # 20% 어둡게
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # 정규화            
            ToTensorV2(),  # albumentations에서 제공하는 PyTorch 텐서 변환
        ]

        if is_train:
            # 훈련용 변환
            self.transform = A.Compose(
                [
                    
                ] + common_transforms
            )
        else:
            self.transform = A.Compose(common_transforms)
            
    def __call__(self, image: np.ndarray) -> torch.Tensor:
        transformed = self.transform(image=image)
        return transformed['image']


class AlbumentationsTransform_02:
    def __init__(self, is_train: bool = True):
        common_transforms = [
            A.Resize(224, 224),
            A.RandomBrightnessContrast(brightness_limit=(-0.2, -0.2), contrast_limit=0, p=1),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2()
        ]
        
        if is_train:
            self.transform = A.Compose(
                [   
                    
                ] + common_transforms
            )
        else:
            self.transform = A.Compose(common_transforms)
            
    def __call__(self, image: np.ndarray) -> torch.Tensor:
        transformed = self.transform(image=image)
        return transformed['image']
    
# %%
# 시각적으로 변환된 이미지를 비교하는 함수
def plot_image_comparison(original_img, torchvision_img, albumentations_img, idx):
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
    ax[0].imshow(original_img)
    ax[0].set_title(f'Original Image {idx}')
    
    ax[1].imshow(torchvision_img.permute(1, 2, 0))  # torch tensor는 (C, H, W)이므로 (H, W, C)로 변환
    ax[1].set_title('Albumentations Unmask')
    
    ax[2].imshow(albumentations_img.permute(1, 2, 0))
    ax[2].set_title('Albumentations No Unmask')
    
    for a in ax:
        a.axis('off')
    
    plt.show()
    
def load_images_from_folder(image_folder, num_images):
    # 이미지 파일 목록 가져오기
    # all_images = [f for f in os.listdir(image_folder) if f.endswith(('JPEG'))]
    all_images = [f for f in os.listdir(image_folder) if f.endswith(('JPEG')) and not f.startswith('._')]
    
    # 이미지 파일이 num_images보다 적으면 모두 가져옴
    if len(all_images) < num_images:
        num_images = len(all_images)
    
    # 이미지 목록에서 랜덤하게 num_images개 선택
    random_images = random.sample(all_images, num_images)
    
    # 선택된 이미지 파일을 로드
    images = []
    for filename in random_images:
        img_path = os.path.join(image_folder, filename)
        # img = Image.open(img_path).convert('RGB')
        img = np.array(Image.open(img_path).convert("RGB"))
        images.append((img, filename))
    
    return images

# %%

# ./test_image 폴더에서 이미지 로드
image_folder = './data/test'
num_images = 10
images = load_images_from_folder(image_folder, num_images)

# Torchvision 및 Albumentations 변환 생성
albumentations_transform_01 = AlbumentationsTransform_01(is_train=True)
albumentations_transform_02 = AlbumentationsTransform_02(is_train=True)

# 10개의 이미지에 대해 변환 및 비교
for idx, (img, filename) in enumerate(images):
    # 원본 이미지
    original_img = img  
    
    # Torchvision과 Albumentations 변환 적용
    albumentations_img_01 = albumentations_transform_01(img)
    albumentations_img_02 = albumentations_transform_02(img)
    
    # 시각적으로 비교
    print(f"Comparing Image: {filename}")
    plot_image_comparison(original_img, albumentations_img_01, albumentations_img_02, idx)

 


마무리

데이터 증강(Augmentation)은 머신러닝 모델의 성능을 향상시키는 중요한 도구입니다. 특히 이미지 데이터셋에서 다양한 Augmentation을 시도해보는 것은 매우 효과적입니다.
이번 포스팅에서 소개한 Unsharp MaskRandomBrightnessContrast는 스케치 이미지에 특히 유효한 기법들이었으니, 유사한 이미지 처리 문제를 다룰 때 한 번 시도해 보시길 추천드립니다.

앞으로도 다양한 데이터 증강 기법에 대해 소개해드리겠습니다.
질문이나 피드백이 있으면 언제든 댓글로 남겨주세요!

 

감사합니다.

728x90
반응형

댓글