[2차 프로젝트] Colab을 활용한 CartoonGAN 사용 - 2

저번 포스팅 내용에는 구글 코랩을 활용해 깃허브 오픈소스 코드를 사용하는 방법과

저희의 프로젝트 GAN모델을 실행시키는 과정을 담았습니다.

 

이번 포스팅 내용에는 구글 코랩이 아닌 코랩의 중요한 부분들을 local repository에서 실행시키고 사용하는 방법과

최종적으로 프로젝트에서 사용한 모델의 구성은 어떻게 되는지 담아보도록 하겠습니다.

 

먼저, 저희 프로젝트에서 사용하는 모델이 이미지를 처리하는 과정을 알아두시면 도움이 될 것 같습니다.

이미지 처리과정은 다음과 같습니다.

 

제가 만들었습니다 불펌 ㄴㄴ

사용자가 이미지를 넣으면, 일단 원본 이미지에서 강아지를 객체 인식하여 객체 바깥을 흰색으로 처리하고,

객체 내부는 검정색으로 채워 배경으로 사용하게 됩니다. 그 후, 해상도가 높은 이미지일수록 더 좋은 변환결과가

나오기 때문에 realESRGAN을 통해 원본 이미지의 품질을 높여줍니다. 그 후, 저희가 수정한 cartoonGAN을 통해

이미지를 만화화시킨 다음, 앞서 만들어졌던 흰색 배경과 합쳐 변환된 이미지를 최종 완성하게 됩니다.

다음은 저희 프로젝트에서 사용한 코드입니다.

먼저 필요한 도구들을 import해줍시다.

위의 네개의 import 부분은 라이브러리가 아니라 local repository에 있는 py파일입니다.

이후 torch부터 실질적인 import라고 생각해주세요.

import CartoonGAN_model_modified as modified_models
from CartoonGAN_train import CartoonGANTrainer
from config import CartoonGANConfig as Config
from dataloader import load_image_dataloader
import torch
import argparse
import torchvision.utils as tvutils
import os
from torchvision import transforms, models

import glob
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import torch
import time
import cv2
import pymysql

!jupyter nbextension enable --py widgetsnbextension
# Clone Real-ESRGAN and enter the Real-ESRGAN
!git clone https://github.com/xinntao/Real-ESRGAN.git
%cd Real-ESRGAN
# Set up the environment
!pip install basicsr
!pip install facexlib
!pip install gfpgan
!pip install -r requirements.txt
!python setup.py develop
# Download the pre-trained model
!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
    
from flask import Flask, request, send_file

import os
import shutil

## 순서

- input image 받기
- sementic segmation 으로 누끼따기 >> 배경으로 사용
- RealESR GAN을 통해 화질과 크기를 더 좋게 만들기
- CartoonGAN model 적용
- 원본사진 크기로 resize >> 적용할 피사체로 사용
- 원본에서 따온 배경(255의 흰색) + 원본 사진크기로 resize된 cartoonGAN결과값
- 완료

## FID_score

- 실행 커맨드

    - python FID_SCORE.py ["data/photo/"] ["data/animation/"] --batch-size [200] --dims [dims]

 

- 해결해야 할 의문점

    - Config.batch_size 를 batch_size에 넣어도 되나?
    - dims 값에 뭘 넣어야 할지 모르겠음. config.py >> 검토 후 찾지 못함. num_worker랑 똑같은건가?

- 실행 오류 해결법

    - 그냥 주소 앞에 두개만 채운다.

 

## 요구사항
- realESR gan 모델 pth파일이 pretrained_model 폴더 내에 있어야함
- checkpoints/CartoonGAN/안에 학습된 cartoonGAN모델.ckpt 가 폴더내에 있어야함
- pytorch cuda 환경구축
- flask 환경구축

## 초기 설정
- C:\Users\(사용자명)\cartoonGAN\semantic-segmentation-pytorch\Real-ESRGAN\upload 에 변환될 이미지를 넣으면 된다.
- C:\Users\(사용자명)\cartoonGAN\semantic-segmentation-pytorch\checkpoints\CartoonGAN 에 사용할 모델 ckpt를 넣으면 된다.

 

---

app = Flask(__name__)

###################### 폴더 지정 ######################

