[미니프로젝트] 닮은꼴 연예인 찾기

임재규·2023년 5월 28일
0

프로젝트

목록 보기
1/11

주제

  • 사용자의 이미지 파일을 입력값으로 받아 닮은꼴 연예인을 출력값으로 보여주기
    • CNN기법으로 미리 학습시킨 연예인 범주 내에서 출력

개요

  • 프로젝트 제목 : Project Face "닮은꼴 찾기"
  • 인원 / 기간 : 6명 / 2023.05.15 ~ 2023.05.18
  • 🛠️ 사용 기술
    • Python
    • Django
    • WSL (Ubuntu-22.04)
    • GitHub

팀프로젝트에서 수행한 작업

이미지 크롤링

  • 모델 학습에 사용될 이미지 크롤링 코드
import os
import requests
from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.action_chains import ActionChains
import time
from PIL import Image
import io

# 크롬드라이버 위치 설정
DRIVER_PATH = '/usr/local/bin/chromedriver'

# 크롬 드라이버 옵션 설정
options = Options()
options.add_argument('--headless')  # 창 안띄우기
options.add_argument('--no-sandbox')  # 리눅스 환경에서 필요한 옵션
options.add_argument('--disable-dev-shm-usage')  # 리눅스 환경에서 필요한 옵션

# 검색어와 폴더명 설정 (원하는 변수로 바꿔서 설정)
search_names = ['경리 얼굴',
'조유리 얼굴',
'해원 얼굴',
'레이 얼굴',
'홍은채 얼굴',
'가을 얼굴',
'우기 얼굴',
'태연 얼굴',
'김숙 얼굴',
'송은이 얼굴']
folder_names = ['경리 얼굴',
'조유리 얼굴',
'해원 얼굴',
'레이 얼굴',
'홍은채 얼굴',
'가을 얼굴',
'우기 얼굴',
'태연 얼굴',
'김숙 얼굴',
'송은이 얼굴']

for search_name, folder_name in zip(search_names, folder_names):
    # 검색어를 이용한 구글 이미지 검색 url
    url = f'https://www.google.com/search?q={search_name}&source=lnms&tbm=isch'
    # 크롬 드라이버 실행
    service = Service(DRIVER_PATH)
    driver = webdriver.Chrome(service=service, options=options)
    # url 접속
    driver.get(url)
    # 페이지 로드를 위한 대기 시간
    time.sleep(2)
    # 이미지 로딩을 위한 스크롤 다운
    last_height = driver.execute_script("return document.body.scrollHeight")
    while True:
        driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
        time.sleep(1)
        new_height = driver.execute_script("return document.body.scrollHeight")
        if new_height == last_height:
            break
        last_height = new_height
    # 이미지 링크 추출
    soup = BeautifulSoup(driver.page_source, 'html.parser')
    img_tags = soup.find_all('img')
    urls = []
    for img in img_tags:
        try:
            url = img['src']
            if 'http' in url:
                urls.append(url)
        except:
            pass
    # 이미지 다운로드
    os.makedirs(f'./{folder_name}', exist_ok=True)
    count = 0
    for url in urls:
        try:
            response = requests.get(url, stream=True)
            image = Image.open(io.BytesIO(response.content))
            width, height = image.size
            # 이미지 추출 과정에서 너무 작은 사이즈가 추출되는 경우를 제외하기 위해 조건문 사용
            if width >= 30 and height >= 30:
                file_name = f'./{folder_name}/{count}.jpg'
                with open(file_name, 'wb') as out_file:
                    out_file.write(response.content)
                print(f'{file_name} saved')
                count += 1
                if count == 200:
                    break
        except:
            pass
    # 크롬 드라이버 종료
    driver.quit()

장고로 로컬에서 구현해보기

  • 학습된 모델을 가지고 프로젝트 로컬에서 구현해보기
    • 이미지를 입력값으로 주면 학습시킨 데이터를 기반으로 닮은꼴 연예인을 출력하는지 확인

mini_project 파일 생성

mkdir mini_project && cd mini_project
# 장고 프로젝트 생성
django-admin startproject config .
# 장고 앱 생성
django-admin startapp star
  • 경로 지정
# config - settings.py
INSTALLED_APPS = [
    ...
       
    'star',       # 추가
]
# config - urls.py
from django.contrib import admin
from django.urls import path, include


urlpatterns = [
    path('admin/', admin.site.urls),
    path('', include('star.urls')), 
]
# star - urls.py
from django.urls import path, include
from . import views
from star.views import classify_image
from .views import rank

app_name = 'star'

urlpatterns = [
    path('', classify_image, name='classify_image'),
    path('ranking/', rank, name='rank'),
    path('sub_menu_1/', views.sub_menu_1, name='sub_menu_1'),
    #[주소]/sub_menu_1로 가면 sub_menu_1함수로!
]
  • 이미지 업로드 폼 생성
