April 2, 2024

파이썬 SwinIR 업스케일 API 서버와 클라이언트 웹사이트 만들기

이미지 업스케일이란? 저해상도 이미지를 고해상도로 변환하여 이미지의 품질을 향상시키는 과정이며 이번 구현에서는 SwinIR을 사용하여 업스케일 API 서버와 클라이언트를 만들어보겠습니다. SwinIR은 이미지 복원을 위한 최신 딥러닝 모델 중 하나로, 단순 리사이즈가 아닌, AI 업스케일링이며 일반적인 리사이즈보다 훨씬 좋은 성능을 보입니다.

업스케일 API

프로젝트 개요

이번 프로젝트에서는 파이썬을 사용하여 이미지 업스케일 API 서버와 클라이언트 웹사이트를 구축하겠습니다.

인터넷에 흔히 있는 이미지 업스케일 API 서비스를 제작하기 위한 첫 걸음으로, 서버 사이드는 Flask와 SwinIR를 사용하여 구현하고 클라이언트는 간단하게 HTML, CSS, JS, Axios를 사용하여 웹사이트를 생성하겠습니다.

이미지를 업로드하면 서버로 이미지를 전송한 뒤 서버에서 이미지를 업스케일링 하고 클라이언트가 받는 방식입니다.

서버 사이드: Flask, SwinIR

클라이언트 사이드: HTML, CSS, JavaScript, Axios

SwinIR 설치

우선 깃허브에서 SwinIR 코드와 사전 학습된 모델을 다운받아 경로의 model_zoo 폴더에 넣어줍니다.

GitHub – JingyunLiang/SwinIR: SwinIR: Image Restoration Using Swin Transformer (official repository)

Release Pretrained models, supplementary and visual results · JingyunLiang/SwinIR · GitHub

코드와 모델을 다운받으셨으면 간단한 REST API를 구축합니다. 이 API는 이미지를 받아 업스케일 후 반환하는 기능을 합니다.

업스케일 API 구현

from flask import Flask, request, send_file
import torch
from models.network_swinir import SwinIR as net
from utils import utils_image as util
import io

app = Flask(__name__)

# SwinIR 모델 설정
model_path = 'model_zoo/모델 이름'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 모델 불러오기
model = net(upscale=2, in_chans=3, img_size=64, window_size=8,
            img_range=1., depths=[6,6,6,6,6,6], embed_dim=60,
            num_heads=[6,6,6,6,6,6], mlp_ratio=2, upsampler='pixelshuffledirect',
            resi_connection='1conv').to(device)
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()

@app.route('/upscale', methods=['POST'])
def upscale_image():
    # 이미지 파일 받기
    file = request.files['image']

    # PIL 형식으로 변환
    img_L = util.imread_uint(file.stream, n_channels=3)
    img_L = util.uint2tensor4(img_L).to(device)

    # 업스케일링 실행
    with torch.no_grad():
        img_E = model(img_L)

    # 결과 이미지를 바이트로 변환
    img_E = util.tensor2uint(img_E)
    img_E_pil = util.tensor2img(img_E)
    byte_io = io.BytesIO()
    img_E_pil.save(byte_io, 'PNG')
    byte_io.seek(0)

    # 업스케일된 이미지 반환
    return send_file(byte_io, mimetype='image/png')

if __name__ == '__main__':
    app.run(debug=True)

위 코드는 Flask 웹 서버를 설정하고 /upscale 엔드포인트를 통해 이미지를 받아 SwinIR 모델로 업스케일링한 후 결과를 반환합니다. 모델 경로를 설정하고 필요한 SwinIR 관련 파일들이 올바른 위치에 있어야 합니다.

클라이언트 웹사이트

클라이언트 웹사이트 구현은 HTML과 JavaScript를 사용하여 사용자가 이미지를 업로드하고 업스케일링된 결과를 받을 수 있는 인터페이스를 만듭니다.

<!DOCTYPE html>
<html lang="ko">
<head>
    <meta charset="UTF-8">
    <title>Image Upsacle</title>
</head>
<body>
    <input type="file" id="imageInput">
    <button onclick="uploadImage()">업로드</button>
    <script src="https://cdn.jsdelivr.net/npm/axios/dist/axios.min.js"></script>
    <script>
        function uploadImage() {
            const input = document.getElementById('imageInput');
            const data = new FormData();
            data.append('image', input.files[0]);
            
            axios.post('http://localhost:5000/upscale', data, {
                responseType: 'blob'
            })
            .then(response => {
                const url = URL.createObjectURL(new Blob([response.data]));
                const link = document.createElement('a');
                link.href = url;
                link.setAttribute('download', 'upscaled.png');
                document.body.appendChild(link);
                link.click();
            })
            .catch(error => console.error(error));
        }
    </script>
</body>
</html>

결론

이번 프로젝트를 통해 간단한 SwinIR 이미지 업스케일링 서비스를 구축할 첫걸음을 걸었습니다. 위와 같은 예시를 활용해서 나만의 업스케일 서비스를 만들거나 실제 서비스 실현까지 해보세요. 감사합니다.

python