Gradio : 이미지를 사용하는 AI 웹 앱 만들기

등장

 

 

 

지난번에 이어서 Grrrrrrrrrrrrrrrrrrrrrradio를 이용해서

이미지를 선택하면 해당 이미지를 AI 모델에 넣은 후 결과를 보여주는 웹앱을 만들어 보고자 한다.

 

 

 

 

내가 사용한 AI 모델은 PReNet 이다.

PReNet은 이미지에서 비로 인한 노이즈를 제거해주는 모델이다. 자세한 사항은 아래 깃허브에ㅎㅎ,,

GitHub - csdwren/PReNet: Progressive Image Deraining Networks: A Better and Simpler Baseline (CVPR 2019)

 

 

 

 

깃허브를 보고 선행학습된 모델과 기본적인 코드를 가져오고, 내 컴퓨터에 맞게 환경을 설정하였다.

코드는 다음과 같다.

 

 

 

필요 라이브러리 임포트

import gradio as gr
import torch
from torch.autograd import Variable
from utils import *
from networks import *
import cv2
import os

여기서 utils 와 networks 는 PReNet의 깃허브 코드이다.

 

 

AI 모델 불러오기 및 GPU 가속 설정

model = PRN_r(6, 1)
model = model.cuda()
model.load_state_dict(torch.load('./net_latest.pth'))
iscuda = torch.cuda.is_available()

PReNet 의 깃허브에 있는 여러가지 모델중에 나는 PRN_r 모델을 정하였다. 

 

 

 

Test Data 불러오기 

data_path = './datasets/Rain100H/'
file_list = [data_path+file_name for file_name in os.listdir(data_path)]

모델의 학습에 사용된 테스트 데이터 중 일부를 불러온 것인데, 이는 나중에 사이트 화면에서 선택되는 이미지로 사용할 예정이다.

 

 

 

딥러닝 모델을 이용하여 입력받은 이미지의 노이즈를 제거하는 함수 

def deNosing(noise_img):
    y = np.array(noise_img)
    y = cv2.resize(y, (int(500), int(500)), interpolation=cv2.INTER_CUBIC)

    y = normalize(np.float32(y))
    y = np.expand_dims(y.transpose(2, 0, 1), 0)
    y = Variable(torch.Tensor(y))

    if iscuda:
        y = y.cuda()

    with torch.no_grad():
        if iscuda:
            torch.cuda.synchronize()

        out, _ = model(y)
        out = torch.clamp(out, 0., 1.)

        if iscuda:
            torch.cuda.synchronize()

    if iscuda:
        save_out = np.uint8(255 * out.data.cpu().numpy().squeeze())
    else:
        save_out = np.uint8(255 * out.data.numpy().squeeze())

    save_out = save_out.transpose(1, 2, 0)
    save_out = cv2.resize(save_out,(noise_img.width,noise_img.height),interpolation=cv2.INTER_CUBIC)

    return save_out

이제 입력으로 받는 이미지를, 모델의 입력 사이즈와 동일하게 리사이즈 해주고,

모델에 넣어서 노이즈 제거를 실행한 후 결과를 반환하는 함수이다.

 

 

 

Gradio 설정

demo = gr.Interface(fn=deNosing,
                    inputs=gr.components.Image(type='pil'),
                    outputs=gr.components.Image(type='pil'),
                    examples=file_list,
                    examples_per_page = 30,
                    title = 'Hello, My App with AI',
                    )
                    
demo.launch(debug=True, share=False)

gradio 홈페이지를 참고하면 여러가지 설정 값들을 볼 수 있다.

여기서 share 값을 True로 주면 외부에서 접속할 수 있는 public link 를 만들어 준다.

 

 

 

결과 화면

 

 

이미지의 경우는 홈페이지에 예시도 잘 되어 있고, 생각보다 만드는데 큰 어려움은 없었다.

 추후에 동영상을 이용하는 경우도 포스팅 하도록 하겠다.