허깅페이스를 이용한 이미지 분류 Fine Tuning 방법 (전체 소스코드 있음)

허깅페이스로 이미지 분류 AI 만들기 - 콩 잎 질병 진단 프로젝트

=> 전체 소스코드는 맨 아래에 있습니다. 

=> 실습 동영상 : https://youtu.be/k7mCnlSfUFs

 

 

들어가며

최근 AI 개발에서 허깅페이스(Hugging Face)가 대세입니다. TensorFlow Hub보다 사용하기 쉽고, 최신 모델을 바로 활용할 수 있어서 실무에서도 많이 사용되고 있죠. 오늘은 허깅페이스를 이용해서 콩 잎 질병을 분류하는 이미지 분류 AI를 만들어보겠습니다.

프로젝트 개요

• 목표: 콩 잎 이미지를 보고 질병 종류 판별
• 데이터셋: AI-Lab-Makerere/beans (3개 클래스, 약 1,300장)
• 클래스: angular_leaf_spot, bean_rust, healthy
• 모델: ConvNeXt-tiny (사전학습된 모델)
• 방법: Fine-tuning (파인튜닝)
• 환경: Google Colab (무료 GPU 사용)
• 예상 정확도: 90% 이상

전체 흐름

이 프로젝트는 다음과 같은 순서로 진행됩니다:

1. 환경 설정 및 라이브러리 설치
2. 데이터셋 로드 및 확인
3. 데이터 시각화
4. 모델 선택 및 로드
5. 데이터 전처리
6. 학습 설정
7. 모델 학습 (Fine-tuning)
8. 성능 평가
9. 예측 및 테스트
10. 모델 저장

1. 환경 설정

Google Colab에서 작업하며, 필요한 라이브러리는 transformers, datasets, torch 등입니다. GPU가 할당되었는지 확인하는 것이 중요한데, GPU가 있으면 학습 시간이 3-5분 정도로 매우 빨라집니다.

2. 데이터셋 로드

허깅페이스의 가장 큰 장점은 load_dataset 함수 하나로 데이터를 불러올 수 있다는 점입니다.

Beans 데이터셋은 자동으로 train, validation, test 세 개의 split으로 나뉘어져 있어 편리합니다. Train 1,034장, Validation 133장, Test 128장으로 구성되어 있고, 각 클래스별 분포도 비교적 균등합니다.

3. 데이터 시각화

실제 데이터가 어떻게 생겼는지 확인하는 단계입니다. 콩 잎 사진들을 보면 질병에 따라 반점의 색깔과 패턴이 다른 것을 육안으로도 확인할 수 있습니다.

4. 모델 선택

여러 모델 옵션 중에서 선택할 수 있습니다:

• ConvNeXt-tiny: 속도와 정확도의 균형이 좋음 (추천)
• ConvNeXt-base: 더 높은 정확도, 조금 느림
• EfficientNet-B0: 가장 빠르고 경량
• ViT-base: Vision Transformer, 학습용으로 좋음
• MobileNet: 모바일 배포용

오늘은 ConvNeXt-tiny를 사용합니다. 이 모델은 ImageNet으로 사전학습되어 있어 이미 수백만 장의 이미지에서 특징을 추출하는 방법을 알고 있습니다.

모델을 로드할 때 중요한 파라미터:

• num_labels=3: 우리 데이터셋의 클래스 개수
• ignore_mismatched_sizes=True: 출력층 크기가 달라도 자동 조정
• id2label, label2id: 클래스 번호와 이름 매핑

5. 데이터 전처리

Processor가 자동으로 이미지를 전처리합니다:

• 크기 조정 (224x224로 통일)
• 정규화 (ImageNet 평균/표준편차 사용)
• 텐서 변환

이게 허깅페이스의 큰 장점입니다. TensorFlow에서는 이런 전처리를 일일이 코딩해야 하는데, 여기서는 Processor가 모델에 맞게 자동으로 처리해줍니다.

6. 학습 설정 (TrainingArguments)

학습 하이퍼파라미터를 설정하는 단계입니다. 주요 파라미터를 살펴보겠습니다:

Epoch 설정

• num_train_epochs=15: 데이터가 적으므로 epoch를 많이 설정합니다. 데이터가 많으면 5-10 정도면 충분합니다.

Batch Size

• per_device_train_batch_size=32: GPU 메모리에 따라 조정합니다. 메모리 부족 에러가 나면 16이나 8로 줄이세요.

