본문 바로가기

카테고리 없음

StyleGAN XL를 사용하여 ImageNet 이미지를 다양하게 생성하는 방법

※ 본 게시글은 필자의 개인 경험을 토대로 작성된 것으로, 잘못된 정보를 포함하고 있을 수 있습니다. ※

 

  사전 학습된(pre-trained) StyleGAN 모델을 사용하면 다양한 그럴싸하고 자연스러운 이미지를 만들어 낼 수 있다. 대표적으로 ImageNet에 대하여 학습된 StyleGAN XL 모델을 사용하면 512 X 512 크기ImageNet 클래스의 다양한 이미지를 만들 수 있는데, 이번 포스팅에서는 해당 내용을 다루고자 한다.

 

  ▶ GitHub StyleGAN XL 저장소: https://github.com/autonomousvision/stylegan-xl

 

GitHub - autonomousvision/stylegan-xl: [SIGGRAPH'22] StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets

[SIGGRAPH'22] StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets - GitHub - autonomousvision/stylegan-xl: [SIGGRAPH'22] StyleGAN-XL: Scaling StyleGAN to Large Diverse Datasets

github.com

 

※ GitHub에서 StyleGAN XL 소스 코드 다운로드해서 불러오기 ※

 

  이러한 Stylegan XL 저장소를 GitHub에서 클론(clone)하여 다운로드하고 가지고 있자. 다음과 같이 git clone 명령어를 사용하여 전체 소스 코드를 불러올 수 있다.

git clone https://github.com/autonomousvision/stylegan-xl


  이후에 Jupyter Notebook에서 해당 소스 코드를 사용하기 위해 sys 모듈을 이용하여 경로에 추가한다.

 

import sys

sys.path.append("./stylegan-xl")

 

※ 사용할 GPU 설정하기 ※

 

  필자의 경우 다음과 같이 GPU 번호 0, 1, 2, 3번을 사용하기로 Jupyter Notebook 상에서 설정할 것이다.

 

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

 

※ StyleGAN XL 모델 불러오기 ※

 

  이어서 실질적으로 생성자(generator)를 사용하기 위한 필수적인 유틸리티를 다음과 같이 실행하여 불러온다.

 

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
"""Generate images using pretrained network pickle."""

import os
import re
from typing import List, Optional, Tuple, Union

import click
import dnnlib
import numpy as np
import PIL.Image
import torch

import legacy
from torch_utils import gen_utils

#----------------------------------------------------------------------------

def parse_range(s: Union[str, List]) -> List[int]:
    '''Parse a comma separated list of numbers or ranges and return a list of ints.
    Example: '1,2,5-10' returns [1, 2, 5, 6, 7]
    '''
    if isinstance(s, list): return s
    ranges = []
    range_re = re.compile(r'^(\d+)-(\d+)$')
    for p in s.split(','):
        m = range_re.match(p)
        if m:
            ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
        else:
            ranges.append(int(p))
    return ranges

#----------------------------------------------------------------------------

def parse_vec2(s: Union[str, Tuple[float, float]]) -> Tuple[float, float]:
    '''Parse a floating point 2-vector of syntax 'a,b'.
    Example:
        '0,1' returns (0,1)
    '''
    if isinstance(s, tuple): return s
    parts = s.split(',')
    if len(parts) == 2:
        return (float(parts[0]), float(parts[1]))
    raise ValueError(f'cannot parse 2-vector {s}')

#----------------------------------------------------------------------------

def make_transform(translate: Tuple[float,float], angle: float):
    m = np.eye(3)
    s = np.sin(angle/360.0*np.pi*2)
    c = np.cos(angle/360.0*np.pi*2)
    m[0][0] = c
    m[0][1] = s
    m[0][2] = translate[0]
    m[1][0] = -s
    m[1][1] = c
    m[1][2] = translate[1]
    return m

#----------------------------------------------------------------------------

 

 

  이후에 사전에 다운로드 받은 imagenet512.pkl 파일을 불러온다. 해당 사전 학습된 모델 파일 또한 마찬가지로 StyleGAN XL 공식 GitHub에서 받을 수 있다.

import pickle

device = torch.device('cuda')

# 사전 학습된(pre-trained) StyleGAN 모델 불러오기
print('Loading networks from ...')
with dnnlib.util.open_url("./stylegan-xl/imagenet512.pkl") as f:
    G = legacy.load_network_pkl(f)['G_ema']
    G = G.eval().requires_grad_(False).to(device)

# Construct an inverse rotation/translation matrix and pass it to the generator.
rotate = 0
translate = [0,0]

if hasattr(G.synthesis, 'input'):
    m = make_transform(translate, rotate)
    m = np.linalg.inv(m)
    G.synthesis.input.transform.copy_(torch.from_numpy(m))

 

  다음과 같이 생성자(generator) 네트워크가 정상적으로 불러와진 것을 확인할 수 있다.

 

 

※ StyleGAN XL로 이미지 생성해보기 ※


  적절한 배치 크기(batch size)를 설정한 뒤에, 다음과 같이 이미지를 생성해 볼 수 있다. 필자는 간단히 배치 사이즈를 4로 설정하고, truncation trick의 값으로 0.7을 설정했다. 참고로 truncation_psi 값이 0에 가까울수록 클래스 평균에 가까운 데이터가 생성된다. 참고로 ImageNet-1k에서 클래스 인덱스 471는 대포(cannon)에 해당한다.

 

batch_sz = 4
truncation_psi = 0.7
seed = 1234
centroids_path = None
class_idx = 471

# 잠재 벡터(latent vector) w 생성 및 이미지 생성
w = gen_utils.get_w_from_seed(G, batch_sz, device, truncation_psi, seed=seed, centroids_path=centroids_path, class_idx=class_idx)
img = gen_utils.w_to_img(G, w, to_np=True)

# 생성된 이미지 저장
outdir = "StyleGAN_example"
os.makedirs(outdir, exist_ok=True)
PIL.Image.fromarray(gen_utils.create_image_grid(img), 'RGB').save(f'{outdir}/seed{seed:04d}.png')


  결과적으로 만들어진 이미지를 확인해 보자. 다음과 같이 4개의 생성된 이미지가 격자 형태로 만들어져 출력된 것을 확인할 수 있다. 조금 어색해 보이는 부분들이 있지만, 상당히 대포(cannon)의 형상을 잘 보이고 있다.

 

 

  이제 모든 클래스에 대하여 각각 20장씩 만들어보자. 이렇게 하게 되면, 총 2만장의 이미지를 만들게 될 것이다. 필자의 TITAN RTX에서는 전부 돌아가기 위해 약 2시간 정도가 소요된다. 즉, 2만 장의 이미지를 생성하기 위해 결과적으로 2시간이 소요된 것이다. 혹은 만약에 총 10만 장의 이미지를 만든다면, 10시간 정도가 소요될 것이다.

 

import time

!rm -rf ./imagenet_sampling_1
!rm -rf ./imagenet_sampling_1_grid

outdir = "./imagenet_sampling_1"
outdir_grid = "./imagenet_sampling_1_grid"

os.makedirs(outdir, exist_ok=True)
os.makedirs(outdir_grid, exist_ok=True)

made_cnt = 0
truncation_psi = 1.0
start_time = time.time()

# 각 클래스에 대하여
for class_idx in range(1000):
    # 배치 크기(batch size)만큼 잠재 벡터(latent vector) 생성
    batch_sz = 20
    seed = random.randint(0, int(2 ** 32 - 1))
    ws = gen_utils.get_w_from_seed(G, batch_sz, device, truncation_psi, seed=seed, centroids_path=centroids_path, class_idx=class_idx)
    ws = ws.to(device)
    print(class_idx, ws.shape)

    # 무작위로(50%) 값 증폭하기
    for i in range(len(ws)):
        random_data = random.randint(1, 2)
        if random_data == 1:
            ws[i] *= 1

    # latent vector를 사용하여 이미지 생성하기
    imgs = gen_utils.w_to_img(G, ws, to_np=True)

    # 격자 형태로 이미지 저장하기
    PIL.Image.fromarray(gen_utils.create_image_grid(imgs), 'RGB').save(f'{outdir_grid}/class_{class_idx}_seed{seed}.png')
    # 각 이미지를 개별적으로 저장하기
    for i, img in enumerate(imgs):
        PIL.Image.fromarray(gen_utils.create_image_grid(np.expand_dims(img, axis=0)), 'RGB').save(f'{outdir}/{made_cnt}.png')
        made_cnt += 1

    print(f"[{class_idx}/{1000}] {time.time() - start_time:.2f} seconds elapsed.")

 

 

  예를 들어 class 2에 대하여 만들어진 격자 이미지는 다음과 같다. 참고로 ImageNet-1k에서 두 번째 클래스는 백상아리(white shark) 클래스에 해당한다. 20장의 이미지가 차곡차곡 잘 담긴 것을 확인할 수 있다. 아래의 이미지들은 전부 다 생성된 이미지이며, 실제로 존재하지 않는 동물 사진이다. 만들어진 것들이다. 심지어 위에 보이는 코드처럼 truncation trick을 위한 파라미터의 값을 1.0으로 설정했으니, truncation 적용을 하지 않고 만들어진 이미지인데도 우수하다.

 

 

  코드를 실행한 결과는 다음과 같다. 1분정도 기다리면, 10개의 클래스에 대하여 처리가 된다. 즉, 1,000개의 모든 클래스에 대해서 다 처리되기 위해서는 넉넉히 2시간 정도는 잡아야 한다.

 

 

※ StyleGAN XL을 이용해 Latent Vector를 변경해보기 ※

 

  필자는 이후에 각 클래스의 평균적인 정보(평균 latent vector)를 빼는(제거하는) 방식으로  한 번 의미없는 이미지를 만들어 보는 작업도 해보았다. 이를 위해, 가장 먼저 truncation trick의 값으로 0을 설정한 뒤에, 각 클래스에 대하여 평균 벡터를 계산할 수 있었다.

# 각 클래스에 대한 평균 latent vector 계산하여 기록
batch_sz = 1
truncation_psi = 0

class_centers = []
for class_idx in range(1000):
    w = gen_utils.get_w_from_seed(G, batch_sz, device, truncation_psi, seed=seed, centroids_path=centroids_path, class_idx=class_idx)
    class_centers.append(w)

 

  이후에 다음과 같이 각 클래스마다 20개의 이미지를 생성할 수 있다. 단, 각 이미지에 대하여 해당 이미지의 클래스에 대한 평균적인 정보를 제거하는 방식을 선택했다.

 

import time

!rm -rf ./imagenet_sampling_1
!rm -rf ./imagenet_sampling_1_grid

outdir = "./imagenet_sampling_1"
outdir_grid = "./imagenet_sampling_1_grid"

os.makedirs(outdir, exist_ok=True)
os.makedirs(outdir_grid, exist_ok=True)

made_cnt = 0
truncation_psi = 1.0
start_time = time.time()

# 각 클래스에 대하여
for class_idx in range(1000):
    # 배치 크기(batch size)만큼 잠재 벡터(latent vector) 생성
    batch_sz = 20
    seed = random.randint(0, int(2 ** 32 - 1))
    ws = gen_utils.get_w_from_seed(G, batch_sz, device, truncation_psi, seed=seed, centroids_path=centroids_path, class_idx=class_idx)
    ws = ws.to(device)
    print(class_idx, ws.shape)

    ws -= class_centers[class_idx]
    # 무작위로(50%) 값 증폭하기
    for i in range(len(ws)):
        random_data = random.randint(1, 2)
        if random_data == 1:
            ws[i] *= 3

    # latent vector를 사용하여 이미지 생성하기
    imgs = gen_utils.w_to_img(G, ws, to_np=True)

    # 격자 형태로 이미지 저장하기
    PIL.Image.fromarray(gen_utils.create_image_grid(imgs), 'RGB').save(f'{outdir_grid}/class_{class_idx}_seed{seed}.png')
    # 각 이미지를 개별적으로 저장하기
    for i, img in enumerate(imgs):
        PIL.Image.fromarray(gen_utils.create_image_grid(np.expand_dims(img, axis=0)), 'RGB').save(f'{outdir}/{made_cnt}.png')
        made_cnt += 1

    print(f"[{class_idx}/{1000}] {time.time() - start_time:.2f} seconds elapsed.")

 

  결과적으로 다음과 같이 텍스처(texture)와 같은 이미지가 생성되는 것을 확인할 수 있다.

 

 

  더 나아가 의미 없는(특정한 클래스에 대한 정보를 가지고 있지 않은) latent vector를 서로 다른 클래스끼리 섞어서 새로운 의미없는 이미지를 더욱 다양하게 만들어 볼 수 있다.

 

import time
import random

!rm -rf ./imagenet_sampling_2
!rm -rf ./imagenet_sampling_2_grid

outdir = "./imagenet_sampling_2"
outdir_grid = "./imagenet_sampling_2_grid"

os.makedirs(outdir, exist_ok=True)
os.makedirs(outdir_grid, exist_ok=True)

made_cnt = 0
truncation_psi = 1.0
start_time = time.time()
max_iters = 1000

for iter in range(max_iters):
    # 배치 크기(batch size)만큼 잠재 벡터(latent vector) 생성
    batch_sz = 20
    k = 5 # 섞을 데이터 개수
    results = []
    for i in range(batch_sz):
        result = None # 하나의 이미지에 대한 latent vector
        sampled = random.sample(range(1000), k) # 1000개의 클래스에서 k개를 뽑아 mixing
        for class_idx in sampled:
            seed = random.randint(0, int(2 ** 32 - 1))
            w = gen_utils.get_w_from_seed(G, 1, device, truncation_psi, seed=seed, centroids_path=centroids_path, class_idx=class_idx)
            w = w.to(device)
            w -= class_centers[class_idx]
            if result == None:
                result = (w / k)
            else:
                result += (w / k)
        results.append(result)
    ws = torch.cat(results, dim=0).to(device)

    # 무작위로(50%) 값 증폭하기
    for i in range(len(ws)):
        random_data = random.randint(1, 2)
        if random_data == 1:
            ws[i] *= 3

    # latent vector를 사용하여 이미지 생성하기
    imgs = gen_utils.w_to_img(G, ws, to_np=True)

    # 격자 형태로 이미지 저장하기
    seed = random.randint(0, int(2 ** 32 - 1))
    PIL.Image.fromarray(gen_utils.create_image_grid(imgs), 'RGB').save(f'{outdir_grid}/class_{class_idx}_seed{seed}.png')
    # 각 이미지를 개별적으로 저장하기
    for i, img in enumerate(imgs):
        PIL.Image.fromarray(gen_utils.create_image_grid(np.expand_dims(img, axis=0)), 'RGB').save(f'{outdir}/{made_cnt}.png')
        made_cnt += 1

    print(f"[{iter}/{max_iters}] {time.time() - start_time:.2f} seconds elapsed.")

 

  만들어지는 이미지를 확인해 보면, 더욱 다채로운 텍스처(texture)가 섞인 다양한 이미지가 생성되는 것을 확인할 수 있다.