17. 딥러닝 + 유니티 + VR 프로젝트 : Draw and painting

 

 

프로젝트 개요

VR게임상에서 그림을 그리고 그 그림을 딥러닝 모델로 판별하고 다시 VR환경에 관련 오브젝트를 띄운다. 이후, 스프레이, 색연필로 오브젝트로 색칠하여 전시하는 프로젝트.

---

딥러닝 (데이터셋, 모델선언, 트레이닝)

Google에서 제공하는 quickdraw dataset .npy형태를 받아 선언한 resnet 모델로 트레이닝.

 

batch : 256

 

10개의 카테고리의 경우

	CNN	RESNET34
5epoch	93%	94%
10epoch		96%

20개의 카테고리의 경우

	CNN	RESNET34
5epoch	85%	88%
10epoch	86%	90%

최종적으로 선택한 파라미터들

epoch : 10

model : resnet34

batch size : 256

lr : 0.1

optimizer : SGD

loss function : cross-entropy

accuracy : 90%

 

def resnet34(numclasses, pretrained=False):
    model = models.resnet34(pretrained)
    conv1_out_channels = model.conv1.out_channels
    model.conv1 = nn.Conv2d(1, conv1_out_channels, kernel_size=3,
                            stride=1, padding=1, bias=False)
    model.maxpool = nn.MaxPool2d(kernel_size=2)
    fc_features = model.fc.in_features
    model.fc = nn.Linear(fc_features, numclasses)
    return model

    def train():
        net.train()
        loss_avg = 0.0
        correct = 0
        # info printed in terminal
        # info printed in terminal
        data_loader = tqdm(train_loader, desc='Training')
        # data_loader = train_loader  # info log in logtext
        for batch_idx, (data, target) in enumerate(data_loader):
            if args.ngpu > 0:
                data, target = torch.autograd.Variable(data.cuda()), \
                    torch.autograd.Variable(target.cuda())
            else:
                data, target = torch.autograd.Variable(data.cpu()), \
                    torch.autograd.Variable(target.cpu())

            data = data.view(-1, 1, args.image_size, args.image_size)
            data /= 255.0

            # forward
            output = net(data)

            # backward
            optimizer.zero_grad()
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()

            # accuracy
            pred = output.data.max(1)[1]
            correct += float(pred.eq(target.data).sum())

            # exponential moving average
            loss_avg = loss_avg*0.2+float(loss)*0.8

        state['train_loss'] = loss_avg
        state['train_accuracy'] = correct/len(train_loader.dataset)

---

유니티와 파이썬의 연결

유니티에서 그린 그림을 파이썬에게 넘겨주기 위한 코드

유니티에서 강화학습에 사용되는 Ml-agent를 이용하여 강화학습을 하는게 아니라 통신에만 사용.

연결 모듈을 만들어 재사용성이 용이하게 구성

 

유니티상에서 설정 28X28의 형태로 데이터를 파이썬상에 보낸다.

 

연결 모듈 코드

env_name = "../DrawAndPainting_VR/Build/DrawAndPainting_VR"
# env_name = "C:/Users/user/Documents/GitHub/IMG_Transport_Sample/Build/IMG_Transport_Sample"  # Name of the Unity environment binary to launch
train_mode = False  # Whether to run the environment in training or inference mode

import matplotlib.pyplot as plt
import numpy as np
import sys
import cv2
import datetime

from PIL import Image
from matplotlib import cm

from mlagents.envs.environment import UnityEnvironment

def ConnectToUnity():
    print("ConnectToUnity")

    # 파이썬 버전 검사
    print("Python version:")
    print(sys.version)
    if (sys.version_info[0] < 3):
        raise Exception("ERROR: ML-Agents Toolkit (v0.3 onwards) requires Python 3")

    # 파일의 경로를 가져와서 유니티를 킨다.
    # print("파일의 경로를 가져와서 유니티를 킨다.")
    env = UnityEnvironment(file_name=env_name)

    return env

def DisconnectToUnity(env):
    print("DisconnectToUnity")
    env.close()

def GetImgFromUnity(env_info):
    print("GetImgFromUnity")

    img = np.uint8(255 *np.array(env_info.visual_observations[0]))

    #img reshape
    img = np.reshape(img, (28, 28))
    return img


# 학습정보 리셋
def EnvReset(env):
    # 첫번째 브레인을 가저온다.
    #print("첫번째 브레인을 가저온다.")
    default_brain = env.brain_names[0]
    brain = env.brains[default_brain]

    # 학습 정보 리셋
    env_info = env.reset(train_mode=train_mode)[default_brain]

    return env_info

# 유니티에서 while 문에대한 bool 값을 정해줌
def PyLoopBoolFromUnity(env_info):
    #print(env_info.vector_observations[0][0])
    return env_info.vector_observations[0][0]