# star - forms.py
from django import forms
class ImageUploadForm(forms.Form):    
    image = forms.ImageField()
  • Classify_image - 입력된 이미지를 모델을 통해 분류하는 기능을 구현
  • Post - 분류 결과를 저장하기 위한 데이터베이스 기능을 구현
# star - models.py
import tensorflow as tf
from django.db import models
from PIL import Image
import numpy as np
import os

class ImageClassifier:
    def __init__(self, model_path):
        self.model = tf.keras.models.load_model(model_path)
    def classify_image(self, image):
        class_label=[   '강해린',
                        '김채원',
                        '뉴진스 하니',
                        '안유진',
                        '아린',
                        '장원영',
                        '카즈하',
                        '차은우',
                        '방탄소년단 뷔',
                        '정준하',
                        '정형돈',
                        '서강준',
                        '카리나',
                        '마동석',
                        '박보검',
                        '손석구',
                        '아이브 가을',
                        '이광수',
                        '침착맨',
                        '르세라핌 사쿠라',]
        img = Image.open(image)
        img = img.resize((224, 224))  # 모델에 맞는 이미지 크기로 조정
        img = np.array(img) / 255.0  # 이미지 정규화
        img = np.expand_dims(img, axis=0)  # 배치 차원 추가
        predictions = self.model.predict(img)
        pred_class = np.argmax(predictions, axis=1)
        class_index = np.argmax(predictions[0])  # 가장 높은 확률을 가진 클래스 인덱스
        confidence = predictions[0][class_index]  # 분류 확률
        result = class_label[pred_class[0]]
        result = {
            'class_label': result,
            'confidence': confidence,
        }
        return result

class Post(models.Model):
    confidence = models.FloatField()
    result = models.CharField(max_length=100)

    def __str__(self):
        return self.result
  • classify_image - 이미지를 분류하고 저장하는 기능
  • rank - 데이터베이스에 저장된 분류 결과를 조회하여 ranking 템플릿에 랜더링하고 결과 전달
# star - view.py
import os
from PIL import Image
from .models import Post
from django.shortcuts import render
from django.http import HttpResponse
from .forms import ImageUploadForm
from .models import ImageClassifier,Post
import numpy as np
import glob

model_path = 'star/team3_new.h5'  # 실제 모델 파일의 경로로 수정
# 이미지 분류기 초기화
classifier = ImageClassifier(model_path)

def classify_image(request):
    # 이미지 분류 요청 처리
    if request.method == 'POST':
        form = ImageUploadForm(request.POST, request.FILES)
        if form.is_valid():
            image = form.cleaned_data['image']
            
            # 이미지 분류
            result = classifier.classify_image(image)
            
            # 분류 결과 저장
            Post.objects.create(confidence=result['confidence'], result=result['class_label'])
            
            # 결과를 result.html 템플릿에 전달하여 렌더링
            return render(request, 'result.html', {'result': result})
    else:
        form = ImageUploadForm()
    
    # 이미지 분류 페이지 렌더링
    return render(request, 'index.html', {'form': form})

def rank(request):
    # 모든 분류 결과 조회
    results = Post.objects.all()
    
    # 결과를 ranking.html 템플릿에 전달하여 렌더링
    return render(request, 'ranking.html', {'results': results})

def sub_menu_1(request):
    # sub_menu_1.html 템플릿 렌더링
    return render(request, 'sub_menu_1.html')

웹페이지 구현

  • 초기화면
  • 시작버튼 누른 후 화면
  • 결과 화면

느낀점

팀 프로젝트에서 내가 한 부분은 모델에 사용될 이미지를 크롤링 코드를 구현한 것과 이미지를 넣었을 때, 장고에서 우리가 만든 모델이 잘 돌아갈 수 있도록 만드는 것이었다.
크롤링할때, 검색어를 통해 이미지를 크롤링하다보니, 그 사람과 맞지 않는 얼굴도 많았고, 얼굴 사진이 아닌 전신 사진 등이 데이터로 담겼다. 이를 전처리하는 과정에서 잘되지 않아, 모델에 성능이 떨어졌던게 아쉽다.
HTML에 관련된 것은 처음 다뤄봐서 웹페이지를 깔끔하게 하지 못한 것이 아쉽다.
설계 당시, 데이터 분류 결과들을 데이터베이스에 모아 닮은꼴 연예인 별로 가장 닮은 확률 값을 가지고 랭킹페이지를 구현하려고 했는데, 구현하지 못했다.
(연예인 별이 아닌 가장 높은 값을 구현하는 데는 성공했지만, 팀 프로젝트 과정에서 이런 페이지가 필요 있을까?라는 생각에 제외됐다.)
팀 프로젝트 당시에 GPU가 없는 PC를 통해 학습시키는 과정을 거쳐 시간이 많이 걸렸는데, 자원이 충분했다면 더 많은 이미지를 통해 출력되는 연예인들도 다양성을 더 높여 만들어보고 싶다.

profile
공부 기록

0개의 댓글