Learning Rate

• learning_rate=5e-5: 0.00005로 매우 작습니다. 이건 파인튜닝의 핵심입니다. 사전학습된 모델을 조금씩만 조정해야 하므로 learning rate를 작게 설정합니다.

Warmup

• warmup_ratio=0.1: 처음 10%의 스텝 동안은 learning rate를 천천히 올립니다. 급격한 변화를 방지합니다.

평가 및 저장

• eval_strategy="epoch": 매 epoch마다 validation set으로 평가
• save_strategy="epoch": 매 epoch마다 체크포인트 저장
• load_best_model_at_end=True: 학습 끝나고 가장 좋았던 모델 자동 로드

최적화

• fp16=True: GPU 사용 시 혼합 정밀도 학습으로 메모리 절약 및 속도 향상

7. 모델 학습 (Fine-tuning)

Trainer를 생성하고 train() 메서드를 호출하면 학습이 시작됩니다.

Fine-tuning이란?
사전학습된 모델 전체를 우리 데이터에 맞게 조금씩 조정하는 것입니다. 이는 Feature Extraction(특성 추출)과 다릅니다:

• Feature Extraction: 사전학습 부분 고정, 마지막 분류층만 학습
• Fine-tuning: 사전학습 부분 포함, 전체 모델 학습 (우리가 한 것!)

Fine-tuning이 더 높은 정확도를 달성할 수 있지만, 조금 더 시간이 걸립니다.

학습 중에는:

• 매 10 step마다 loss 출력
• 매 epoch마다 validation accuracy 출력
• Loss가 점점 줄어들고 accuracy가 올라가는 것 확인

GPU 기준 3-5분, CPU 기준 15-30분 정도 소요됩니다.

8. 모델 평가

학습이 끝나면 두 가지 데이터셋으로 평가합니다:

Validation Set

• 학습 중에 모델 선택에 사용된 데이터
• 최종 모델이 이 데이터로 선택되었으므로 약간 낙관적일 수 있음

Test Set

• 학습 과정에서 한 번도 본 적 없는 데이터
• 실제 성능을 더 정확하게 반영

Classification Report
클래스별 precision, recall, f1-score를 보여줍니다:

• Precision: 모델이 A라고 예측한 것 중 실제로 A인 비율
• Recall: 실제 A인 것 중 모델이 A라고 맞춘 비율
• F1-score: Precision과 Recall의 조화평균

Confusion Matrix
어떤 클래스끼리 헷갈리는지 시각적으로 보여줍니다. 대각선이 진할수록 잘 분류한 것입니다.

9. 예측 함수

predict_bean_disease 함수는 세 가지 방식으로 이미지를 받습니다:

1. 테스트 데이터 인덱스: predict_bean_disease(0)
2. 파일 경로: predict_bean_disease("my_image.jpg")
3. URL: predict_bean_disease("https://example.com/image.jpg")

예측 결과는:

• 예측된 클래스명
• 확신도 (Confidence)
• 각 클래스별 확률
• 실제 정답과 비교 (테스트 데이터인 경우)

10. 모델 저장

두 가지를 저장합니다:

Model (모델)

• 학습된 가중치
• 모델 구조
• 예측 로직

Processor (전처리기)

• 이미지 전처리 방법
• 크기 조정, 정규화 설정

왜 둘 다 저장하나?

비유하자면:

• Processor = 주방에서 재료 손질하는 사람
• Model = 요리사

재료만 있어도 안 되고, 요리사만 있어도 안 됩니다. 둘 다 필요합니다.

나중에 모델을 불러올 때도 둘 다 필요합니다:

전체 소스 코드

# ============================================
# Beans Dataset 완전 실습 코드
# Dataset: AI-Lab-Makerere/beans
# 콩 잎 질병 분류 (3 classes)
# ============================================

# 설치
!pip install -q transformers datasets pillow torch torchvision scikit-learn matplotlib

# ============================================
# 1. 임포트
# ============================================

import torch
from transformers import (
    AutoImageProcessor, 
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer
)
from datasets import load_dataset
import numpy as np
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns

print(f"PyTorch version: {torch.__version__}")
print(f"GPU available: {'Yes' if torch.cuda.is_available() else 'No'}")
if torch.cuda.is_available():
    print(f"GPU name: {torch.cuda.get_device_name(0)}")

# ============================================
# 2. 데이터셋 로드
# ============================================

