[11-12] 데이터 수집 + 모델

1. 추가 스케치 데이터 수집

1. Children - (Anger)

2. Children - (Sadness)

3. 20대 대학생 - (Happiness)

4. 20대 대학생 - (Happiness)

5. 20대 대학생 - (Sadness)

2. 모델 코드 - ResNet

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image

from datasets import load_dataset

ds = load_dataset("xodhks/EmoSet118K_MonetStyle")

emotion_classes = ["happiness", "sadness", "fear", "surprise", "anger", "disgust"]

transform = transforms.Compose([
    transforms.Resize((224, 224)), 
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

class SketchDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image = Image.open(self.data[idx]).convert("RGB")  
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

# 모델 정의 (ResNet-18 사용)
def create_resnet_model(num_classes):
    model = models.resnet18(pretrained=True) 
    model.fc = nn.Linear(model.fc.in_features, num_classes) 
    return model

# 학습 및 평가 함수
def train_model(model, dataloader, criterion, optimizer, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
        
        epoch_loss = running_loss / len(dataloader.dataset)
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_classes = len(emotion_classes)
model = create_resnet_model(num_classes).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

train_data = ["path/to/image1.jpg", "path/to/image2.jpg", ...]  # 이미지 경로 리스트
train_labels = [0, 1, ...]  # 각 이미지에 대한 라벨 (0~5까지의 정수로 매핑)
train_dataset = SketchDataset(train_data, train_labels, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 모델 학습
train_model(model, train_loader, criterion, optimizer, num_epochs=10)

• 추가 수정 → 이미지를 제대로 못받음, 크기 조정?

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
from datasets import load_dataset
from PIL import Image

# 감정 클래스 목록 및 매핑
emotion_classes = ["Anger", "Disgust", "Fear", "Happiness", "Sadness", "Surprise"]
label_mapping = {emotion: idx for idx, emotion in enumerate(emotion_classes)}

# 데이터 전처리 (이미지 크기 조정 및 정규화)
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # ResNet 입력 크기
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 사용자 정의 데이터셋 클래스
class EmoSetDataset(Dataset):
    def __init__(self, dataset, transform=None):
        self.dataset = dataset
        self.transform = transform

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        image = Image.fromarray(self.dataset[idx]["image"]).convert("RGB")
        label = label_mapping[self.dataset[idx]["emotion"]]
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

# 데이터셋 로드
ds = load_dataset("xodhks/EmoSet118K_MonetStyle")

# 데이터셋 및 데이터로더 생성
train_dataset = EmoSetDataset(ds["train"], transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 모델 정의 (ResNet-18 사용)
def create_resnet_model(num_classes):
    model = models.resnet18(pretrained=True)
    model.fc = nn.Linear(model.fc.in_features, num_classes)  # 마지막 레이어 수정
    return model

# 모델, 손실 함수, 최적화 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = create_resnet_model(num_classes=len(emotion_classes)).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 학습 및 평가 함수 정의
def train_model(model, dataloader, criterion, optimizer, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for inputs, labels in dataloader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * inputs.size(0)
        
        epoch_loss = running_loss / len(dataloader.dataset)
        print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")

# 학습 실행
train_model(model, train_loader, criterion, optimizer, num_epochs=10)

# 모델 저장 예시
torch.save(model.state_dict(), "emotion_classification_resnet18.pth")