250930 [ Day 61 ] - Project (6)

시작하며

오늘은 기존에 테스트로 만들어둔 이항 분류 모델을 6개의 클래스를 가진 다항 분류 모델로 수정한 뒤, 그에 맞춰 파이프라인도 약간 수정해서 클래스당 약 1000개 정도의 데이터를 사용해서 학습시킨 프로토타입 모델을 만들었다. 일단은 새로운 API와 웹과의 연동을 확인해보기 위한 용도라 모델의 정확도보다는 빠른 구현과 API 응답 데이터에 우선순위를 두고 작업하였고, 내일 연동이 정상적으로 작동하면 본격적인 모델 학습을 시작할 것 같다.

Prototype Model v1

Data

• train : 클래스별 800
• valid : 클래스별 150

Model

• ResNet18

Epoch

• 5

Learning Rate

• 0.001

Batch Size

• 32

Apple MPS (Metal Performance Shaders) 사용
선택된 디바이스: mps

[Epoch] 1 / 5
[batch :   10], Loss : 0.771349 (  288 /  4800)
[batch :   20], Loss : 0.885182 (  608 /  4800)
[batch :   30], Loss : 0.889852 (  928 /  4800)
[batch :   40], Loss : 0.701717 ( 1248 /  4800)
[batch :   50], Loss : 0.861263 ( 1568 /  4800)
[batch :   60], Loss : 0.310380 ( 1888 /  4800)
[batch :   70], Loss : 0.768648 ( 2208 /  4800)
[batch :   80], Loss : 0.618361 ( 2528 /  4800)
[batch :   90], Loss : 0.591268 ( 2848 /  4800)
[batch :  100], Loss : 0.357009 ( 3168 /  4800)
[batch :  110], Loss : 0.537911 ( 3488 /  4800)
[batch :  120], Loss : 0.352613 ( 3808 /  4800)
[batch :  130], Loss : 0.508844 ( 4128 /  4800)
[batch :  140], Loss : 0.411743 ( 4448 /  4800)
[batch :  150], Loss : 0.272760 ( 4768 /  4800)
Validation Results: Accuracy: 0.682 (68.2%), Avg loss: 0.8979

[Epoch] 2 / 5
[batch :   10], Loss : 0.423951 (  288 /  4800)
[batch :   20], Loss : 0.342776 (  608 /  4800)
[batch :   30], Loss : 0.422824 (  928 /  4800)
[batch :   40], Loss : 0.248441 ( 1248 /  4800)
[batch :   50], Loss : 0.388495 ( 1568 /  4800)
[batch :   60], Loss : 0.640245 ( 1888 /  4800)
[batch :   70], Loss : 0.132340 ( 2208 /  4800)
[batch :   80], Loss : 0.429424 ( 2528 /  4800)
[batch :   90], Loss : 0.659897 ( 2848 /  4800)
[batch :  100], Loss : 0.102553 ( 3168 /  4800)
[batch :  110], Loss : 0.503341 ( 3488 /  4800)
[batch :  120], Loss : 0.712427 ( 3808 /  4800)
[batch :  130], Loss : 0.121946 ( 4128 /  4800)
[batch :  140], Loss : 0.212581 ( 4448 /  4800)
[batch :  150], Loss : 0.336336 ( 4768 /  4800)
Validation Results: Accuracy: 0.763 (76.3%), Avg loss: 0.6591

[Epoch] 3 / 5
[batch :   10], Loss : 0.246793 (  288 /  4800)
[batch :   20], Loss : 0.194673 (  608 /  4800)
[batch :   30], Loss : 0.131795 (  928 /  4800)
[batch :   40], Loss : 0.329142 ( 1248 /  4800)
[batch :   50], Loss : 0.223976 ( 1568 /  4800)
[batch :   60], Loss : 0.498322 ( 1888 /  4800)
[batch :   70], Loss : 0.155122 ( 2208 /  4800)
[batch :   80], Loss : 0.585494 ( 2528 /  4800)
[batch :   90], Loss : 0.352118 ( 2848 /  4800)
[batch :  100], Loss : 0.086463 ( 3168 /  4800)
[batch :  110], Loss : 0.256145 ( 3488 /  4800)
[batch :  120], Loss : 0.187811 ( 3808 /  4800)
[batch :  130], Loss : 0.201910 ( 4128 /  4800)
[batch :  140], Loss : 0.340605 ( 4448 /  4800)
[batch :  150], Loss : 0.095940 ( 4768 /  4800)
Validation Results: Accuracy: 0.702 (70.2%), Avg loss: 0.9834

