2.2.4. 이미지 처리 과정 통합

1️⃣ 이미지 처리 과정

Untitled

이전까지 따로 진행되었던 각 과정이 한 번의 코드 실행만으로 진행되도록 통합
원본 사진과 스타일 종류가 입력되면 아래의 과정을 순서대로 거친 후 최종 이미지가 png 형식으로 저장됨
1. image segmentation 모델을 통해 반려견 영역/그 외 영역 분할
2. 반려견을 제외한 나머지 부분을 흰 색으로 변경
3. 원하는 스타일의 cartoonGAN 모델을 통해 그림으로 변환
4. 분할 과정에서 사용한 mask를 이용해 이미지 내 alpha channel 추가
5. real-esrGAN 모델을 통해 해상도 개선
위 그림에서 해상도 개선 전과 후의 차이를 쉽게 확인할 수 있도록, 해당 과정의 이미지만 확대했음

2️⃣ 코드

필요한 파일들이 존재하는 디렉토리 내에서 필요한 모듈이나 라이브러리 설치 등 환경 설정 후코드 실행 (Github 참고)

# 모듈 및 라이브러리 호출
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Layer, InputSpec
import PIL
import cv2
import matplotlib.pyplot as plt
import torch
from torchvision import transforms, models

# 변환 대상 사진과 원하는 스타일 지정
image_path = '#변환할 사진 경로#'
style = '#원하는 스타일명#'

print('강아지 영역 추출 시작')
# segmentation 코드는 이전과 동일하므로 생략
print('강아지 영역 추출 완료')

print('그림그리기 시작')
# cartonGAN 모델 구성 및 호출 함수 코드는 이전과 동일하므로 생략

# pretrained 모델을 이용할 경우
if (style == "shinkai") | (style == "hayao") | (style == "hosoda") | (style == "paprika"):
    PRETRAINED_WEIGHT_DIR = '#pretrained 모델 가중치 파일 경로#'
    cartoonGAN_model = load_model(style)
    input_image = cv2.cvtColor(result, cv2.COLOR_RGB2BGR) # result는 이전까지의 과정을 거친 이미지
    input_image = np.expand_dims(input_image, axis=0)
    output_image = cartoonGAN_model.predict(input_image)
    output_image = output_image[0]
    output_image = output_image[:,:,[2,1,0]]
    output_image = output_image * 0.5 + 0.5
    print('그림그리기 완료')
    print('배경 투명화 시작')
    alpha_output = cv2.cvtColor(output_image, cv2.COLOR_BGR2BGRA)
    resize_alpha = cv2.resize(alpha_scale, (output_image.shape[1],output_image.shape[0]))
    alpha_output[:,:,3] = resize_alpha
    alpha_output_2= alpha_output*255
    alpha_output_2 = alpha_output_2.astype(np.uint8)
    alpha_output_2 = cv2.cvtColor(alpha_output_2, cv2.COLOR_BGRA2RGBA)
    cv2.imwrite('#이미지를 저장할 경로#', alpha_output_2)
    print('배경 투명화 완료')

# 직접 학습한 모델 이용할 경우(모델 수나 경로에 맞춰 수정 필요)
else : 
    model_dir = '#저장한 모델 폴더 경로#'
    model = tf.saved_model.load(model_dir)
    f = model.signatures["serving_default"]
    input_image = result # result는 이전까지의 과정을 거친 이미지
    input_image = np.expand_dims(input_image, 0).astype(np.float32) / 127.5 -1
    output_image = f(tf.constant(input_image))['output_1']
    output_image = ((output_image.numpy().squeeze() + 1) * 127.5).astype(np.uint8)
    print('그림그리기 완료')
    print('배경 투명화 시작')
    alpha_output = cv2.cvtColor(output_image, cv2.COLOR_RGB2RGBA)
    resize_alpha = cv2.resize(alpha, (output_image.shape[1],output_image.shape[0]))
    alpha_output[:,:,3] = resize_alpha
    alpha_output_2= alpha_output
    alpha_output_2 = alpha_output_2.astype(np.uint8)
    alpha_output_2 = cv2.cvtColor(alpha_output_2, cv2.COLOR_RGBA2BGRA)
    cv2.imwrite('#이미지를 저장할 경로#', alpha_output_2)
    print('배경 투명화 완료')
    

print('해상도 높이기 시작')
!python inference_realesrgan.py \\
-n RealESRGAN_x4plus_anime_6B \\
-i #알파 채널이 추가된 이미지가 존재하는 경로#
-o #최종 결과 이미지가 저장될 경로#
print('해상도 높이기 완료')

# 아래부터는 각 과정간 결과물을 출력해보기 위한 코드로 실제 서비스에서는 불필요
last_output = cv2.imread('#최종 결과 이미지가 저장된 경로', cv2.IMREAD_UNCHANGED)
last_output = cv2.cvtColor(last_output, cv2.COLOR_BGRA2RGBA)

print('모든 결과 확인')
plt.figure(figsize=(20,15))
ax = plt.subplot(1,5,1)
ax.imshow(img)
ax.axis('off')
ax.set_title('raw')
ax = plt.subplot(1,5,2)
ax.imshow(result)
ax.axis('off')
ax.set_title('segmentation')
ax = plt.subplot(1,5,3)
ax.imshow(output_image)
ax.axis('off')
ax.set_title('cartonnGAN')
ax = plt.subplot(1,5,4)
ax.imshow(alpha_output)
ax.axis('off')
ax.set_title('alpha')
ax = plt.subplot(1,5,5)
ax.imshow(last_output)
ax.axis('off')
ax.set_title('read_esrGAN')

3️⃣ 결과 확인

이미지 좌측부터, 원본 사진, 분할, 변환, 배경 투명화, 해상도 개선 순

Untitled