[지능] Vision AI 기반 컨베이어 벨트 객체 인식 딥러닝 모델 최적화 2일차

2024.12.04

 

2일차에는 1일차 때 라벨링 한 샘플 데이터를 YOLOv6-n 모델을 이용하여 학습시킨 뒤 prediction을 해보았다.

 

 

 

샘플 데이터의 화질이 매우 좋지 않아서 그런지 확실히 객체를 정확하게 탐지하지 못하는 것 같다.

 

파이썬 코드를 작성하여 수집한 이미지 데이터에 학습하여 생성한 YOLOv6 모델을 적용시켜 객체 탐지 박스가 생성된 이미지를 생성하여 저장하도록 하였다.

 

import cv2
import requests
from requests.auth import HTTPBasicAuth
import os
import glob
import json
import random
from collections import Counter

# ------------------------------
# 1. 클래스별 색상 매핑 및 동적 색상 생성 함수 정의
# ------------------------------
# 초기 클래스별 색상 매핑 (원하는 색상으로 변경 가능)
class_colors = {
    'RASPBERRY PICO': (255, 0, 0),    # 파란색
    'USB': (0, 255, 0),               # 초록색
    'OSCILLATOR': (0, 0, 255),        # 빨간색
    'HOLE': (255, 255, 0),            # 시안색
    'CHIPSET': (255, 0, 255),         # 자홍색
    'BOOTSEL': (0, 255, 255)          # 노란색
}

def get_class_color(class_name):
    """
    클래스 이름에 따라 색상을 반환합니다.
    새로운 클래스가 등장하면 랜덤 색상을 생성하여 추가합니다.
    """
    if class_name in class_colors:
        return class_colors[class_name]
    else:
        # 랜덤 색상 생성
        color = tuple(random.randint(0, 255) for _ in range(3))
        class_colors[class_name] = color
        return color

# ------------------------------
# 2. 레전드(legend) 그리기 함수 정의
# ------------------------------
def draw_legend(img, class_counts, class_colors, font, font_scale, text_thickness):
    """
    이미지에 클래스별 객체 수를 표시하는 레전드를 그립니다.
    """
    box_size = 20
    padding = 5
    margin = 10

    num_classes = len(class_counts)
    legend_height = num_classes * (box_size + padding) + padding
    legend_width = 0
    for class_name, count in class_counts.items():
        label = f"{class_name} ({count})"
        (text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, text_thickness)
        if text_width + box_size + 3 * padding > legend_width:
            legend_width = text_width + box_size + 3 * padding
    legend_width += margin  # 추가 여백

    # 레전드 배경 사각형 그리기 (흰색)
    legend_bg_start = (margin, margin)
    legend_bg_end = (margin + legend_width, margin + legend_height)
    cv2.rectangle(img, legend_bg_start, legend_bg_end, (255, 255, 255), cv2.FILLED)

    # 레전드 항목 그리기
    current_y = margin + padding
    for class_name, count in class_counts.items():
        color = get_class_color(class_name)
        label = f"{class_name} ({count})"

        # 색상 박스 그리기
        box_start = (margin + padding, current_y)
        box_end = (margin + padding + box_size, current_y + box_size)
        cv2.rectangle(img, box_start, box_end, color, cv2.FILLED)

        # 텍스트 그리기
        text_position = (margin + padding + box_size + padding, current_y + box_size - 5)
        cv2.putText(img, label, text_position, font, font_scale, (0, 0, 0), text_thickness, cv2.LINE_AA)

        # 다음 항목을 위해 y 위치 조정
        current_y += box_size + padding

# ------------------------------
# 3. API 요청을 통해 박스 데이터 가져오기
# ------------------------------
def get_box_data(image_path, url, username, access_key):
    """
    이미지를 API에 전송하여 박스 데이터를 가져옵니다.
    """
    try:
        with open(image_path, "rb") as image_file:
            image_data = image_file.read()

        response = requests.post(
            url=url,
            auth=HTTPBasicAuth(username, access_key),
            headers={"Content-Type": "image/jpeg"},
            data=image_data,
        )

        response.raise_for_status()  # HTTP 오류가 발생하면 예외를 발생시킵니다.

        return response.json().get("objects", [])
    except Exception as e:
        print(f"API 요청 중 오류 발생 ({image_path}): {e}")
        return []