[Epoch] 4 / 5
[batch :   10], Loss : 0.167587 (  288 /  4800)
[batch :   20], Loss : 0.118853 (  608 /  4800)
[batch :   30], Loss : 0.052107 (  928 /  4800)
[batch :   40], Loss : 0.113755 ( 1248 /  4800)
[batch :   50], Loss : 0.282270 ( 1568 /  4800)
[batch :   60], Loss : 0.141293 ( 1888 /  4800)
[batch :   70], Loss : 0.452946 ( 2208 /  4800)
[batch :   80], Loss : 0.173623 ( 2528 /  4800)
[batch :   90], Loss : 0.232137 ( 2848 /  4800)
[batch :  100], Loss : 0.095230 ( 3168 /  4800)
[batch :  110], Loss : 0.147951 ( 3488 /  4800)
[batch :  120], Loss : 0.110777 ( 3808 /  4800)
[batch :  130], Loss : 0.125662 ( 4128 /  4800)
[batch :  140], Loss : 0.079683 ( 4448 /  4800)
[batch :  150], Loss : 0.020562 ( 4768 /  4800)
Validation Results: Accuracy: 0.747 (74.7%), Avg loss: 1.0222

[Epoch] 5 / 5
[batch :   10], Loss : 0.051510 (  288 /  4800)
[batch :   20], Loss : 0.009294 (  608 /  4800)
[batch :   30], Loss : 0.077664 (  928 /  4800)
[batch :   40], Loss : 0.026722 ( 1248 /  4800)
[batch :   50], Loss : 0.369626 ( 1568 /  4800)
[batch :   60], Loss : 0.370070 ( 1888 /  4800)
[batch :   70], Loss : 0.243725 ( 2208 /  4800)
[batch :   80], Loss : 0.125985 ( 2528 /  4800)
[batch :   90], Loss : 0.108147 ( 2848 /  4800)
[batch :  100], Loss : 0.027245 ( 3168 /  4800)
[batch :  110], Loss : 0.075170 ( 3488 /  4800)
[batch :  120], Loss : 0.081495 ( 3808 /  4800)
[batch :  130], Loss : 0.119835 ( 4128 /  4800)
[batch :  140], Loss : 0.223253 ( 4448 /  4800)
[batch :  150], Loss : 0.082168 ( 4768 /  4800)
Validation Results: Accuracy: 0.758 (75.8%), Avg loss: 0.8861

역시나 에포크 수와 학습 데이터 수가 적고 파인 튜닝을 전혀 하지 않았기 때문에 손실과 정확도가 매우 나빴다. 빨리 연동이 모두 끝나고 본격적인 학습을 시켜보고 싶다.

Model

# AI Model 구현 및 학습

import torch
import torchvision.models as models
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from torch.utils.data import DataLoader

# 분류할 클래스 수
num_classes = 6

# 이미 학습되어 있는 resnet18 모델 불러오기
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(512, num_classes)