print("\nLoading Beans dataset...")
dataset = load_dataset("AI-Lab-Makerere/beans")

print(f"\nDataset loaded successfully!")
print(f"Dataset structure: {dataset}")
print(f"\nSplits:")
print(f"  - Train: {len(dataset['train'])} images")
print(f"  - Validation: {len(dataset['validation'])} images")
print(f"  - Test: {len(dataset['test'])} images")

# 클래스 정보
class_names = dataset['train'].features['labels'].names
num_classes = len(class_names)

print(f"\nClasses ({num_classes}):")
for i, name in enumerate(class_names):
    print(f"  {i}. {name}")

# 데이터 분포 확인
from collections import Counter

train_labels = [sample['labels'] for sample in dataset['train']]
val_labels = [sample['labels'] for sample in dataset['validation']]
test_labels = [sample['labels'] for sample in dataset['test']]

print(f"\nClass distribution:")
print(f"Train: {dict(Counter(train_labels))}")
print(f"Val: {dict(Counter(val_labels))}")
print(f"Test: {dict(Counter(test_labels))}")

# ============================================
# 3. 데이터 시각화
# ============================================

def show_samples(dataset, num_samples=6):
    """샘플 이미지 시각화"""
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    axes = axes.flatten()

    indices = np.random.choice(len(dataset), num_samples, replace=False)

    for i, idx in enumerate(indices):
        sample = dataset[int(idx)]
        image = sample['image']
        label = class_names[sample['labels']]

        axes[i].imshow(image)
        axes[i].set_title(f'{label}', fontsize=12, fontweight='bold')
        axes[i].axis('off')

    plt.tight_layout()
    plt.show()

print("\nShowing sample images from training set:")
show_samples(dataset['train'], num_samples=6)

# ============================================
# 4. 모델 선택 및 로드
# ============================================

# 여러 모델 옵션
MODEL_OPTIONS = {
    "convnext-tiny": "facebook/convnext-tiny-224",      # 추천! 빠르고 정확
    "convnext-base": "facebook/convnext-base-224",      # 더 높은 정확도
    "efficientnet-b0": "google/efficientnet-b0",        # 경량
    "vit-base": "google/vit-base-patch16-224",          # Vision Transformer
    "mobilenet": "google/mobilenet_v2_1.0_224",         # 가장 빠름
}

# 사용할 모델 선택
selected_model = "convnext-tiny"  # 원하는 모델로 변경 가능
model_name = MODEL_OPTIONS[selected_model]

print(f"\nSelected model: {model_name}")

# Image Processor 로드
processor = AutoImageProcessor.from_pretrained(model_name)
print(f"Processor loaded")

# 모델 로드
model = AutoModelForImageClassification.from_pretrained(
    model_name,
    num_labels=num_classes,
    id2label={i: label for i, label in enumerate(class_names)},
    label2id={label: i for i, label in enumerate(class_names)},
    ignore_mismatched_sizes=True
)
print(f"Model loaded with {num_classes} output classes")

# 모델 파라미터 수 확인
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model parameters:")
print(f"  - Total: {total_params:,}")
print(f"  - Trainable: {trainable_params:,}")

# ============================================
# 5. 데이터 전처리
# ============================================

def transform(examples):
    """이미지 전처리 함수"""
    images = examples['image']
    inputs = processor(images, return_tensors='pt')
    inputs['labels'] = examples['labels']
    return inputs

print("\nApplying transformations...")

# Transform 적용
dataset['train'].set_transform(transform)
dataset['validation'].set_transform(transform)
dataset['test'].set_transform(transform)

print("Transformations applied to all splits")

# ============================================
# 6. 평가 메트릭 정의
# ============================================

def compute_metrics(eval_pred):
    """평가 메트릭 계산"""
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)

    accuracy = accuracy_score(labels, predictions)

    return {
        'accuracy': accuracy,
    }

# ============================================
# 7. 학습 설정
# ============================================

