개발/Python

이미지에서 사물을 분석해주는 모델(VGG16)을 사용해 API로 만들어보자

joon95 2024. 12. 11. 18:13
반응형

언젠가 AI를 해봐야지 생각하면서.. 미루고 있었다.

perplexity(GPT) 와 채팅을 하면서 문득 사전에 학습된 모델을 알게되어

VGG16 모델을 통해 이미지를 분석하는 테스트를 해보았다.

 

VGG16

먼저 VGG16 이란

- 키보드, 동물, 연필 등의 정보를 1000개의 Class로 분류 가능

- 3x3 커널을 사용

- 모델 이미지 입력 크기 224x224

 

GPU (CUDA) 세팅

CPU 를 사용했더니 너무 느려서 노트북에 내장되어 있는 그래픽카드를 셋업했다.

필자의 그래픽 카드정보는 NVIDIA GeForce GTX 1650 Ti with Max-Q Design 이며,

NVIDIA 사이트에서 Cuda를 지원하는지 체크하여 가능한 것임을 알았다.

 

TensorFlow GPU 지원-소프트웨어 요구사항

다음 NVIDIA® 소프트웨어가 시스템에 설치되어 있어야 합니다.

 

세팅 시작

1. 그래픽카드 드라이버 업데이트

업무용PC라 한번도 안했던 것이라 바로 고.

 

2. CUDA 설치

CUDA ToolKit 버전 11.2

 

CUDA Toolkit 11.2 Downloads

Get CUDA Toolkit 11.2 for Linux and Windows.

developer.nvidia.com

현재 최신버전 12.6 여서 설치했더니, 에러가 나서 알아보니

텐서플로우에서 지원하는 안정된 버전을 쓰라고 한다.

 

3. cuDNN 설치

cuDNN 8.1 (cuda 11.2 지원)

NVIDIA 사이트에서 cuDNN 메뉴에 들어가 로그인을 해야만 다운로드가 가능하였다.

 

4. 라이브러리 다운로드

pip install tensorflow-gpu --user

 

4. 실제 GPU 사용하는지 Test Code 검증

import tensorflow as tf

# GPU 사용 가능 여부 확인
print("GPU 사용 가능 여부:", tf.test.is_built_with_cuda())

# 사용 가능한 GPU 목록 출력
print("사용 가능한 GPU 목록:", tf.config.list_physical_devices('GPU'))

# GPU 메모리 할당 테스트
if tf.test.is_built_with_cuda():
    # 간단한 GPU 연산 수행
    with tf.device('/GPU:0'):
        a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
        b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
        c = tf.matmul(a, b)
        print("GPU 연산 결과:", c)
else:
    print("GPU를 사용할 수 없습니다.")

# TensorFlow 버전 출력
print("TensorFlow 버전:", tf.__version__)

# CUDA 버전 출력 (사용 가능한 경우)
if tf.test.is_built_with_cuda():
    print("CUDA 버전:", tf.sysconfig.get_build_info()["cuda_version"])

위 파이선 코드를 실행하면 작업관리자에서 GPU 메모리 사용량이 올라간다.

 

 

VGG16 활용 코드

from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image
import numpy as np

# VGG16 모델 로드 (완전 연결 층 포함)
model = VGG16(weights='imagenet')

def identify_objects(img_path, top_n=5):
    # 이미지 로드 및 전처리
    img = image.load_img(img_path, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)

    # 예측 수행
    preds = model.predict(x)
    
    # 예측 결과 디코딩
    results = decode_predictions(preds, top=top_n)[0]

    # 결과 출력
    print(f"이미지 '{img_path}'에서 식별된 객체들:")
    for i, (imagenet_id, label, score) in enumerate(results):
        print(f"{i+1}. {label}: {score:.2%}")

# 사용 예시
image_path = 'C:\\ML\\tf\\bg.jpg'
identify_objects(image_path)

이제 위 소스를 저장하고 실행하면 C:\ML\tf\bg.jpg 파일을 읽어 나오게된다.

이미지 'C:\azureML\tf\bg1.jpeg'에서 식별된 객체들:
1. patio: 55.62%
2. park_bench: 15.68%
3. shopping_cart: 6.67%
4. rocking_chair: 3.07%
5. shopping_basket: 1.81%

 

API로 만들기(Flask)

이제 편하게 사용하기 위해 API로 만들어보자.

from flask import Flask, request, jsonify
from werkzeug.utils import secure_filename
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image
import numpy as np
import os

app = Flask(__name__)

# VGG16 모델 로드
model = VGG16(weights='imagenet')

# 업로드된 파일을 저장할 디렉토리 설정
UPLOAD_FOLDER = 'uploads'
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

def analyze_single_image(filepath):
    img = image.load_img(filepath, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)

    preds = model.predict(x)
    results = decode_predictions(preds, top=5)[0]

    return [
        {'label': label, 'score': float(score)}
        for (imagenet_id, label, score) in results
    ]

@app.route('/analyze', methods=['POST'])
def analyze_images():
    if 'files' not in request.files:
        return jsonify({'error': '파일이 없습니다'}), 400
    
    files = request.files.getlist('files')
    if not files or files[0].filename == '':
        return jsonify({'error': '선택된 파일이 없습니다'}), 400
    
    results = {}
    for file in files:
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        results[filename] = analyze_single_image(filepath)
        
        # 임시 파일 삭제
        os.remove(filepath)

    return jsonify(results)

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

POST http://localhost:5000/analyze

files 로 사진을 업로드하면 된다.

응답값:

{
    "bg.jpg": [
        {
            "label": "desk",
            "score": 0.27180394530296326
        },
        {
            "label": "guillotine",
            "score": 0.2136017382144928
        },
        {
            "label": "file",
            "score": 0.06649696826934814
        },
        {
            "label": "apiary",
            "score": 0.04862062633037567
        },
        {
            "label": "crate",
            "score": 0.026243386790156364
        }
    ],
    "bg2.png": [
        {
            "label": "necklace",
            "score": 0.2780897617340088
        },
        {
            "label": "hook",
            "score": 0.07838109880685806
        },
        {
            "label": "Christmas_stocking",
            "score": 0.042882807552814484
        },
        {
            "label": "letter_opener",
            "score": 0.03623732924461365
        },
        {
            "label": "safety_pin",
            "score": 0.03522750362753868
        }
    ],
    "bg3.png": [
        {
            "label": "paintbrush",
            "score": 0.7162522077560425
        },
        {
            "label": "rubber_eraser",
            "score": 0.04436235502362251
        },
        {
            "label": "ballpoint",
            "score": 0.024050327017903328
        },
        {
            "label": "face_powder",
            "score": 0.023242339491844177
        },
        {
            "label": "wooden_spoon",
            "score": 0.016569746658205986
        }
    ]
}

 

반응형