# 데이터 불러오기
transform = transforms.Compose([
    # 데이터 전처리 (224x224) 사이즈로 통일
    transforms.Resize((224,224)),
    # 텐서로 변환
    transforms.ToTensor(),
    # ImageNet 통계로 정규화
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 데이터 로더 생성 함수 (학습 시에만 호출)
def create_data_loaders():
    # 데이터 불러오기
    train_dataset = datasets.ImageFolder("datasets/train", transform=transform)
    test_dataset = datasets.ImageFolder("datasets/test", transform=transform)
    valid_dataset = datasets.ImageFolder("datasets/valid", transform=transform)

    # 배치 사이즈 설정
    batch_size = 32

    # 데이터 로더 생성
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True, # 에포크마다 섞어서 훈련
        num_workers=0 # 병렬처리할 cpu 코어 / 윈도우 multiprocessing error시 0으로 처리
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0
    )
    valid_loader = DataLoader(
        valid_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=0
    )

    return train_loader, test_loader, valid_loader

# 사용가능한 디바이스 확인
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA GPU 사용")
    print(f"GPU 이름: {torch.cuda.get_device_name(0)}")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
    print("Apple MPS (Metal Performance Shaders) 사용")
else:
    device = torch.device("cpu")
    print("CPU 사용")

print(f"선택된 디바이스: {device}")

# 모델을 디바이스로 이동
model = model.to(device)

# 손실 함수 정의
criterion = nn.CrossEntropyLoss()

# 옵티마이저 설정
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 학습 루프 함수 정의
def train_loop(data_loader, model, criterion, optimizer):
    model.train() # 학습 모드
    size = len(data_loader.dataset)
    running_loss = 0.

    for batch_idx, (data, target) in enumerate(data_loader):
        device = next(model.parameters()).device
        data, target = data.to(device), target.long().to(device)

        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * data.size(0)

        if (batch_idx+1) % 10 == 0:
            current = batch_idx * len(data)
            print(f"[batch : {batch_idx+1: 4d}], Loss : {loss.item():>7f} ({current:>5d} / {size:>5d})")
    
    epoch_loss = running_loss / size
    return epoch_loss

# 검증 루프 함수 정의
def valid_loop(data_loader, model, criterion):
    model.eval() # 평가 모드
    size = len(data_loader.dataset)
    valid_loss = 0.
    correct = 0

    with torch.no_grad():
        for data, target in data_loader:
            device = next(model.parameters()).device
            data, target = data.to(device), target.long().to(device)

            outputs = model(data)
            loss = criterion(outputs, target)
            valid_loss += loss.item() * data.size(0)

            # 정확도 계산
            predictions = torch.argmax(outputs, dim=1)
            correct += (predictions == target).sum().item()
        
    # 평균 계산
    avg_loss = valid_loss / size
    accuracy = correct / size

    print(f"Validation Results: Accuracy: {accuracy:.3f} ({100*accuracy:.1f}%), Avg loss: {avg_loss:.4f}")
    return avg_loss, accuracy

# 학습 반복 수
epochs = 5

# 학습 실행
if __name__ == "__main__":
    # 데이터 로더 생성
    train_loader, test_loader, valid_loader = create_data_loaders()

    for epoch in range(epochs):
        print(f"\n[Epoch] {epoch+1} / {epochs}")
        train_loss = train_loop(train_loader, model, criterion, optimizer)
        valid_loss, valid_acc = valid_loop(valid_loader, model, criterion)
    print("\n학습 및 검증 완료!")

    # 모델 저장
    torch.save(model.state_dict(), "prototype_model_v1.pth")
    print("\n모델 저장 완료!")

YOLO

# YOLO 구현

from ultralytics import YOLO
import numpy as np

# 모델 로딩을 한번만 하기 위해 클래스로 구현
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}")

Pipeline

# YOLO + ResNet Model 파이프라인

from model import transform
from yolo_detector import YOLODetector
import cv2 as cv
import torch
import torchvision.models as models
import torch.nn as nn
from PIL import Image

