내일배움캠프 - My Little Hero 개발일지

Dongwoo Kim·2022년 5월 25일

AI 개발일지 내일배움캠프 스파르타코딩클럽

내일배움캠프 - 프로젝트

목록 보기

9/12

스파르타 코딩클럽

내일배움캠프 AI 웹개발자양성과정 2회차

사물인식 팀 프로젝트 개발일지

0. 프로젝트 정보

1) 프로젝트 명

My Little Hero - 나의 닮은 마블 캐릭터 찾기

2) 기간

2022.05.18.~25.

3) 프로젝트 간단설명

머신러닝 모델을 이용하여 사용자 이미지와 비슷한 마블 캐릭터를 찾아줌

4) 역할 분담

공통 : 각자 CNN모델 1개씩 학습시키기
최민기 : VGG16
김진수 : MobileNetV2
김동우 : InceptionV3
이윤지 : ResNet50
UI + 기능
최민기 : 로그인 / 회원가입
김진수 : 헤더 디자인 + 이미지 업로드 (모달창)
김동우 : 메인페이지 + 결과 보여주기
이윤지 : 히스토리페이지 + 결과 보여주기

1. 데이터셋 & 머신러닝 모델

1) 데이터셋

총 8개의 classes, 3035 files로 구성된 마블 캐릭터 이미지 데이터셋
https://www.kaggle.com/datasets/hchen13/marvel-heroes

2) MobileNetV2

팀원별 CNN모델 학습 결과 VGG16을 제외한 모델들에서 validation accuracy가 0.65정도로 비슷한 성능을 냄. 따라서 그 중에서 가볍고 학습속도가 빠른 MobileNetV2를 이용

MobileNetV2 학습 그래프, epoch 20 이후부터는 overfitting이 일어나서 그전까지만 학습한 모델을 이용했다.

3) Tools

keras, flask, mongoDB, javascript, html, css

2. Frontend

0) 공통 기능

: 로그인한상태가 아니면 (localstorage에 token값이 없으면) 로그인 화면으로 강제 이동

$(document).ready(function () {
  if (!localStorage.getItem("token")) {
    location.href = "/templates/log_in.html";
  }
});

1) 메인페이지 UI 기능

: 각 종 버튼 기능 및 스크롤 이동 기능

// 업로드 버튼 함수
function upload_user_img() {
  save_result();

  $("#select_view").css("display", "none");
  $("#result_view").css("display", "block");
  $(".img_modal_box").css("display", "none");
}

// 다시하기 버튼 함수
function restart_upload() {
  $("#select_view").css("display", "block");
  $("#result_view").css("display", "none");
}

// 메인 마블 캐릭터 변경시 스크롤 이동 함수
function move_info() {
  let offset = $(".main_body").offset(); //해당 위치 반환
  $("html, body").animate({ scrollTop: offset.top - 120 }, 400);
}

2) 이미지 업로드 모달창

: 사용자가 이미지를 업로드하기위한 모달창 기능

// 이미지 업로드 모달창 띄우는 함수
function upload_img_modal() {
  $(".img_modal_box").css("display", "block");
}

// 모달창에서 업로드할 이미지 미리보기
function preview(event) {
  let reader = new FileReader();
  reader.onload = (event) => {
    let img = document.querySelector("#image_file");
    img.setAttribute("src", event.target.result);

    $('.main_img[name="user"]').attr("src", event.target.result);
  };

  reader.readAsDataURL(event.target.files[0]);
}

// 이미지 업로드 모달창 끄기 함수
function quit_img_modal(event) {
  const modal_box = document.querySelector(".img_modal_box");

  if (event.target == modal_box) {
    $(".img_modal_box").css("display", "none");
  }
}

3) 이미지 예측 결과 API

: 사용자가 업로드한 이미지를 서버로 전송하고 예측결과 받아서 화면에 띄워줌