training_args = TrainingArguments(
    output_dir='./beans-classifier-results',

    # 학습 설정
    num_train_epochs=15,                      # Epoch 수
    per_device_train_batch_size=32,           # 배치 크기
    per_device_eval_batch_size=32,

    # 옵티마이저 설정
    learning_rate=5e-5,                       # 학습률
    warmup_ratio=0.1,                         # Warmup 비율
    weight_decay=0.01,                        # Weight decay

    # 평가 및 저장
    eval_strategy="epoch",                    # 매 epoch마다 평가
    save_strategy="epoch",                    # 매 epoch마다 저장
    load_best_model_at_end=True,              # 최고 성능 모델 로드
    metric_for_best_model="accuracy",         # 최적 모델 선택 기준

    # 로깅
    logging_dir='./logs',
    logging_steps=10,

    # 기타
    remove_unused_columns=False,
    push_to_hub=False,
    report_to="none",

    # GPU 최적화
    fp16=torch.cuda.is_available(),
)

print("\nTraining arguments configured")

# ============================================
# 8. Trainer 생성
# ============================================

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset['train'],
    eval_dataset=dataset['validation'],
    compute_metrics=compute_metrics,
)

print("Trainer created")

# ============================================
# 9. 학습 시작
# ============================================

print("\n" + "="*60)
print("TRAINING STARTED")
print("="*60)

train_result = trainer.train()

print("\n" + "="*60)
print("TRAINING COMPLETED")
print("="*60)

# 학습 결과 출력
print(f"\nTraining Results:")
print(f"  - Training time: {train_result.metrics['train_runtime']:.2f} seconds")
print(f"  - Training samples/second: {train_result.metrics['train_samples_per_second']:.2f}")
print(f"  - Final training loss: {train_result.metrics['train_loss']:.4f}")

# ============================================
# 10. 모델 평가
# ============================================

print("\n" + "="*60)
print("EVALUATION")
print("="*60)

# Validation set 평가
print("\n[1] Validation Set:")
val_results = trainer.evaluate(dataset['validation'])
print(f"   Accuracy: {val_results['eval_accuracy']:.2%}")
print(f"   Loss: {val_results['eval_loss']:.4f}")

# Test set 평가
print("\n[2] Test Set:")
test_results = trainer.evaluate(dataset['test'])
print(f"   Accuracy: {test_results['eval_accuracy']:.2%}")
print(f"   Loss: {test_results['eval_loss']:.4f}")

# ============================================
# 11. 상세 분석 - Confusion Matrix
# ============================================

print("\nGenerating detailed classification report...")

# 테스트 셋에 대한 예측
predictions = trainer.predict(dataset['test'])
y_pred = np.argmax(predictions.predictions, axis=1)
y_true = predictions.label_ids

# Classification Report
print("\n" + "="*60)
print("Classification Report (Test Set)")
print("="*60)
print(classification_report(
    y_true, 
    y_pred, 
    target_names=class_names,
    digits=4
))

# Confusion Matrix
cm = confusion_matrix(y_true, y_pred)

plt.figure(figsize=(10, 8))
sns.heatmap(
    cm, 
    annot=True, 
    fmt='d', 
    cmap='Blues',
    xticklabels=class_names,
    yticklabels=class_names
)
plt.title('Confusion Matrix - Test Set', fontsize=14, fontweight='bold')
plt.ylabel('True Label')
plt.xlabel('Predicted Label')
plt.tight_layout()
plt.show()

# ============================================
# 12. 예측 함수
# ============================================

def predict_bean_disease(image_input, show_image=True):
    """
    콩 잎 질병 예측 함수

    Parameters:
    -----------
    image_input : int, str, or PIL.Image
        - int: dataset['test']의 인덱스
        - str: 이미지 파일 경로 또는 URL
        - PIL.Image: PIL Image 객체
    show_image : bool
        이미지를 출력할지 여부

    Returns:
    --------
    predicted_label : str
        예측된 클래스명
    confidence : float
        예측 확신도 (0-1)
    """

    # 이미지 로드
    true_label = None
    if isinstance(image_input, int):
        # Dataset index
        sample = dataset['test'][image_input]
        image = sample['image']
        true_label = class_names[sample['labels']]
    elif isinstance(image_input, str):
        # 파일 경로 또는 URL
        if image_input.startswith('http'):
            import requests
            from io import BytesIO
            response = requests.get(image_input)
            image = Image.open(BytesIO(response.content)).convert('RGB')
        else:
            image = Image.open(image_input).convert('RGB')
    else:
        # PIL Image
        image = image_input

    # 전처리
    inputs = processor(images=image, return_tensors="pt")

    # GPU로 이동 (가능한 경우)
    if torch.cuda.is_available():
        inputs = {k: v.to('cuda') for k, v in inputs.items()}
        model.to('cuda')

    # 예측
    model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=-1)[0]
        predicted_idx = probs.argmax().item()

    predicted_label = class_names[predicted_idx]
    confidence = probs[predicted_idx].item()

    # 결과 출력
    print(f"\n{'='*50}")
    print(f"PREDICTION RESULTS")
    print(f"{'='*50}")

    if true_label:
        is_correct = predicted_label == true_label
        print(f"True Label:      {true_label}")
        print(f"Predicted Label: {predicted_label}")
        print(f"Confidence:      {confidence:.1%}")
        print(f"Result:          {'CORRECT' if is_correct else 'WRONG'}")
    else:
        print(f"Predicted Label: {predicted_label}")
        print(f"Confidence:      {confidence:.1%}")

    print(f"\nAll Class Probabilities:")
    for i, (class_name, prob) in enumerate(zip(class_names, probs)):
        print(f"  {class_name:20s}: {prob:6.1%}")

    # 이미지 표시
    if show_image:
        plt.figure(figsize=(8, 8))
        plt.imshow(image)
        title = f"Predicted: {predicted_label} ({confidence:.1%})"
        if true_label:
            title += f"\nTrue: {true_label}"
        plt.title(title, fontsize=14, fontweight='bold')
        plt.axis('off')
        plt.tight_layout()
        plt.show()

    return predicted_label, confidence

