Coding 공부/Python

YOLOv3와 ChatGPT를 이용한 이미지 객체 인식 및 영양 정보 분석

CBJH 2024. 7. 23.
728x90
반응형

이 포스트에서는 YOLOv3와 ChatGPT를 이용해 이미지 객체를 인식하고, 인식된 객체의 영양 정보를 분석하는 방법을 소개합니다. 이를 위해 YOLOv3 모델을 사용하여 이미지를 분석하고, ChatGPT API를 통해 영양 정보를 받아 JSON 파일로 저장합니다.

import cv2
import numpy as np
import json
import os
import aiohttp
import asyncio
from pathlib import Path
from dotenv import load_dotenv  # python-dotenv 불러오기
from collections import Counter

# .env 파일에서 환경 변수 로드
load_dotenv()

# YOLO 설정 파일과 가중치 파일 경로
yolo_config = 'path/to/yolov3.cfg'
yolo_weights = 'path/to/yolov3.weights'
yolo_names = 'path/to/coco.names'

# YOLO 네트워크 불러오기
net = cv2.dnn.readNet(yolo_weights, yolo_config)
layer_names = net.getLayerNames()
output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]

# coco.names 파일 읽기
with open(yolo_names, 'r') as f:
    classes = [line.strip() for line in f.readlines()]

# 이미지 불러오기
image_path = 'img/test01.jpg'
image = cv2.imread(image_path)
height, width, channels = image.shape

# YOLO를 위한 이미지 전처리
blob = cv2.dnn.blobFromImage(image, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
net.setInput(blob)
outs = net.forward(output_layers)

# 결과 분석
class_ids = []
confidences = []
boxes = []
for out in outs:
    for detection in out:
        scores = detection[5:]
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        if confidence > 0.5:
            center_x = int(detection[0] * width)
            center_y = int(detection[1] * height)
            w = int(detection[2] * width)
            h = int(detection[3] * height)
            x = int(center_x - w / 2)
            y = int(center_y - h / 2)
            boxes.append([x, y, w, h])
            confidences.append(float(confidence))
            class_ids.append(class_id)

# NMS 적용
indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)

# 객체 리스트 생성 및 개수 계산
objects = []
for i in range(len(boxes)):
    if i in indexes:
        x, y, w, h = boxes[i]
        label = str(classes[class_ids[i]])
        objects.append(label)
        cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)
        cv2.putText(image, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 객체 개수 계산
object_counts = Counter(objects)

# 객체와 개수 리스트를 한국어로 변환
object_counts_korean = [f"{label} {count}개" for label, count in object_counts.items()]

# 콘솔에 출력
print("Detected objects:", objects)
print("Object counts:", object_counts)
print("Object counts (Korean):", object_counts_korean)

# ChatGPT API 호출 비동기 함수
async def fetch_nutrition_info(prompt):
    api_url = 'https://api.openai.com/v1/chat/completions'
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
    }
    data = {
        'model': 'gpt-4',
        'messages': [
            {'role': 'system', 'content': 'You are a helpful assistant who can understand and respond in Korean.'},
            {'role': 'user', 'content': prompt}
        ],
        'max_tokens': 300
    }

    async with aiohttp.ClientSession() as session:
        async with session.post(api_url, headers=headers, json=data) as response:
            return await response.json()

# 비동기 메인 함수
async def main():
    prompt = f"다음 중 식재료만을 구별하여 분석하고 영양 정보를 제공해 주세요: {', '.join(object_counts_korean)}"
    nutrition_info = await fetch_nutrition_info(prompt)

    # 콘솔에 ChatGPT 응답 출력
    print("ChatGPT Response:", nutrition_info)

    # JSON 저장
    result_path = Path('result/nutrition_info.json')
    result_path.parent.mkdir(parents=True, exist_ok=True)
    with open(result_path, 'w') as f:
        json.dump(nutrition_info, f, indent=4)

    print(f"Nutrition information saved to {result_path}")

    # 객체 인식 결과 이미지를 저장
    result_image_path = Path('result/detected_objects.jpg')
    result_image_path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(result_image_path), image)

    print(f"Detected objects image saved to {result_image_path}")

# 비동기 함수 실행
if __name__ == '__main__':
    asyncio.run(main())

detection 배열의 구조

detection 배열은 객체 검출 모델의 출력 값입니다. 보통 YOLO 모델을 사용할 때 각 검출마다 다음과 같은 정보를 포함하는 배열을 제공합니다:

  1. detection[0]: 검출된 객체의 중심 x 좌표 (normalized, 0~1)
  2. detection[1]: 검출된 객체의 중심 y 좌표 (normalized, 0~1)
  3. detection[2]: 검출된 객체의 너비 (normalized, 0~1)
  4. detection[3]: 검출된 객체의 높이 (normalized, 0~1)
  5. detection[4]: 객체 검출의 신뢰도
  6. detection[5:]: 각 클래스에 대한 신뢰도 (클래스 점수)

 

 

주요 기능 및 함수 설명

1. 환경 변수 로드

.env 파일에서 API 키와 같은 중요한 환경 변수를 로드합니다. 이는 보안과 편의성을 위해 사용됩니다.
 

2. YOLO 설정 파일과 가중치 파일 로드

YOLOv3 모델의 설정 파일(.cfg)과 가중치 파일(.weights), 그리고 클래스 이름 파일(.names)을 로드합니다.
다양한 사람이 학습시키고 업로드한 YOLOv3 모델이 많지만, 가장 대중적이고 신뢰할 수 있는 Darknet GitHub 리포지터리에서 설정 파일과 가중치 파일을 다운로드 했습니다.
 