# 유니티에서 이미지 체크할지 bool 값을 정해줌
def ImgCheckBoolFromUnity(env_info):
    #print(env_info.vector_observations[0][1])
    return env_info.vector_observations[0][1]

 

하나의 이미지를 체크하는 코드

env_name = "../DrawAndPainting_VR/Build/DrawAndPainting_VR"
# env_name = "C:/Users/user/Documents/GitHub/IMG_Transport_Sample/Build/IMG_Transport_Sample"  # Name of the Unity environment binary to launch
train_mode = False  # Whether to run the environment in training or inference mode

import os
import cv2
import sys
import json
import argparse
import torch.nn as nn
import torch.nn.functional as F

from tqdm import tqdm
from Model.nets import resnet34
from torchvision import transforms

# dataset
from DataUtils.load_data import QD_Dataset
from mlagents.envs.environment import UnityEnvironment

import torch
import numpy as np
import matplotlib.pyplot as plt
import testUnityCommunication

# model - resnet34
from Model.nets import resnet34
#from Model.nets import convnet


if __name__ == '__main__':
    def loadCnnModel():
        #load cnn model
        net = resnet34(20)
        PATH = "./Checkpoints/model.pytorch"
        net.load_state_dict(torch.load(PATH))
        return net

    def PredictSingleImage(net, data):
        net.eval()
        loss_avg = 0.0
        correct = 0

        tensorData = torch.from_numpy(data).type(torch.FloatTensor)
        tensorData = tensorData.view(-1, 1, 28, 28)
        tensorData /= 255.0

        # forward
        output = net(tensorData)
        print(output)

        #predict
        _,output_index = torch.max(output, 1)
        print(output_index)

        result = int(output_index[0])

        return result

    # connect Image
    env = testUnityCommunication.ConnectToUnity()

    # 학습정보 리셋
    env_info = testUnityCommunication.EnvReset(env)
    # 파이썬 구문 반복 Bool
    isPyLoop = testUnityCommunication.PyLoopBoolFromUnity(env_info)

    # 브레인 가져오기
    default_brain = env.brain_names[0]
    brain = env.brains[default_brain]

    while isPyLoop:
        # 이미지 체크할지 bool
        isImgCheck = testUnityCommunication.ImgCheckBoolFromUnity(env_info)

        #print (isImgCheck)
        # 브레인 엑션 사이즈
        action_size = brain.vector_action_space_size

        if isImgCheck:  # 이미지 인식을 할경우 (해당 이미지 번호를 유니티에 반환 1~20)
            net = loadCnnModel()

            # getImage 1 : 이미지를 불러오기
            # singleImagePATH = "./mydraw/glasses4.png"
            # gray = cv2.imread(singleImagePATH, cv2.IMREAD_GRAYSCALE)
            # resizeImage = cv2.resize(gray, (28, 28))

            # get Image 2 :
            resizeImage = testUnityCommunication.GetImgFromUnity(env_info)

            # convert numpy
            transform = np.array(resizeImage)
            # 색 반전
            data = np.where(transform < 25, 255, np.where(transform > 25, 0, transform))

            # image show
            plt.imshow(data, cmap="gray")
            #plt.show()

            # 인식한 애 보여주기
            #print(data)

            # predict
            result = PredictSingleImage(net, data)+1
            env_info = env.step(result)[default_brain]
            #print("env.step 실핼 IF")

        else :  # 이미지 인식을 안할경우 (0을 반환 )
            env_info = env.step(0)[default_brain]
            # print("env.step 실핼 else")

        # 학습정보 리셋
        env_info = testUnityCommunication.EnvReset(env)

        # 파이썬 구문 반복 Bool
        isPyLoop = testUnityCommunication.PyLoopBoolFromUnity(env_info)

    testUnityCommunication.DisconnectToUnity(env)

---

VR 환경구성 (전체 흐름)

 

---

시연영상

 

 그림그리기

 

색칠하기

그 외 기능 영상은 생략

---

아쉬운점

일반적인 이미지처리 모델의 한계상 구글의 낙서데이터는 내가 원하는 수준까지의 정확도를 끌어 올릴 순 없었다. 그리고 이러한 단점에도 unity에서 스트로크 데이터를 보내는데 어려움이 있어 시간 관계상 이미지 처리모델을 쓸 수 밖에 없었다는 점이 아쉬웠다.

 

그리고 유니티 환경에서는 개인적인 사정으로 VR환경에 신경을 많이 쓰지못해 퀄리티가 많이 떨어져 그 부분이 아쉬웠다. 조금 더 다듬었다면 더 괜찮은 프로젝트가 되지 않았을까 생각한다.

---

github.com/Kwonkunkun/DrawAndPainting_VR

Kwonkunkun/DrawAndPainting_VR

딥러닝 관련 코드들이 같이 안올라가 있는데.. 그건 나중에 수정해서 깃에 추가할 계획이다. < -추가
https://github.com/Kwonkunkun/DrawAndPainting_Pytorch