# ------------------------------
# 4. 이미지에 박스 및 레이블 그리기
# ------------------------------
def draw_boxes_on_image(img, boxes, font, font_scale, text_thickness, confidence_threshold):
    """
    이미지에 박스와 레이블을 그립니다.
    """
    occupied_text_regions = []
    class_counts = Counter(obj['class'] for obj in boxes if obj['score'] >= confidence_threshold)

    for obj in boxes:
        class_name = obj['class']
        score = obj['score']
        box = obj['box']  # [x1, y1, x2, y2]

        if score < confidence_threshold:
            continue

        start_point = (box[0], box[1])  # 좌측 상단 (x1, y1)
        end_point = (box[2], box[3])    # 우측 하단 (x2, y2)

        color = get_class_color(class_name)
        thickness = 2  # 선 두께

        # 박스 그리기
        cv2.rectangle(img, start_point, end_point, color, thickness)

        # 텍스트 준비 (클래스명과 점수 표시)
        label = f"{class_name}: {score:.2f}"

        # 텍스트 크기 계산
        (text_width, text_height), baseline = cv2.getTextSize(label, font, font_scale, text_thickness)

        # 텍스트 배경 사각형 그리기 (가독성을 위해)
        # 텍스트를 박스 위쪽에 먼저 배치, 공간이 없으면 아래쪽에 배치
        text_bg_x1 = start_point[0]
        text_bg_y1 = start_point[1] - text_height - 10
        if text_bg_y1 < 0:
            text_bg_y1 = start_point[1] + text_height + 10
        text_bg_x2 = start_point[0] + text_width + 4
        text_bg_y2 = start_point[1] if text_bg_y1 < start_point[1] else start_point[1] + text_height + 20
        text_bg_start = (text_bg_x1, text_bg_y1)
        text_bg_end = (text_bg_x2, text_bg_y2)

        # 텍스트 배경 사각형 영역이 겹치지 않도록 조정
        overlap = False
        for region in occupied_text_regions:
            if not (text_bg_end[0] < region[0] or
                    text_bg_start[0] > region[0] + region[2] or
                    text_bg_end[1] < region[1] or
                    text_bg_start[1] > region[1] + region[3]):
                overlap = True
                break

        if overlap:
            # 텍스트를 박스 아래쪽으로 이동
            text_bg_y1 = start_point[1] + text_height + 10
            text_bg_y2 = start_point[1] + text_height + 20
            text_bg_start = (text_bg_x1, text_bg_y1)
            text_bg_end = (text_bg_x2, text_bg_y2)

            # 다시 겹치는지 확인
            overlap = False
            for region in occupied_text_regions:
                if not (text_bg_end[0] < region[0] or
                        text_bg_start[0] > region[0] + region[2] or
                        text_bg_end[1] < region[1] or
                        text_bg_start[1] > region[1] + region[3]):
                    overlap = True
                    break

            if overlap:
                # 추가적인 조정 (오프셋을 더 추가)
                text_bg_y1 += text_height + 5
                text_bg_y2 += text_height + 5
                text_bg_start = (text_bg_x1, text_bg_y1)
                text_bg_end = (text_bg_x2, text_bg_y2)

        # 텍스트 배경 사각형 그리기
        cv2.rectangle(img, text_bg_start, text_bg_end, color, cv2.FILLED)

        # 텍스트 위치 설정
        text_position = (text_bg_start[0] + 2, text_bg_end[1] - 5)

        # 텍스트 그리기 (흰색)
        cv2.putText(img, label, text_position, font, font_scale, (255, 255, 255), text_thickness, cv2.LINE_AA)

        # 점유된 위치 저장 (텍스트 배경 사각형)
        occupied_text_regions.append((text_bg_start[0], text_bg_start[1],
                                      text_bg_x2 - text_bg_x1, text_bg_y2 - text_bg_y1))

    # 레전드 그리기
    draw_legend(img, class_counts, class_colors, font, font_scale, text_thickness)

    return img