#cartoon_model_path = cartoonGAN 모델 ckpt가 있는 폴더를 가리킵니다. 모델 변경시 path를 변경하세요.
#test_img_path = cartoonGAN을 적용시킬 사진이 있는 곳
#upload_folder = realESR GAN을 적용시킬 사진이 있는 곳

#realESR GAN이 적용되어 나오는 사진이 있는 곳 
ESR_output_folder = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/Real-ESRGAN/results/'
#realESR GAN이 적용되어 나오는 사진을 복사하여 CartoonGAN의 재료가 될 사진이 있는 곳 
GAN_input_folder = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/data/test/'
#semantic segmentation의 재료가 될 카툰갠 적용사진이 있는 곳
GAN_output_folder = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/generated_images/CartoonGAN/'

#cartoonGAN 모델이 있는곳
cartoon_model_path = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/checkpoints/CartoonGAN/MIXcheckpoint-epoch-300.ckpt'
#카툰갠을 적용시킬 사진이 있는곳(GAN_input폴더와는 다름)
test_image_path = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/data/'

#realESR 업로드 이름 폴더 제어
upload_folder = 'upload'
#realESR 결과 이미지가 저장되는 폴더
result_folder = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/Real-ESRGAN/results'


###################### 폴더 지정 ######################\

#업로드 폴더와 cartoon img path가 없을 경우 만들어주는 역할 
if not os.path.exists(upload_folder) :
    print('이미지를 업로드할 폴더가 없어요!')
    print('realESRGAN에 적용시킬 upload img를 올려줄 폴더를 만들고있어요.')
    os.mkdir(upload_folder)
    
#이 코드가 진짜로 필요한가?
# if not os.path.exists(cartoon_img_path) :
#     print('이미지를 업로드할 폴더가 없어요!')
#     print('CartoonGAN에 적용시킬 upload img가 될 폴더를 만들고있어요.')
#     os.mkdir(cartoon_img_path)

#사용자 지정함수를 정의하는 형태로 모델을 돌림

###################### <RealESR GAN def> ##########################

def real_ESR():   
    import os
    import shutil  
    print('realESR GAN 적용에 필요한 모든 조건을 만족시켜주셔서 감사합니다!')
    print('realESR GAN을 적용시키는 중입니다. 시간이 조금 걸리니 양해해주세요.')
    #realESR GAN 작동 시키는 부분
    !python inference_realesrgan.py -n RealESRGAN_x4plus -i upload --outscale 3.5 --face_enhance
    
    #Testing 0 0014가 프린트된다. 0은 index 0014는 파일이름
    #변환된 파일은 0014_out.jpg 형태로 저장된다.
    
###################### </RealESR GAN def> ##########################

###################### <Cartoon GAN def> ##########################

def load_model(generator, discriminator, checkpoint_path):
    checkpoint = torch.load(checkpoint_path)
    generator.load_state_dict(checkpoint['generator_state_dict'])
    discriminator.load_state_dict(checkpoint['discriminator_state_dict'])


def load_generator(generator, checkpoint_path):
    checkpoint = torch.load(checkpoint_path, map_location=Config.device)
    generator.load_state_dict(checkpoint['generator_state_dict'])


def generate_and_save_images(generator, test_image_loader, save_path):
    # for each image in test_image_loader, generate image and save
    generator.eval()
    torch_to_image = transforms.Compose([
        transforms.Normalize(mean=(-1, -1, -1), std=(2, 2, 2)),  # [-1, 1] to [0, 1]
        transforms.ToPILImage()
    ])

    image_ix = 0
    with torch.no_grad():
        for test_images, _ in test_image_loader:
            test_images = test_images.to(Config.device)
            generated_images = generator(test_images).detach().cpu()

            for i in range(len(generated_images)):
                image = generated_images[i]
                image = torch_to_image(image)
                image.save(os.path.join(save_path, '{0}.jpg'.format(image_ix)))
                image_ix += 1