// 메인 결과 화면을 보여주는 API
async function save_result() {
  let token = localStorage.getItem("token");
  let file = $("#modal_input_img")[0].files[0];
  let form_data = new FormData();

  form_data.append("user_img", file);
  form_data.append("token", token);

  $.ajax({
    type: "POST",
    url: "http://127.0.0.1:5000/main/result",
    data: form_data,
    cache: false,
    contentType: false,
    processData: false,
    success: function (response) {
      dummy_heros = response["results"];

      // 1위 마블 캐릭터
      $("#main_hero_img").attr("src", dummy_heros[0].hero_img);
      $("#main_hero").attr("name", dummy_heros[0].hero);
      $("#main_hero").text("1위 -" + dummy_heros[0].hero);
      $("#main_desc").text(dummy_heros[0].description);

      // 2위 마블 캐릭터
      $('.hero_sub_box[name="2"] .hero_sub_img').attr(
        "name",
        dummy_heros[1].hero
      );
      $('.hero_sub_box[name="2"] .hero_sub_img').attr(
        "src",
        dummy_heros[1].hero_img
      );
      $('.hero_sub_box[name="2"] div').text(
        dummy_heros[1].rank + "위 - " + dummy_heros[1].hero
      );

      // 3위 마블 캐릭터
      $('.hero_sub_box[name="3"] .hero_sub_img').attr(
        "name",
        dummy_heros[2].hero
      );
      $('.hero_sub_box[name="3"] .hero_sub_img').attr(
        "src",
        dummy_heros[2].hero_img
      );
      $('.hero_sub_box[name="3"] div').text(
        dummy_heros[2].rank + "위 - " + dummy_heros[2].hero
      );
    },
  });

3. Backend

1) common.py : 공통 기능

: 사용자 이메일과 저장 시간을 파일 이름으로 하여 이미지 파일을 저장

# 사용자 이미지를 저장
# user_img: 사용자 이미지 파일
# email: 사용자 이메일
# return: 저장한 경로
def save_user_img(user_img, email):
    extension = user_img.filename.split('.')[-1]
    today = datetime.now()
    mytime = today.strftime('%Y-%m-%d-%H-%M-%S')
    filename = f'{mytime}'

    save_to = f'static/images/user/{email}_{filename}.{extension}'
    user_img.save(save_to)

    return save_to

2) model.py : 머신러닝 모델 기능

: 입력받은 이미지 주소로 저장한 이미지를 불러와서 전처리 후 학습 시킨 모델로 예측, 예측 확률이 높은 순으로 정렬하여 반환

# 머신러닝 모델 기능 함수

from tensorflow.keras.preprocessing.image import load_img, img_to_array
import numpy as np

from tensorflow.keras.models import load_model

model = load_model('static/models/MbileNetV2_2_2.h5')

hero_class = [
    '블랙위도우',
    '캡틴아메리카',
    '닥터스트레인지',
    '헐크',
    '아이언맨',
    '로키',
    '스파이더맨',
    '타노스']


# 사용자 이미지로 닮은 마블 캐릭터 예측
# user_img: 사용자 이미지 주소
# return: 마블 캐릭터별 예측 확률 (확률이 높은 순으로 정렬)
def predict_hero(user_img):
    img_dir = user_img
    image = load_img(img_dir, target_size=(224, 224))
    input_arr = img_to_array(image)
    input_arr = np.array([input_arr])
    input_arr = input_arr / 255
>
    predictions = model.predict(input_arr)

    results = []
    for i in range(8):
        result = {
            'hero': hero_class[i],
            'accuracy': predictions[0][i]
        }
        results.append(result)

    results.sort(key=acc_of_result, reverse=True)

    return results


# 닮은 캐릭터 예측 결과를 정렬하기위한 key 함수
def acc_of_result(result):
    return result['accuracy']

3) main.py : 메인페이지 기능

# 사용자 이미지로 예측 결과 저장
# user_img: 사용자 이미지 파일
# user_info: 사용자 정보
# return: result - 가장 닮은 마블 캐릭터 정보 3개
def predict_img(user_img, user_info):
    # 이미지 저장하기
    user_email = user_info['email']
    dir = common.save_user_img(user_img, user_email)

    # 닮은 마블 캐릭터 예측 & 상위 3개 선정
    results = model.predict_hero(dir)[0:3]  # DB에 저장할 결과 데이터

    # DB에 저장
    for result in results:
        result['accuracy'] = result['accuracy'].item()
        result['email'] = user_email
        result['user_img'] = dir

        print(result)
        db.results.insert_one(result)

    # 클라이언트에 보내줄 결과 리스트
    view_results = []
    for i, result in enumerate(results, start=1):
        rank = i
        hero_info = common.get_hero_info(result['hero'])
        view_result = {
            'rank': rank,
            'hero': result['hero'],
            'description': hero_info['description'],
            'hero_img': hero_info['hero_img']
        }
        view_results.append(view_result)

    return view_results

4) app.py : API 기능

# 메인 결과 보여주기 API
@app.route('/main/result', methods=['POST'])
def main_result():
    token = request.form['token']
    user_img = request.files['user_img']

    # 토큰으로부터 유저 정보 가져오기
    user_info = common.get_user_from_token(token)

    # 예측 결과 가져오기
    results = main.predict_img(user_img, user_info)

    return jsonify({'results': results})