포켓몬 분류 데이터셋

Ryuzy 2025. 2. 17. 17:10

728x90

1. 포켓몬 분류 데이터셋

1. Pokemon Generation One 데이터셋

Pokemon Generation One 데이터셋은 포켓몬 시리즈의 첫 번째 세대(Generation One)에 등장하는 151마리 포켓몬의 이미지와 정보를 포함한 데이터셋입니다. 이 데이터셋은 주로 컴퓨터 비전과 머신러닝 작업(예: 이미지 분류, 객체 감지, 스타일 전환 등)에 활용됩니다. 이 데이터 세트에는 151개의 폴더가 포함되어 있으며, 각 폴더에는 각 포켓몬 1세대당 하나씩 포함되어 있으며 각 폴더에는 각 포켓몬당 60개의 이미지가 포함되어 있습니다. 총 10,000개 이상의 이미지가 있습니다.

2. Complete Pokemon Image Dataset

Complete Pokemon Image 데이터셋은 Kaggle에서 제공되는 데이터셋으로 모두 https://pokemondb.net/ 에서 스크랩된 것입니다. 포켓몬 세대 1부터 8까지의 2,500개 이상의 깨끗하게 라벨링된 이미지를 포함하고 있습니다. 이 데이터셋은 각 포켓몬의 다양한 형태와 자세를 포괄하며, 이미지 분류, 객체 인식, 생성 모델 등 다양한 딥러닝 및 머신러닝 작업에 유용하게 활용될 수 있습니다. 이미지들은 고품질로 제공되어 모델 학습 시 정확한 특징 추출에 도움이 됩니다.

!kaggle datasets download -d thedagger/pokemon-generation-one

!unzip -q /content/pokemon-generation-one.zip

!kaggle datasets download -d hlrhegemony/pokemon-image-dataset

!unzip -q /content/pokemon-image-dataset.zip

!mv dataset train # dataset의 데이터를 train 폴더(생성)로 이동

!rm -rf train/dataset

!mv images validation

train_labels = os.listdir('train')
print(train_labels)
print(len(train_labels))

val_labels = os.listdir('validation')
print(val_labels)
print(len(val_labels))

import shutil

for val_label in val_labels:
    if val_label not in train_labels:
        shutil.rmtree(os.path.join('validation', val_label))

val_labels = os.listdir('validation')
len(val_labels)

for train_label in train_labels:
    if train_label not in val_labels:
        shutil.rmtree(os.path.join('train', train_label))
        print(train_label, ' 삭제!')

val_labels = os.listdir('validation')
len(val_labels)

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

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor()
    ]),
    'validation': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])
}

image_datasets = {
    'train': datasets.ImageFolder('train', data_transforms['train']),
    'validation': datasets.ImageFolder('validation', data_transforms['validation'])
}

dataloaders = {
    'train': DataLoader(
        image_datasets['train'],
        batch_size=32,
        shuffle=True
    ),
    'validation': DataLoader(
        image_datasets['validation'],
        batch_size=32,
        shuffle=False
    )
}

len(image_datasets['train']), len(image_datasets['validation'])

imgs, labels = next(iter(dataloaders['train']))

fig, axes = plt.subplots(4, 8, figsize=(16, 8))

for ax, img, label in zip(axes.flatten(), imgs, labels):
    ax.imshow(img.permute(1, 2, 0))
    ax.set_title(label.item())
    ax.axis('off')

image_datasets['train'].classes[21]

2. EfficientNet

EfficientNet은 구글에서 개발한 합성곱 신경망(CNN) 모델로, 모델의 크기(depth), 너비(width), 해상도(resolution)를 균형 있게 조정하는 컴파운드 스케일링(compound scaling) 기법을 사용하여 효율성과 성능을 동시에 향상시킨 것이 특징입니다. 기존에는 모델의 크기를 단순히 깊게 만들거나(width나 resolution을 개별적으로 확장) 했지만, EfficientNet은 세 가지 요소를 균형 있게 확장하여 연산량을 최적화하면서도 높은 정확도를 유지할 수 있도록 설계되었습니다. 특히, EfficientNet-B0을 기본 모델로 하고, B1~B7까지 단계적으로 확장하여 다양한 컴퓨팅 리소스 환경에서 최적의 성능을 낼 수 있도록 제공됩니다. EfficientNet은 이미지 분류, 객체 탐지 등 다양한 컴퓨터 비전 작업에서 기존 모델들보다 더 적은 연산량으로도 뛰어난 성능을 보여주며, 실무에서도 널리 활용됩니다.

from torchvision.models import efficientnet_b4, EfficientNet_B4_Weights
from torchvision.models._api import WeightsEnum