3. 이미지 전처리

이미지를 YOLO 모델에 맞게 전처리합니다. 이미지를 블롭(blob) 형식으로 변환하여 YOLO 모델의 입력으로 설정합니다. 이 과정을 통해 이미지를 정규화하고 크기를 조정합니다.
 

4. 결과 분석 및 객체 인식

YOLO의 출력 결과를 분석하여 객체의 클래스와 위치를 식별합니다. confidence가 0.5 이상인 경우에만 객체로 인식합니다.
 

5. NMS(Non-Maximum Suppression) 적용

NMS를 적용하여 중복되는 박스를 제거합니다. 이를 통해 보다 정확한 객체 인식을 할 수 있습니다.
 

6. 객체 리스트 생성 및 개수 계산

 

인식된 객체를 리스트로 저장하고, Counter 클래스를 이용해 객체의 개수를 계산합니다.

 

7. ChatGPT API 호출

ChatGPT API를 비동기적으로 호출하여 객체 인식 결과를 기반으로 영양 정보를 요청합니다.

 

8. 비동기 메인 함수

메인 함수에서는 객체 인식 결과를 ChatGPT에 보내고, 응답받은 영양 정보를 JSON 파일로 저장합니다. 또한, 객체 인식 결과 이미지를 저장합니다.

 

이미지 전처리와 NMS에 대한 설명

이미지 전처리는 YOLO 모델의 성능을 극대화하기 위해 필수적인 과정입니다. 이미지를 블롭(blob) 형식으로 변환하여 모델의 입력 크기에 맞추고, 정규화하여 모델이 쉽게 인식할 수 있도록 합니다. 블롭(blob) 형식은 이미지 데이터를 단일 배열로 변환하여 신경망에 입력할 수 있도록 해줍니다.

NMS(Non-Maximum Suppression)는 객체 인식에서 매우 중요한 과정입니다. 여러 박스가 겹치는 경우, 가장 신뢰도가 높은 박스만을 남기고 나머지를 제거합니다. 이를 통해 중복 인식을 방지하고 정확도를 높일 수 있습니다. 예를 들어, 같은 객체를 여러 번 인식할 경우, NMS는 가장 신뢰도 높은 예측을 선택하여 중복을 제거합니다.

객체의 개수 카운트 방법

collections 모듈의 Counter 클래스를 사용하여 객체의 개수를 카운트합니다. 객체 리스트를 생성한 후, Counter를 이용해 각 객체의 빈도를 계산합니다. 이렇게 하면 도넛이 몇 개인지, 컵이 몇 개인지 등을 쉽게 파악할 수 있습니다. 예를 들어, 객체 리스트에 "도넛"이 4번 등장하면, Counter는 이를 자동으로 세어줍니다.

라이브러리 설명 및 사용 이유

1. import aiohttp

aiohttp는 비동기 HTTP 클라이언트/서버 라이브러리입니다. 비동기적으로 HTTP 요청을 보내기 위해 사용됩니다. 이 프로젝트에서는 ChatGPT API를 비동기적으로 호출하여 응답을 받기 위해 사용됩니다.

2. import asyncio

asyncio는 Python에서 비동기 프로그래밍을 가능하게 하는 표준 라이브러리입니다. 비동기 함수 실행 및 이벤트 루프 관리를 위해 사용됩니다. 이 프로젝트에서는 비동기 ChatGPT API 호출을 관리하기 위해 사용됩니다.

3. from pathlib import Path

pathlib는 파일 시스템 경로 조작을 위한 모듈입니다. 경로를 객체 지향적으로 처리할 수 있도록 도와줍니다. 이 프로젝트에서는 결과 파일을 저장할 경로를 관리하기 위해 사용됩니다.

 

디렉토리

결과물

 

입력 사진
출력 사진
콘솔창 출력 결과

이 코드의 한계점

1. 기능 구현만을 위해 작성한 코드라 유연성이 떨어집니다. img 폴더 내의 test01.jpg 파일만 분석 가능합니다.

2. YOLOv3에서 인식할 수 있는 객체는 총 80가지입니다. 따라서 인식할 수 있는 식재료의 가지수가 많지 않아 도넛을 인식하는 코드를 작성했습니다.

 

리스트 : 사람, 자전거, 자동차, 오토바이, 비행기, 버스, 기차, 트럭, 보트, 신호등, 소화전, 정지 표지판, 주차 요금기, 벤치, 새, 고양이, 개, 말, 양, 소, 코끼리, 곰, 얼룩말, 기린, 배낭, 우산, 핸드백, 넥타이, 여행 가방, 프리스비, 스키, 스노보드, 운동 공, 연, 야구 방망이, 야구 글러브, 스케이트보드, 서핑보드, 테니스 라켓, 병, 와인 잔, 컵, 포크, 칼, 숟가락, 그릇, 바나나, 사과, 샌드위치, 오렌지, 브로콜리, 당근, 핫도그, 피자, 도넛, 케이크, 의자, 소파, 화분, 침대, 식탁, 변기, 텔레비전 모니터, 노트북, 마우스, 리모컨, 키보드, 휴대전화, 전자레인지, 오븐, 토스터, 싱크대, 냉장고, 책, 시계, 꽃병, 가위, 테디 베어, 헤어 드라이어, 칫솔

 

3. 정확한 영양 정보를 알기 힘듭니다. 대략적인 정보는 알 수 있지만 신뢰도를 높이기 위해선 음식의 양을 정확히 구분할 수 있는 알고리즘이 필요합니다.

댓글