# ============================================
# 13. 테스트 예측 실행
# ============================================

print("\n" + "="*60)
print("TESTING PREDICTIONS")
print("="*60)

# 각 클래스별로 하나씩 예측
print("\n예측 테스트: 각 클래스별 샘플")

for class_idx, class_name in enumerate(class_names):
    # 해당 클래스의 샘플 찾기
    test_samples_of_class = [
        i for i, sample in enumerate(dataset['test']) 
        if sample['labels'] == class_idx
    ]

    if test_samples_of_class:
        sample_idx = test_samples_of_class[0]
        print(f"\n{'='*60}")
        print(f"Testing class: {class_name}")
        predict_bean_disease(sample_idx, show_image=True)

# 랜덤 샘플 추가 테스트
print("\n" + "="*60)
print("Random Sample Predictions")
print("="*60)

import random
random_indices = random.sample(range(len(dataset['test'])), min(3, len(dataset['test'])))

for idx in random_indices:
    predict_bean_disease(idx, show_image=True)

# ============================================
# 14. 모델 저장
# ============================================

save_path = "./beans_disease_classifier"
print(f"\nSaving model to {save_path}...")

trainer.save_model(save_path)
processor.save_pretrained(save_path)

print(f"Model saved successfully!")

# 저장된 파일 확인
import os
saved_files = os.listdir(save_path)
print(f"\nSaved files: {saved_files}")

# ============================================
# 15. 모델 재로드 테스트
# ============================================

print("\nTesting model reload...")

loaded_processor = AutoImageProcessor.from_pretrained(save_path)
loaded_model = AutoModelForImageClassification.from_pretrained(save_path)

print("Model reloaded successfully!")

# ============================================
# 16. 최종 요약
# ============================================

print("\n" + "="*60)
print("FINAL SUMMARY")
print("="*60)

print(f"\nDataset:")
print(f"  - Name: AI-Lab-Makerere/beans")
print(f"  - Classes: {num_classes} ({', '.join(class_names)})")
print(f"  - Train samples: {len(dataset['train'])}")
print(f"  - Val samples: {len(dataset['validation'])}")
print(f"  - Test samples: {len(dataset['test'])}")

print(f"\nModel:")
print(f"  - Architecture: {model_name}")
print(f"  - Total parameters: {total_params:,}")
print(f"  - Trainable parameters: {trainable_params:,}")

print(f"\nPerformance:")
print(f"  - Validation Accuracy: {val_results['eval_accuracy']:.2%}")
print(f"  - Test Accuracy: {test_results['eval_accuracy']:.2%}")

print(f"\nSaved:")
print(f"  - Model path: {save_path}")

print("\n" + "="*60)
print("ALL DONE!")
print("="*60)

# ============================================
# 17. 사용 가이드
# ============================================

print(f"\nHow to use this model:")
print(f"""
# 새 이미지로 예측하기:
predict_bean_disease('path/to/your/image.jpg')

# 또는 URL로:
predict_bean_disease('https://example.com/bean_leaf.jpg')

# 또는 테스트 셋 인덱스로:
predict_bean_disease(42)
""")