250929 [ Day 60 ] - Project (5)

시작하며

오늘부터 본격적인 기능 구현 단계에 들어갔다. 크게 API와 DB 구현이 필요한 백엔드 파트, 웹 디자인 및 구현이 필요한 프론트엔드 파트, 그리고 AI 모델의 파이프라인을 구축하고 학습을 진행하는 머신러닝 파트 이렇게 3가지로 역할을 분담하여 작업을 시작했다. 그중에서 나는 머신러닝 파트를 담당하게 되었다.

YOLO

# YOLO 구현

from ultralytics import YOLO

# 모델 로딩을 한번만 하기 위해 클래스로 구현
class YOLODetector:
    # YOLO 초기화
    def __init__(self, model_path="yolo11n.pt"):
        self.model = YOLO(model_path)
    
    # 객체 탐지 함수
    def detect_objects(self, img_path):
        # 이미지 로드 (모델 적용 시 리스트 자동 생성)
        yolo_results = self.model(img_path)
        # 이미지 리스트에서 0번 이미지 로드
        detection = yolo_results[0]
        # 객체 정보를 저장할 리스트 생성
        detected_objects = []

        # 구조 확인
        print(f"타입 확인 : {type(yolo_results)}")
        print(f"결과 개수 : {len(yolo_results)}")

        # 검출된 객체 확인
        if detection.boxes is None:
            print("검출된 객체가 없습니다")
            return detected_objects
        print(f"검출된 객체 수: {len(detection.boxes)}")

        # 각 검출된 객체에 대해 정보 추출
        for box in detection.boxes:
            # 좌표 추출
            coords = box.xyxy[0].cpu().numpy()
            x1, y1, x2, y2 = coords
            # 정확도/신뢰도
            confidence = box.conf[0].cpu().numpy()
            # 객체의 클래스 정보
            class_id = int(box.cls[0].cpu().numpy())
            # 객체 정보를 딕셔너리로 저장
            object_info = {
                "bbox" : [int(x1), int(y1), int(x2), int(y2)],
                "confidence" : float(confidence),
                "class_id" : class_id
            }
            # 객체 정보를 리스트에 저장
            detected_objects.append(object_info)

            # 결과 확인용 출력
            class_name = detection.names[class_id]
            print(f"검출 : {class_name}, 신뢰도 : {confidence:.2f}")

        return detected_objects

# YOLO 실행
# import시 실행 방지
if __name__ == "__main__":
    # 1. 검출기 생성
    detector = YOLODetector()

    # 2. 이미지 검출 실행
    img = "datasets/yolo_test/p6.jpg"
    results = detector.detect_objects(img)

    # 3. 결과 출력
    print(f"\n=== 최종 결과 ===")
    print(f"총 {len(results)}개 객체 검출")

    for i, obj in enumerate(results):
        print(f"객체{i+1} : {obj}")

YOLO + ResNet 파이프라인

# YOLO + ResNet Model 파이프라인

from model import model, transform
from yolo_detector import YOLODetector
import cv2 as cv
import torch
from PIL import Image

class YOLOResNetPipeline:
    # 재활용 분류 매핑 (알파벳 순서: Can, Glass, Paper, Plastic, Styrofoam, Vinyl)
    recycling_classes = {
        0: {"category": "캔", "item_type": "캔류", "method": "내용물 비우고 캔 전용 수거함"}, # Can
        1: {"category": "유리", "item_type": "유리병", "method": "뚜껑 분리하고 유리 전용 수거함"}, # Glass
        2: {"category": "종이", "item_type": "종이류", "method": "테이프 제거하고 종이 전용 수거함"}, # Paper
        3: {"category": "플라스틱", "item_type": "플라스틱", "method": "라벨 제거하고 플라스틱 전용 수거함"}, # Plastic
        4: {"category": "스티로폼", "item_type": "스티로폼", "method": "이물질 제거하고 스티로폼 전용 수거함"}, # Styrofoam
        5: {"category": "비닐", "item_type": "비닐류", "method": "이물질 제거하고 비닐 전용 수거함"} # Vinyl
    }

    # 파이프라인 초기화
    def __init__(self):
        # YOLO 초기화
        self.yolo = YOLODetector()
        # ResNet 모델 초기화
        self.resnet = model
        self.transform = transform

    # 객체 처리 함수
    def process_object(self, img_path):
        print(f"이미지 처리 시작: {img_path}")

        # 원본 이미지 로드 및 확인
        original_image = cv.imread(img_path)
        if original_image is None:
            print("이미지를 로드할 수 없습니다!")
            return []
        
        # YOLO 객체 검출
        yolo_results = self.yolo.detect_objects(img_path)
        print(f"YOLO 검출 완료: {len(yolo_results)}개 객체")

        # 객체 부분만 자르기
        for idx, box in enumerate(yolo_results):
            print(f"\n객체 {idx+1} / {len(yolo_results)} 처리 중...")
            # 좌표 추출
            x1, y1, x2, y2 = box["bbox"]
            # 이미지 자르기
            cropped_img = original_image[y1:y2, x1:x2]
            # 저장해서 확인
            # cv.imwrite("crop_test_image.jpg", cropped_img)

            # BGR -> RGB 변환
            cropped_rgb = cv.cvtColor(cropped_img, cv.COLOR_BGR2RGB)
            # PIL 포맷으로 변환
            pil_img = Image.fromarray(cropped_rgb)
            # transform 적용 (tensor로 변환)
            input_tensor = self.transform(pil_img)
            # 배치 차원 추가
            input_batch = input_tensor.unsqueeze(0)
            # 디바이스로 이동
            device = next(self.resnet.parameters()).device
            input_batch = input_batch.to(device)

            # 모델 추론
            self.resnet.eval() # 평가 모드

            with torch.no_grad():
                outputs = self.resnet(input_batch)
                # 확률로 변환
                prob = torch.nn.functional.softmax(outputs[0], dim=0)
                # 가장 높은 확률의 클래스
                predicted_class = torch.argmax(outputs[0]).item()
                # 그 클래스의 신뢰도
                confidence = prob[predicted_class].item()

                # 결과 출력
                print(f"ResNet18 분류: 클래스 {predicted_class}, 신뢰도 {confidence:.3f}")

                # YOLO결과 + ResNet18 결과
                box["resnet_class"] = predicted_class
                box["resnet_confidence"] = confidence

        return yolo_results
        
# 테스트 실행
if __name__ == "__main__":
    pipeline = YOLOResNetPipeline()
    results = pipeline.process_object("datasets/p6.jpg")
    print(f"\n최종 결과: {len(results)}개 객체 검출됨")

마치며

기본적인 모델 구현보다 훨씬 더 복잡했던 파이프라인 구축이었다. 거의 작업 시간의 80%가 공부하거나 찾아보는 시간일 정도로 어려웠던 것 같다. 그래도 어느 정도 생각했던 방향으로 흘러가는 것 같아서 내일이면 파이프라인 테스트와 API로 응답할 정보 처리가 끝나고 모델 학습에 들어갈 수 있을 것 같다.