import torch

import torch.nn as nn

import torch.optim as optim


# YOLO 모델의 정의

class YOLO(nn.Module):

    def __init__(self, num_classes=20):

        super(YOLO, self).__init__()

        # 간단한 합성곱 신경망(CNN)을 사용하여 특징을 추출하는 부분

        self.conv_layers = nn.Sequential(

            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),  # 입력 채널 3, 출력 채널 16, 커널 사이즈 3x3, stride 1, padding 1

            nn.ReLU(),                                             # 활성화 함수 ReLU 적용

            nn.MaxPool2d(2, 2),                                    # 2x2 맥스 풀링

            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),  # 입력 채널 16, 출력 채널 32

            nn.ReLU(),                                             # ReLU 활성화

            nn.MaxPool2d(2, 2),                                    # 2x2 맥스 풀링

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),  # 입력 채널 32, 출력 채널 64

            nn.ReLU(),                                             # ReLU 활성화

            nn.MaxPool2d(2, 2)                                     # 2x2 맥스 풀링

        )

        # 완전 연결(fully connected) 층: 경계 상자와 클래스 예측을 위한 부분

        self.fc_layers = nn.Sequential(

            nn.Flatten(),                                          # 다차원 텐서를 1차원으로 평탄화

            nn.Linear(64 * 8 * 8, 256),                             # 입력 크기: 64채널, 8x8 특징맵 -> 256 뉴런

            nn.ReLU(),                                             # ReLU 활성화

            nn.Linear(256, S * S * (B * 5 + num_classes))           # 최종 출력: 각 그리드당 (B*5 + 클래스 개수) 예측값, 전체 그리드 수 SxS

        )


    # 순전파(forward) 함수 정의

    def forward(self, x):

        x = self.conv_layers(x)  # 합성곱 계층을 통과

        x = self.fc_layers(x)    # 완전 연결 계층을 통과하여 최종 예측값 생성

        return x


# 그리드 크기와 각 그리드당 예측할 경계 상자의 수 설정

S = 7  # 7x7 그리드로 분할

B = 2  # 각 그리드에서 2개의 경계 상자 예측


# 모델과 최적화 기법 초기화

yolo_model = YOLO(num_classes=20)               # YOLO 모델 생성 (클래스 수: 20)

optimizer = optim.Adam(yolo_model.parameters(), lr=0.001)  # Adam 옵티마이저 사용, 학습률 0.001

criterion = nn.MSELoss()                        # 손실 함수로 MSE(평균 제곱 오차) 사용


# 입력 이미지와 목표 출력(라벨) 생성 (여기서는 임의의 랜덤 값 사용)

data = torch.randn(1, 3, 224, 224)  # 1개의 예시 이미지, 3채널, 224x224 크기

labels = torch.randn(1, S * S * (B * 5 + 20))  # 해당 크기에 맞는 임의의 라벨 생성


# 모델 학습 단계

yolo_model.train()      # 모델을 학습 모드로 전환

optimizer.zero_grad()   # 기울기 초기화

outputs = yolo_model(data)  # 모델의 순전파 결과 계산

loss = criterion(outputs, labels)  # 예측값과 라벨 간의 손실 계산

loss.backward()         # 역전파 수행하여 기울기 계산

optimizer.step()        # 옵티마이저를 통해 파라미터 업데이트


# 손실 값 출력

print("Loss: {:.4f}".format(loss.item()))



---
예전에 책보고 저장해둔거.

단순 CNN 사용해서, 특징 추출 후ㅡ Fully 커넥티드 레이어로 예측값 산출함.

학습과정은 기초적인

forward->backWard->Optimizer.step() 절차로 진행하고.

그래서 아마 어떤 이미지 특징을 검증하는지는 모르겠지만 기본적으로 뼈대는 프레임워크 쓰면 되긴함

우선 데이터 준비

-> 특정 이미지 데이터셋 수집
-> 데이터 라벨링 객체 위치와 클래스 정보 라벨링

여기서 모델 선택하고 배포하면될듯?

완전 기초라 실제 사용하려면 기능 추가가 필요하긴할텐데, 그냥 뼈대는 이런구조다 정도로 기억하고 있긴함