# 학습이 완료된 모델 가져오기
model_path = "../backend/models/prototype_model_v1.pth"
def load_trained_model(model_path=model_path):
    # 모델 구조
    model = models.resnet18(pretrained=False)
    model.fc = nn.Linear(512, 6)
    # 학습시킨 가중치 업데이트
    model.load_state_dict(torch.load(model_path, map_location="cpu"))
    model.eval() # 검증 모드
    # 디바이스 설정
    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    # 모델을 디바이스로 이동
    model = model.to(device)
    return model

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 = load_trained_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
    
    # API용 정보 응답 함수
    def format_recycling_response(self, yolo_results, img_path=""):
        # ResNet 분류 결과를 담을 리스트 생성
        recycling_items = []
        # 분류에 실패한 객체를 담을 리스트(피드백 및 DB용)
        unclassified_items = []
        for idx, object in enumerate(yolo_results):
            # ResNet 분류 결과가 있는 경우에만 처리
            if "resnet_class" in object:
                recycling_info = self.recycling_classes.get(object["resnet_class"])
                item = {
                    "item_id": idx + 1,
                    "location": {
                        "bbox": object["bbox"],
                        "confidence": object["confidence"]
                    },
                    "recycling_info": {
                        "category": recycling_info["category"],
                        "item_type": recycling_info["item_type"],
                        "recycling_method": recycling_info["method"],
                        "confidence": object["confidence"]
                    }
                }
                recycling_items.append(item)
            # 분류 실패시 피드백 요청
            else:
                unclassified_item = {
                    "item_id": idx + 1,
                    "location": {
                        "bbox": object["bbox"],
                        "confidence": object["confidence"]
                    },
                    "status": "classification_failed",
                    "feedback_request": {
                        "message": "이 객체의 재활용 분류를 도와주세요!",
                        "options": ["캔", "유리", "종이", "플라스틱", "스티로폼", "비닐"]
                    }
                }
                unclassified_items.append(unclassified_item)

        # API 응답 구성
        response = {
            "status": "success",
            "total_items": len(recycling_items) + len(unclassified_items),
            "classified_items": len(recycling_items),
            "unclassified_items": len(unclassified_items),
            "recycling_items": recycling_items
        }
        if unclassified_items:
            response["feedback_needed"] = unclassified_items
            response["summary"] = f"총 {len(recycling_items)}개 분류 완료, {len(unclassified_items)}개 항목의 사용자 피드백 필요"
        else:
            response["summary"] = f"총 {len(recycling_items)}개의 재활용품이 모두 분류되었습니다!"
        return response

# =============테스트 실행=============
if __name__ == "__main__":
    pipeline = YOLOResNetPipeline()
    # 테스트할 이미지 파일들
    test_images = [
        "datasets/pipe_test/test1.jpg",
        "datasets/pipe_test/test2.jpg",
        "datasets/pipe_test/test3.jpg",
        "datasets/pipe_test/test4.jpg",
        "datasets/pipe_test/test5.jpg"
    ]
    print("YOLO + ResNet 파이프라인 종합 테스트 시작")

    for idx, img_path in enumerate(test_images):
        print(f"\n테스트 {idx+1}/5: {img_path}")
        # 파이프라인 실행
        results = pipeline.process_object(img_path)
        # API 응답 생성
        api_response = pipeline.format_recycling_response(results, img_path)
        # 결과 요약 출력
        print(f"\n결과 요약:")
        print(f"   • YOLO 탐지: {len(results)}개 객체")
        print(f"   • 분류 완료: {api_response["classified_items"]}개")
        print(f"   • 미분류: {api_response["unclassified_items"]}개")
        print(f"   • 요약: {api_response["summary"]}")
        # 상세 분류 결과
        if api_response["recycling_items"]:
            print(f"\n분류 결과:")
            for item in api_response["recycling_items"]:
                category = item["recycling_info"]["category"]
                confidence = item["recycling_info"]["confidence"]
                print(f"   • 객체 {item["item_id"]}: {category} (신뢰도: {confidence:.3f})")
    print("\n전체 테스트 완료!")

마치며

수업 시간에 AI 모델 코드를 따라칠 때는 너무 어렵게만 느껴졌는데 역시나 직접 알아보면서 작성하니까 많이 이해도 되고 상당히 재미있게 작업하는 중이다. 계속 수정해보면서 좋은 결과를 낼 수 있었으면 좋겠다.