개인 공부/인공지능

[kaggle/Flask] 견종 분류 모델 웹에 구현하기

KimNang 2023. 10. 4. 18:26

데이터셋 정보

제목 : Dog Breed Image Classification Dataset

데이터셋 링크

https://www.kaggle.com/datasets/wuttipats/dog-breed-image-classification-dataset/data

 

Dog Breed Image Classification Dataset

A Collection of Dog Images Categorized by Breed

www.kaggle.com

데이터셋 설명

- 이 데이터셋에는 6종류의 개 이미지가 있음

- 파일의 형식은 JPEG이고 이름은 사진 속 개의 품종을 알려줌

- 사진은 견종에 따라 다른 폴더에 저장되어있음

 

견종 이미지 분류 모델 웹에 구현

시작하기 전

- 원래는 이전 글의 견종 이미지 분류 모델을 이용해 Flask로 웹에 구현해보려고 했음

https://kne-coding.tistory.com/265

 

[kaggle/TensorFlow] 견종 이미지 분류 모델 구현 - 2

데이터셋 정보 제목 : Dog Breed Image Classification Dataset 데이터셋 링크 https://www.kaggle.com/datasets/wuttipats/dog-breed-image-classification-dataset/data Dog Breed Image Classification Dataset A Collection of Dog Images Categorized by B

kne-coding.tistory.com

- PyCharm을 이용해 가상환경을 설정하고 코드는 작성해보았으나, tensorflow와 keras 버전 문제인지 계속 오류가 떴음

- 며칠동안 여러 방법을 시도해보다 도저히 해결이 안돼서 결국 PyTorch로 모델을 다시 학습, 저장 후 진행함

 

PyTorch 모델 학습, 저장
Dog_Breed_Classification_PyTorch

 

예측 수행한 이미지

 

Flask 이용해서 웹에 서빙

- 아래의 사진은 해당 프로젝트의 디렉터리 구조

main.py

- 학습된 ResNet 모델을 로드하여 이미지를 입력받아 예측을 수행함

- 예측 수행 후 테스트로 출력하여 결과를 알려줌 (결과 출력 테스트용)

from torchvision import models
import time
import copy
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch
import torchvision.transforms as transforms
import os

from PIL import Image

# 미리 학습된 ResNet 모델 로드
resnet50 = torch.load('./resnet50.pt')
resnet50.eval()

class_names = {0:'Bedlington_terrier',
 1:'Bernese_mountain_dog',
 2:'Dandie_Dinmont',
 3:'Gordon_setter',
 4:'Ibizan_hound',
 5:'Norwegian_elkhound'}

def predict_image(image_path):
    # 이미지를 모델의 입력 크기에 맞게 조정
    preprocess_image = transforms.Compose([
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    image = Image.open(image_path) # 이미지 불러오기
    image = preprocess_image(image)
    image = image.unsqueeze(0) # 배치 차원 추가 (이미지를 하나의 배치로 처리)



    # 모델 및 입력 데이터를 모델의 디바이스로 이동
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    resnet50.to(device)
    image = image.to(device)

    # 모델 예측
    with torch.no_grad():
        output = resnet50(image)

    # 확률 분포에서 가장 높은 확률을 가지는 클래스를 찾습니다.
    _, predicted_class = output.max(1)

    return class_names.get(predicted_class.item(), 'Unknown')


# # 이미지 예측
# for i in range(1, 9):
#     image_path = './image/image' + str(i) + '.jpg'
#     predicted_class = predict_image(image_path)
#
#     print(f'예측된 클래스: {predicted_class}')

 

index.html

- 이미지 업로드를 위한 웹 페이지

- '파일 선택'을 통해 이미지 업로드 후 파일명을 확인하고 '견종 분류 예측하기!'버튼을 클릭함

- 클릭시 result.html 페이지로 이동함

<!DOCTYPE html>
<html>
    <head>
    <title>견종 분류하기</title>
    </head>
    <body>
        <div class="container">
        <h1>견종 분류하기</h1>
        <h2>이미지 업로드</h2>
        {% if filename %}
            <div>
                <img src="{{filename}}">
            </div>
    
            <div>
    
                <p>{{label}} {{probability}}</p>
    
            </div>
        {% endif %}
        <form method="post" action="/result" enctype="multipart/form-data">
            <input type="file" name="file">
            <input type="submit" value="견종 분류 예측하기!">
        </form>
        </div>
    </body>
</html>

 

result.html

- 사용자가 업로드한 이미지와 함께 모델 예측 결과가 표시됨

- 아래의 '다른 사진 예측하기' 버튼을 통해 이미지 업로드 페이지로 이동할 수 있음

<!DOCTYPE html>
<html>
    <head>
        <title>이미지 예측 결과</title>
    </head>
    <body>
        <h1>이미지 예측 결과</h1>
    
        <!-- 업로드된 이미지 표시 -->
        <h2>개 이미지:</h2>
        <img src="{{ uploaded_image }}" alt="Uploaded Image" style="max-width: 400px;">
    
        <!-- 모델 예측 결과 표시 -->
        <h2>예측한 견종:</h2>
        <p>{{ predicted_class }}</p>
    
        <!-- 이미지 업로드 페이지로 이동하는 링크 -->
        <p><a href="/">다른 사진 예측하기</a></p>
    </body>
</html>

 

app.py

- 기본 페이지를 index.html로 설정함

- 해당 파일의 <form>태그를 보면 action=result 이므로 파일을 업로드 하면 @app.route('/result',~) 부분이 실행됨

- 업로드 된 이미지는 파일 경로를 생성하여 static - IMG 폴더에 저장됨

- predicted_class에 이미지 견종 예측 결과를 저장함

- result.html로 해당 결과와 파일 이미지 경로를 포함하여 렌더링함

from flask import Flask, render_template, request, jsonify
from main import predict_image
import torch
import torchvision.models as models
import joblib
import os
import cv2
import numpy as np

app = Flask(__name__)
IMG_FOLDER = os.path.join('static', 'IMG')
app.config['UPLOAD_FOLDER'] = IMG_FOLDER

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/result', methods=['POST'])
def upload():
    uploaded_file = request.files['file']
    if uploaded_file.filename != '':
        file_path = f"./static/IMG/{uploaded_file.filename}"
        print(file_path)
        uploaded_file.save(file_path)
        predicted_class = predict_image(file_path)
        up_img = os.path.join(app.config['UPLOAD_FOLDER'], uploaded_file.filename)



        # result.html을 렌더링하고 업로드된 이미지와 예측 결과 전달
        return render_template('result.html', uploaded_image=up_img, predicted_class=predicted_class)
    else:
        return jsonify({'error': 'No file selected'})


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

 

결과 확인하기
  • 접속시 화면

  • 파일 선택했을때의 상태

 

  • 이미지 예측 결과

 

  • 업로드 이미지 폴더에 저장