# ------------------------------
# 5. 메인 처리 루프
# ------------------------------
def main():
    # API 설정
    URL = ""
    USERNAME = ""  # API 사용자 이름
    ACCESS_KEY = ""  # API 액세스 키

    # 입력 및 출력 디렉토리 설정
    INPUT_DIR = ""  # 실제 입력 디렉토리 경로로 변경
    OUTPUT_DIR = ""  # 실제 출력 디렉토리 경로로 변경

    # 출력 디렉토리가 존재하지 않으면 생성
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    # 지원되는 이미지 확장자
    IMAGE_EXTENSIONS = ('*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff')

    # 모든 이미지 파일 경로 수집
    image_paths = []
    for ext in IMAGE_EXTENSIONS:
        image_paths.extend(glob.glob(os.path.join(INPUT_DIR, ext)))

    print(f"처리할 이미지 수: {len(image_paths)}")

    # 폰트 설정
    font = cv2.FONT_HERSHEY_COMPLEX_SMALL
    font_scale = 0.5
    text_thickness = 1

    # 신뢰도 기준 설정 (필요에 따라 조정)
    confidence_threshold = 0.3

    # 각 이미지 처리
    for idx, image_path in enumerate(image_paths, 1):
        image_name = os.path.basename(image_path)
        print(f"처리 중: {idx}/{len(image_paths)} - {image_name}")

        # API를 통해 박스 데이터 가져오기
        boxes = get_box_data(image_path, URL, USERNAME, ACCESS_KEY)

        if not boxes:
            print(f"박스 데이터를 가져올 수 없습니다: {image_name}")
            # 필요에 따라 기본 이미지를 저장하거나 건너뜁니다.
            output_path = os.path.join(OUTPUT_DIR, image_name)
            cv2.imwrite(output_path, cv2.imread(image_path))
            continue

        # 이미지 불러오기
        img = cv2.imread(image_path)

        # 이미지 로드 확인
        if img is None:
            print(f"이미지를 불러올 수 없습니다: {image_path}")
            continue

        # 박스 및 레이블 그리기
        img_with_boxes = draw_boxes_on_image(img, boxes, font, font_scale, text_thickness, confidence_threshold)

        # 결과 이미지 저장
        output_path = os.path.join(OUTPUT_DIR, image_name)
        cv2.imwrite(output_path, img_with_boxes)

    print("모든 이미지 처리가 완료되었습니다.")

if __name__ == "__main__":
    main()

 

 

이 코드를 이용하여 GRADIO에 적용해 보았다.

 

사이트를 생성하여 이미지 데이터를 첨부하면 YOLOv6-n 모델을 기반으로 객체 탐지 박스를 생성하여 박스가 생성된 이미지를 출력시켜주도록 하였다.

 

import cv2
import gradio as gr
import requests
import numpy as np
from PIL import Image
from requests.auth import HTTPBasicAuth

# 가상의 비전 AI API URL (예: 객체 탐지 API)
VISION_API_URL = ""
TEAM = ""
ACCESS_KEY = ""

# 클래스별 색상 매핑
class_colors = {
    'RASPBERRY PICO': (255, 0, 0),  # 파란색
    'USB': (0, 255, 0),             # 초록색
    'OSCILLATOR': (0, 0, 255),      # 빨간색
    'HOLE': (255, 255, 0),          # 노란색
    'CHIPSET': (255, 0, 255),       # 자홍색
    'BOOTSEL': (0, 255, 255)        # 청록색
}

def get_class_color(class_name):
    """
    클래스 이름에 따라 색상을 반환합니다.
    새로운 클래스가 등장하면 랜덤 색상을 생성하여 추가합니다.
    """
    if class_name in class_colors:
        return class_colors[class_name]
    else:
        # 랜덤 색상 생성
        color = tuple(np.random.randint(0, 256, size=3).tolist())
        class_colors[class_name] = color
        return color

