[4주차] 뒷통수 인식 모델 구현해보기
Category: AI Core Projects
Tags: classification project
당근 이미지가 사람들 뒷통수도 가리면서 지저분하다는 피드백을 받아서 이를 해결해보기 위해 간단한 앞통수/뒷통수 classification 모델을 구현해봤다.
Data Preprocessing
필요한 모듈 전부 불러오기
# 필요한 모듈 불러오기
import os
import shutil
import math
import time
import copy
import random
import torch
import torchvision.transforms as transforms
from torchvision import models
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from torch.optim import lr_scheduler
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
데이터 분할하기 : train/val/test 데이터를 각각 6:2:2 비율로 분배
# 데이터셋 분할하기
dataset_dir = './dataset'
class_list = os.listdir(dataset_dir) # ['backHead', 'face']
splitted_dataset_dir = './splitted'
try:
os.mkdir(splitted_dataset_dir)
except:
shutil.rmtree('./splitted')
os.mkdir(splitted_dataset_dir)
train_dir = os.path.join(splitted_dataset_dir, 'train')
os.mkdir(train_dir)
validation_dir = os.path.join(splitted_dataset_dir, 'val')
os.mkdir(validation_dir)
test_dir = os.path.join(splitted_dataset_dir, 'test')
os.mkdir(test_dir)
for cls in class_list:
os.mkdir(os.path.join(train_dir, cls))
os.mkdir(os.path.join(validation_dir, cls))
os.mkdir(os.path.join(test_dir, cls))
for cls in class_list:
path = os.path.join(dataset_dir, cls)
fnames = os.listdir(path)
random.shuffle(fnames)
train_size = math.floor(len(fnames) * 0.6)
validation_size = math.floor(len(fnames)*0.2)
test_size = math.floor(len(fnames)*0.2)
train_fnames = fnames[:train_size]
print("Train size(",cls,"): ", len(train_fnames))
for fname in train_fnames:
src = os.path.join(path, fname)
dst = os.path.join(os.path.join(train_dir, cls), fname)
shutil.copyfile(src, dst)
validation_fnames = fnames[train_size:(validation_size + train_size)]
print("Validation size(",cls,"): ", len(validation_fnames))
for fname in validation_fnames:
src = os.path.join(path, fname)
dst = os.path.join(os.path.join(validation_dir, cls), fname)
shutil.copyfile(src, dst)
test_fnames = fnames[(train_size+validation_size):(validation_size + train_size +test_size)]
print("Test size(",cls,"): ", len(test_fnames))
for fname in test_fnames:
src = os.path.join(path, fname)
dst = os.path.join(os.path.join(test_dir, cls), fname)
shutil.copyfile(src, dst)
Model Training
먼저 학습을 하기 위한 hyperparameter 선택
USE_CUDA = torch.cuda.is_available()
DEVICE = torch.device("cuda" if USE_CUDA else "cpu")
BATCH_SIZE = 256
EPOCH = 20
학습 데이터 전처리 및 로드하기
data_transforms = {
'train': transforms.Compose([transforms.Resize([64,64]),
transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(),
transforms.RandomCrop(52), transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]),
'val': transforms.Compose([transforms.Resize([64,64]),
transforms.RandomCrop(52), transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])
}
data_dir = './splitted'
image_datasets = {x: ImageFolder(root=os.path.join(data_dir, x), transform=data_transforms[x]) for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=BATCH_SIZE, shuffle=True, num_workers=4) for x in ['train', 'val']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
학습에 사용할 resnet 모델 불러오기
resnet = models.resnet50(pretrained=True)
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 2)
resnet = resnet.to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.Adam(filter(lambda p: p.requires_grad, resnet.parameters()), lr=0.001)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
resnet 모델을 training하는 함수
def train_resnet(model, criterion, optimizer, scheduler, num_epochs=25):
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
print('-------------- epoch {} ----------------'.format(epoch+1))
since = time.time()
for phase in ['train', 'val']:
if phase == 'train':
model.train()
else:
model.eval()
running_loss = 0.0
running_corrects = 0
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(DEVICE)
labels = labels.to(DEVICE)
optimizer.zero_grad()
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
if phase == 'train':
scheduler.step()
epoch_loss = running_loss/dataset_sizes[phase]
epoch_acc = running_corrects.double()/dataset_sizes[phase]
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
if phase == 'val' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
time_elapsed = time.time() - since
print('Completed in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
model.load_state_dict(best_model_wts)
return model
실제 training 시키는 코드
model_resnet50 = train_resnet(resnet, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=EPOCH)
torch.save(model_resnet50, 'resnet50.pt')
Model Testing
모델 테스트를 위한 데이터 전처리
transform_resNet = transforms.Compose([
transforms.Resize([64,64]),
transforms.RandomCrop(52),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_resNet = ImageFolder(root='./splitted/test', transform=transform_resNet)
test_loader_resNet = torch.utils.data.DataLoader(test_resNet, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
모델 testing하는 함수
def evaluate(model, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(DEVICE), target.to(DEVICE)
output = model(data)
test_loss += F.cross_entropy(output,target, reduction='sum').item()
pred = output.max(1, keepdim=True)[1]
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
test_accuracy = 100. * correct / len(test_loader.dataset)
return test_loss, test_accuracy
모델 테스트 loss, accuracy 값 확인하기
resnet50=torch.load('resnet50.pt')
resnet50.eval()
test_loss, test_accuracy = evaluate(resnet50, test_loader_resNet)
print('ResNet test acc: ', test_accuracy)
-> 돌려본 결과 생각보다 정확도가 높게 나왔다
한 장의 이미지를 넣어서 test하는 코드
test_img = Image.open("넣을 이미지").convert('RGB')
plt.imshow(test_img)
plt.show() # 이미지 보기
test_resNet = transform_resNet(test_img)
test_resNet = test_resNet.unsqueeze(0)
test_resNet = test_resNet.to(DEVICE)
model = resnet50.eval()
output = model(test_resNet)
index = output.data.cpu().numpy().argmax()
class_name = class_list[index]
print(class_name) # 분류한 output
-> 어떤 미용실 광고에서 가져온 사진을 넣어봤는데 잘 돌아간다 !
Result
생각보다 성능이 좋게 나와서 바로 영상에도 적용시켜봤다.
저번주 프로젝트 파일에 backhead 옵션을 추가해서 뒷통수가 아닐 때만 키위 이미지를 넣도록 수정한 결과, 뒷통수를 잘 인식하고 사진이 들어가지 않았다!
깔끔하게 잘 가려지는 듯 했지만 문제는 얘가 이젠 앞통수도 헷갈려 하는 것 같다. 마스크 낀 얼굴을 뒷통수로 가끔 인식한다.
내가 생각했을 때 개선해야 할 점들
- 애초에 학습시킨 데이터셋이 너무 적다. (최소 천 장)
- 마스크 쓴 얼굴 데이터 셋도 넣었어야 한다.
- 여전히 이미지(키위) 크기가 너무 왔다갔다 정신없다.
이번 프로젝트 느낀 점
내 첫 딥러닝 프로젝트여서 사실 결과도 못 뽑는거 아닌지 걱정이 많았는데 그래도 발표는 할 수는 있을 정도로 마무리돼서 다행이다. 플젝 기간이 한 달보다 좀 더 길었으면 더 좋은 결과를 만들 수 있었을거 같은데 아쉽다. 이번 방학 때 딥러닝 공부도 할 겸 남은 부분들을 좀 더 고쳐봐야겠다. 😭