[Synology] AI 사진 자동 공유 > PyTorch 학습 (2)

안녕하세요, 이번 글에서는 PyTorch를 이용하여 간단한 분류 모델을 만들고 활용하는 방법에 대해 알아보겠습니다.

우선, 바로 사용할 수 있는 사진들을 그대로 사용하면 이미지의 노이즈가 심하고, 모델 학습이 효과적으로 이루어지지 않을 수 있습니다. 따라서, 첫 단계로 이미지에서 얼굴을 추출하는 작업을 진행하겠습니다. 이 과정은 학습 효율성을 높이기 위해 필요한 단계입니다. (이는 전체 이미지를 대상으로 하는 Object detection 대신, 얼굴 추출 후 이미지 분류 방식을 선택한 것입니다.)

얼굴 인식에는 cv2나 dlib 같은 유명 라이브러리들을 고려해 볼 수 있지만, 제가 원하는 목표인 아기 얼굴을 가장 잘 찾아내는 것은 mediapipe였습니다. 따라서, mediapipe를 사용하기로 결정했습니다. 학습 대상 얼굴이 포함된 폴더와 그렇지 않은 폴더를 구분하여, 두 폴더 내의 모든 얼굴을 찾아 저장하는 과정을 거쳤습니다.

 

먼저, 분류를 원하는 타겟이 있는 폴더와 나머지 폴더로 나눕니다. (필자의 케이스에는 아이가 있는경우 vs 없는경우)

예)

아이가 있는 폴더명 ['아이1', 아이2', ...]

아이가 없는 폴더명 ['결혼식', '졸업식', ...]

 

두개의 list를 class0 과 class1로 할당해서 진행하겠습니다.

 

1. Mediapipe를 이용하여 이미지에서 얼굴 추출 (face_detection)

import cv2
import mediapipe as mp
import os

BASEPATH = '/Volumes/photo/'

class0 = ['타겟 없는 폴더']
class1 = ['타켓 있는 폴더']

class0_folders = [os.path.join(BASEPATH, folder_name) for folder_name in class0]
class1_folders = [os.path.join(BASEPATH, folder_name) for folder_name in class1]

# Initialize MediaPipe Face Detection.
mp_face_detection = mp.solutions.face_detection
face_detection = mp_face_detection.FaceDetection(min_detection_confidence=0.5)  # min_detection_confidence 는 테스트 후 알맞는 숫자 선택

face_counter = 1

save_base = "cropped_faces"
os.makedirs(save_base, exist_ok=True)  # 폴더 없는 경우 생성

# 폴더에서 얼굴 추출 (추출경로, 라벨링)
def process_folder(folder_path, class_num):
    global face_counter
    for img_file in os.listdir(folder_path):
        if img_file.lower().endswith(('.jpg', '.jpeg', '.png')):   # jpg 만 추출
            try:
                image_path = os.path.join(folder_path, img_file)
                image = cv2.imread(image_path)
                if image is None:
                    continue

                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                results = face_detection.process(image_rgb)

                if results.detections:
                    for detection in results.detections:
                        bboxC = detection.location_data.relative_bounding_box
                        ih, iw, _ = image.shape
                        x, y, w, h = int(bboxC.xmin * iw), int(bboxC.ymin * ih), \
                                     int(bboxC.width * iw), int(bboxC.height * ih)
                        cropped_face = image[y:y+h, x:x+w]
                        save_folder = os.path.join(save_base, str(class_num))
                        os.makedirs(save_folder, exist_ok=True)
                        cv2.imwrite(f"{save_folder}/face_{face_counter}.jpg", cropped_face)
                        face_counter += 1
            except Exception as e:
                print(f"Error processing {img_file}: {e}")

# Process each folder
for folder in class0_folders:
    process_folder(folder, 0)

for folder in class1_folders:
    process_folder(folder, 1)

face_detection.close()
print("Face cropping completed.")

 

face detection 후 폴더를 확인해보니, 내부에 이미지가 잘 분류되어 얼굴만 추출되었습니다.

 

 

 

2. PyTorch를 이용하여 추출한 얼굴을 학습하여 pth 모델 파일 생성

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
import torch.nn.functional as F


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  

        self.fc1 = nn.Linear(64 * 8 * 8, 128)  
        self.fc2 = nn.Linear(128, 2)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = F.max_pool2d(F.relu(self.conv3(x)), 2)
        x = x.view(x.size(0), -1)  
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
model = CNN()

# Data Augmentation 으로 overfitting 방지
transform = transforms.Compose([
    transforms.RandomResizedCrop(64),  
    transforms.RandomHorizontalFlip(),  
    transforms.RandomRotation(10),  # 랜덤하게 사진 기울기 조정
    transforms.ToTensor(),  
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  
])


# 학습대상 폴더 지정
train_path = '/Volumes/home/Work/2023/[deleted]/cropped_faces'
full_train_data = datasets.ImageFolder(root=train_path, transform=transform)

# Splitting into train and validation
train_size = int(0.8 * len(full_train_data))  
val_size = len(full_train_data) - train_size  
train_data, val_data = random_split(full_train_data, [train_size, val_size])

train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
val_loader = DataLoader(val_data, batch_size=32, shuffle=True)
# Model, Loss Function, Optimizer
model = CNN()
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Early stopping parameters
early_stopping_patience = 5
min_val_loss = float('inf')
patience_counter = 0

# Training with validation and early stopping
num_epochs = 30
for epoch in range(num_epochs):
    model.train()
    total_loss = 0

    # Training loop
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        if (i+1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item()}')

    # Validation loop
    model.eval()
    with torch.no_grad():
        val_loss = 0
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            val_loss += loss.item()

        avg_val_loss = val_loss / len(val_loader)
        print(f'Epoch {epoch+1}/{num_epochs}, Validation Loss: {avg_val_loss}')

        # Early stopping check
        if avg_val_loss < min_val_loss:
            min_val_loss = avg_val_loss
            patience_counter = 0
        else:
            patience_counter += 1
            if patience_counter >= early_stopping_patience:
                print("Early stopping triggered")
                # Save the model when early stopping is triggered
                torch.save(model.state_dict(), 'model.pth')
                break
# Save the model after completing the training loop
torch.save(model.state_dict(), 'model.pth')
print("Model saved after completing all epochs.")

 

 

실제 학습 결과, Loss 값은 0.1xx 정도로, 모델 학습이 잘 이루어졌음을 확인할 수 있었습니다. 학습된 모델이 제대로 작동하는지 확인하기 위해, Inference를 진행했습니다.

 

테스트 결과, 마지막 아이의 사진이 'Predicted class: 0'으로 정확하게 분류된 것을 확인할 수 있었습니다. 이로써, model.pth 파일을 성공적으로 생성했습니다. 이 모델 파일을 활용하여 할 수 있는 일들은 매우 다양할 것으로 기대됩니다. 하지만, 우선은 NAS 공유폴더에 이미지를 자동으로 공유하는 작업부터 시작하려 합니다.

다음 글에서는 NAS 서버에 Docker를 이용하여 Ubuntu를 설치하고 설정하는 과정에 대해 다루어보겠습니다.

감사합니다.