def process_image(image):
    # 이미지를 OpenCV 형식으로 변환
    image = np.array(image)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # 이미지를 API에 전송할 수 있는 형식으로 변환
    _, img_encoded = cv2.imencode(".jpg", image)

    # API 호출 및 결과 받기
    headers = {
        "Content-Type": "application/octet-stream"
    }
    auth = HTTPBasicAuth(TEAM, ACCESS_KEY)

    # API에 POST 요청 전송
    response = requests.post(
        VISION_API_URL,
        data=img_encoded.tobytes(),
        headers=headers,
        auth=auth
    )

    # 요청이 성공했는지 확인
    if response.status_code == 200:
        results = response.json()
        print("API 응답:", results)  # 디버깅을 위해 추가
    else:
        print(f"Error: {response.status_code}")
        results = {}

    # API 결과를 바탕으로 박스 그리기
    if 'objects' in results:
        predictions = results['objects']
    elif 'predictions' in results:
        predictions = results['predictions']
    elif 'detections' in results:
        predictions = results['detections']
    else:
        print("API 응답에 예측 결과가 없습니다.")
        predictions = []

    if predictions:
        height, width = image.shape[:2]
        for pred in predictions:
            label = pred.get('class', 'N/A')
            score = pred.get('score', 0)
            bbox = pred.get('box', [0, 0, 0, 0])

            # 바운딩 박스 좌표 추출
            xmin, ymin, xmax, ymax = bbox

            # 좌표가 픽셀 값인 경우 스케일링 불필요
            # 좌표가 이미지 범위를 벗어나지 않도록 조정
            xmin = max(0, min(int(xmin), width - 1))
            ymin = max(0, min(int(ymin), height - 1))
            xmax = max(0, min(int(xmax), width - 1))
            ymax = max(0, min(int(ymax), height - 1))

            # 클래스별 색상 가져오기
            color = get_class_color(label)

            # 바운딩 박스 그리기
            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)

            # 텍스트 준비
            text = f"{label}: {score:.2f}"

            # 텍스트 크기 계산하여 배경 사각형 생성
            (text_width, text_height), baseline = cv2.getTextSize(
                text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
            )
            cv2.rectangle(
                image,
                (xmin, ymin - text_height - baseline),
                (xmin + text_width, ymin),
                color,
                thickness=cv2.FILLED,
            )

            # 바운딩 박스 위에 텍스트 표시
            cv2.putText(
                image,
                text,
                (xmin, ymin - baseline),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 0, 0),
                2,
            )
    else:
        print("이미지에서 객체를 감지하지 못했습니다.")
        cv2.putText(
            image,
            "No objects detected",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 0, 255),
            2,
        )

    # BGR 이미지를 RGB로 변환
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return Image.fromarray(image)

# Gradio 인터페이스 설정
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs="image",
    title="Vision AI Object Detection",
    description="Upload an image to detect objects using Vision AI.",
)

# 인터페이스 실행
iface.launch()

 

 

인터페이스 실행 코드를 다음과 같이 바꾸면 링크를 생성하여 모바일로도 실행을 시킬 수 있다.

iface.launch(share=True)

 

 

 

 

그 후에는 동영상에서 YOLOv6-n 모델을 기반으로 객체를 탐지해보는 실습을 진행하였다.

동영상에 적용을 하기 위해서는 프레임 단위로 나누어 이미지에서 객체 탐지를 하도록 한다.

 

https://youtube.com/shorts/zxFuit2jMSM

 

https://youtube.com/shorts/4on4FaZsE44?feature=share

 

video-ai-inference-practice.ipynb

0.08MB

 

streamlit을 이용하여 실시간 YOLOv8 추론 사이트도 생성하였다.

 

 

streamlit-web-rtc.py

0.00MB

 

 

라벨링_결과_가설_해결방안_F_1.pdf

0.33MB

 

컨베이어_벨트_객체_인식_딥러닝_모델_최적화_명은환_v.3.pdf

11.96MB