model = efficientnet_b4(weights=EfficientNet_B4_Weights.IMAGENET1K_V1).to(device)
model

for param in model.parameters():
    param.requires_grad = False

model.classifier = nn.Sequential(
    nn.Linear(1792, 512),
    nn.ReLU(),
    nn.Linear(512, 147)
).to(device)

model

optimizer = optim.Adam(model.classifier.parameters(), lr=0.001)
epochs = 10

for epoch in range(epochs):
    for phase in ['train', 'validation']:
        if phase == 'train':
            model.train()
        else:
            model.eval()

        sum_losses = 0
        sum_accs = 0

        for x_batch, y_batch in dataloaders[phase]:
            x_batch = x_batch.to(device)
            y_batch = y_batch.to(device)

            y_pred = model(x_batch)
            loss = nn.CrossEntropyLoss()(y_pred, y_batch)

            if phase == 'train':
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            sum_losses = sum_losses + loss
            y_prob = nn.Softmax(1)(y_pred)
            y_pred_index = torch.argmax(y_prob, axis=1)
            acc = (y_batch == y_pred_index).float().sum() / len(y_batch) * 100
            sum_accs = sum_accs + acc

        avg_loss = sum_losses / len(dataloaders[phase])
        avg_acc = sum_accs / len(dataloaders[phase])
        print(f'{phase:10s}: Epoch {epoch+1:4d}/{epochs} Loss: {avg_loss:.4f} Accuracy: {avg_acc:.2f}%')

# 학습된 모델 파일을 저장
torch.save(model.state_dict(), 'model.pth')

model = models.efficientnet_b4().to(device)

model.classifier = nn.Sequential(
    nn.Linear(1792, 512),
    nn.ReLU(),
    nn.Linear(512, 147)
).to(device)

model

model.load_state_dict(torch.load('model.pth'))

model.eval()

from PIL import Image

img1 = Image.open('validation/Ditto/1.jpg')
img2 = Image.open('validation/Golem/4.jpg')

fig, axes = plt.subplots(1, 2, figsize=(12, 6))
axes[0].imshow(img1)
axes[0].axis('off')
axes[1].imshow(img2)
axes[1].axis('off')
plt.show()

img1_input = data_transforms['validation'](img1)
img2_input = data_transforms['validation'](img2)
print(img1_input.shape)
print(img2_input.shape)

test_batch = torch.stack([img1_input, img2_input])
test_batch = test_batch.to(device)
test_batch.shape

y_pred = model(test_batch)
y_pred

y_prob = nn.Softmax(1)(y_pred)
y_prob

probs, idx = torch.topk(y_prob, k=3)
print(probs)
print(idx)

fig, axes = plt.subplots(1, 2, figsize=(15, 6))

axes[0].set_title('{:.2f}% {}, {:.2f}% {}, {:.2f}% {}'.format(
    probs[0, 0] * 100,
    image_datasets['validation'].classes[idx[0, 0]],
    probs[0, 1] * 100,
    image_datasets['validation'].classes[idx[0, 1]],
    probs[0, 2] * 100,
    image_datasets['validation'].classes[idx[0, 2]],
))
axes[0].imshow(img1)
axes[0].axis('off')

axes[1].set_title('{:.2f}% {}, {:.2f}% {}, {:.2f}% {}'.format(
    probs[1, 0] * 100,
    image_datasets['validation'].classes[idx[1, 0]],
    probs[1, 1] * 100,
    image_datasets['validation'].classes[idx[1, 1]],
    probs[1, 2] * 100,
    image_datasets['validation'].classes[idx[1, 2]],
))
axes[1].imshow(img2)
axes[1].axis('off')

mypic = Image.open('본인사진.jpg')
plt.imshow(mypic)
plt.axis('off')

mypic_input = data_transforms['validation'](mypic)
print(mypic_input.shape)

mypic_input = mypic_input.unsqueeze(0).to(device)
print(mypic_input.shape)

y_pred = model(mypic_input)
y_pred

y_prob = nn.Softmax(1)(y_pred)
y_prob

probs, idx = torch.topk(y_prob, k=3)
print(probs)
print(idx)

plt.title('{:.2f}% {}, {:.2f}% {}, {:.2f}% {}'.format(
    probs[0, 0] * 100,
    image_datasets['validation'].classes[idx[0, 0]],
    probs[0, 1] * 100,
    image_datasets['validation'].classes[idx[0, 1]],
    probs[0, 2] * 100,
    image_datasets['validation'].classes[idx[0, 2]],
))
plt.imshow(mypic)
plt.axis('off')

728x90