def cartoon_GAN():

    device = torch.device('cpu')

    Generator = modified_models.Generator
    #Discriminator = modified_models.Discriminator
    #FeatureExtractor = modified_models.FeatureExtractor

    generator = Generator().to(device)

    generator.eval()

    print('모델을 불러오고있어요.')
    
    load_generator(generator, cartoon_model_path)
    
    print('모델이 불러와졌어요. 좋은 모델이네요!')
    
    #카툰화 시킬 이미지를 불러오고 기존에 config파일에서 지정해 놨던 대로 이미지를 변환한다.
    print('변환할 이미지를 탐색합니다.')
    test_images = load_image_dataloader(root_dir=test_image_path, batch_size=1, shuffle=False)
    
    print('변환할 이미지를 탐색완료!. Cartoon GAN을 통해 이미지를 변환하는 작업을 시작할게요.')
    
    image_batch, _ = next(iter(test_images))
    image_batch = image_batch.to(Config.device)

    with torch.no_grad():
        new_images = generator(image_batch).detach().cpu()

    #각각 무슨 이미지인지 모르겠음
    tvutils.save_image(image_batch, 'test_images.jpg', nrow=4, padding=2, normalize=True, range=(-1, 1))
    tvutils.save_image(new_images, 'generated_images.jpg', nrow=4, padding=2, normalize=True, range=(-1, 1))

    print("변환된 개리커쳐를 저장하고 있어요.")
    # 저장경로
    generate_and_save_images(generator, test_images, 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/generated_images/CartoonGAN/')

###################### </Cartoon GAN def> #########################
    
    
################## <semantic segmentation 및 flask 코드 실행> ########################
@app.route('/')
def final_code():

    ##### 이미지 불러와서 저장하기####
    
    # 0.사용자가 입력한 이미지 선명하게 변경 후 
#     if(input_img.shape < nXn):
#         print('입력된 이미지가 작아 real_ESR과정을 진행합니다.')
#         real_ESR()
#         print('realESR GAN적용이 끝났어요.')
#         print('realESR 결과 이미지를 cartoonGAN에 적용시키기위한 준비중이에요.')
        
    #real_ESR()

    
    
    ### shutill library & glob이용, result에 있는 내용을 data폴더로 복사
    
    
    #개선 전의 코드
    #origin = r'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/Real-ESRGAN/results/{}'.format(ESR_output_folder)
    #copy = r'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/data/test/{}'.format(ESR_output_folder)    
    
    #origin = r'{}'.format(ESR_output_folder)
    #copy = r'{}'.format(GAN_input_folder)    
    #shutil.mvdir(origin,copy)
    
    file_list = os.listdir(ESR_output_folder)
    mov_files = []
    for file in file_list:
        mov_files.append(file)
    
    for movfi in mov_files:
        shutil.copy(ESR_output_folder+'\\'+movfi, GAN_input_folder+'\\'+movfi)
        print(ESR_output_folder + '에 있는' + file + '이' + GAN_input_folder + '로 복사하였습니다.')
    
    print('파일 복사를 완료했어요.')
    
    #ㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡㅡ
      
    # 1.변환함수 실행후 변환 이미지 저장.
    print('모든 전처리 작업이 완료됐어요.')
    print('cartoonGAN 작업을 시작할게요.')
    cartoon_GAN()
    
    print('cartoonGAN 적용이 끝나고 누끼따는 과정에 들어갈게요.')
    print('이제 원본의 누끼딴 배경과 카툰간이 적용된 이미지를 합칠거라고 지시하셨어요.')
    
    # 2. 받아온 이미지와 하얀 배경 합성후 (강:검, 배:흰 이름:background 저장)
    model = models.segmentation.deeplabv3_resnet101(pretrained=True).eval()
    
    # 누끼딸 객체 종류들이 있는 라벨
    #객체에 따라 다른 색깔을 부여하게 된다.
    labels = ['background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
    cmap = plt.cm.get_cmap('tab20c')
    colors = (cmap(np.arange(cmap.N)) * 255).astype(np.int)[:, :3].tolist()
    np.random.seed(2020)
    np.random.shuffle(colors)
    colors.insert(0, [0, 0, 0])
    colors = np.array(colors, dtype=np.uint8)
    palette_map = np.empty((10, 0, 3), dtype=np.uint8)

    legend = []
    for i in range(21):
        legend.append(mpatches.Patch(color=np.array(colors[i]) / 255., label='%d: %s' % (i, labels[i])))
        c = np.full((10, 10, 3), colors[i], dtype=np.uint8)
    
    #누끼따는 기능.
    def segment(net, img):
        preprocess = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            ),
        ])

        input_tensor = preprocess(img)
        input_batch = input_tensor.unsqueeze(0)

        if torch.cuda.is_available():
            input_batch = input_batch.to('cuda')
            model.to('cuda')

        output = model(input_batch)['out'][0] 

        output_predictions = output.argmax(0).byte().cpu().numpy() 

        r = Image.fromarray(output_predictions).resize((img.shape[1], img.shape[0]))
        r.putpalette(colors)

        return r, output_predictions

    
    ################ 바꿀 사진 넣는 부분 ###############
    
    #imagename = realESR을 통해 선명해진 사진
    #Image.open에 glob해 온것을 반복문을 돌리자
    print('배경을 제공해줄 원본 이미지를 누끼따고있어요.')
    
    #imagename = '002.jpg'
    img = np.array(Image.open('C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/data/test/img_out.jpg'))
    fg_h, fg_w, _ = img.shape
    segment_map, pred = segment(model, img)
    
    #imagename2 = '0.jpg'
    img2 = np.array(Image.open('C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/generated_images/CartoonGAN/0.jpg'))
    fg_h2, fg_w2, _ = img2.shape
    segment_map2, pred2 = segment(model, img2)
    imgs = cv2.resize(img2, dsize=(fg_w, int(fg_w * fg_h2 / fg_w2)))
    fg_h2, fg_w2, _ = img2.shape
    
    margin = (fg_h2 - fg_h) // 2

    if margin > 0:
        img2 = img2[margin:-margin, :, :]
    else:
        img2 = cv2.copyMakeBorder(img2, top=abs(margin), bottom=abs(margin), left=0, right=0, borderType=cv2.BORDER_REPLICATE)

    img2 = cv2.resize(img2, dsize=(fg_w, fg_h))
    
    
    ################ 배경 사진 넣는 부분###############
    print('하얀색 배경을 적용시키고 있어요.')
    background = np.array(Image.open('C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/data/white.jpg'))
    bg_h, bg_w, _ = background.shape

    background = cv2.resize(background, dsize=(fg_w, int(fg_w * bg_h / bg_w)))
    
    bg_h, bg_w, _ = background.shape

    margin = (bg_h - fg_h) // 2

    if margin > 0:
        background = background[margin:-margin, :, :]
    else:
        background = cv2.copyMakeBorder(background, top=abs(margin), bottom=abs(margin), left=0, right=0, borderType=cv2.BORDER_REPLICATE)

    background = cv2.resize(background, dsize=(fg_w, fg_h))
    
    ################ 카툰간을 누끼딴거 하얀색 배경과 합성  ################

    #누끼딸 객체 지정하는 부분
    mask1 = (pred == 12).astype(float) * 255 # 강아지
    mask2 = (pred == 8).astype(float) * 255 # 고양이

    mask = mask1 + mask2
    _, alpha = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY)

    alpha = cv2.GaussianBlur(alpha, (7, 7), 0).astype(float)

    alpha = alpha / 255. # (height, width)
    alpha = np.repeat(np.expand_dims(alpha, axis=2), 3, axis=2) 
    # 사진 누끼 딴 부분
    foreground = cv2.multiply(alpha, img2.astype(float))
    #배경에 사진 부분 도려낸 나머지
    background = cv2.multiply(1. - alpha, background.astype(float))  

    ################ 결과 ################
    result = cv2.add(foreground, background).astype(np.uint8)
    print('하얀색 배경을 적용시켰어요.')
    
    ##### 현재 시간으로 파일명 지정 ####
    filename = '0'

    # 저장하기
    Image.fromarray(result).save('C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/result/'+filename+'.jpg')
    
    ######### 저장경로지정 ########
    #path = 'C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/result/'+filename
    
    print('이미지 저장이 완료됐어요.')
    
    ######## db저장 #########
#     db = pymysql.connect(host='project-db-stu.ddns.net',port=3307,user='ysjj',passwd='1234',db='ysjj',charset='utf8')
#     cursor = db.cursor()
#     sql = """insert into image(image_path,user_id) values(%s,'sunju')"""
#     cursor.execute(sql,path+'.jpg')
#     db.commit()
#     print('프로젝트 서버에 이미지 저장이 완료됐어요.')
#     db.close()
    
    return send_file(('C:/Users/smhrd/cartoonGAN/semantic-segmentation-pytorch/result/'+filename+'.jpg'), mimetype='image/jpeg'
                     )
    print('Flask return값으로 최종 이미지를 전송완료했습니다.')
    
    
if __name__ == '__main__':
    app.run('172.30.1.4', '9000')
################## </semantic segmentation> ########################