[AICE PROFESSIONAL 준비] 토마토 질병 분류

캐글의 [Tomato Leaf Disease Detection]을 이용하여 토마토 질병 분류 모델을 만들어봤다.

 

https://www.kaggle.com/datasets/kaustubhb999/tomatoleaf

Tomato leaf disease detection

 

데이터 로드 및 확인

import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.metrics import classification_report
from sklearn.decomposition import PCA

def load_tomato_data(data_dir):
    images = []
    labels = []
    classes = os.listdir(data_dir)

    for idx, cls in enumerate(classes):
        path = os.path.join(data_dir, cls)
        for img_name in os.listdir(path):
            img = cv2.imread(os.path.join(path, img_name))
            if img is not None:
                img = cv2.resize(img, (64, 64))
                images.append(img.flatten())
                labels.append(idx)

    return np.array(images), np.array(labels), classes
    
X, y, class_names = load_tomato_data("./data/tomato/train")
class_names

 

# 데이터 스케일링
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train = X_train / 255.0
X_test = X_test / 255.0

 

# 차원 축소
pca = PCA(n_components=100)
X_train_pca = pca.fit_transform(X_train)
X_test_pca = pca.transform(X_test)

 

# 모델 학습
svc_model = SVC(kernel="rbf", C=10, gamma="scale")
svc_model.fit(X_train_pca, y_train)

 

# 평가
y_pred = svc_model.predict(X_test_pca)
print(classification_report(y_test, y_pred, target_names=class_names))

 

# 검증 데이터로 모델 성능 확인
X_val, y_val, class_names = load_tomato_data("./data/tomato/val")
X_val = X_val / 255.0
X_val_pca = pca.transform(X_val)
y_val_pred = svc_model.predict(X_val_pca)

 

print(classification_report(y_val, y_val_pred, target_names=class_names))

 

 

딥러닝으로 전환

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, random_split
import os

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

 

# 데이터 전처리
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]
        )
    ]
)

 

# 데이터셋 로드
data_dir = "./data/tomato/train"

# 폴더 구조별로 자동 라벨링
full_dataset = datasets.ImageFolder(root=data_dir, transform=transform)
print(f"이미지 개수: {len(full_dataset)}")
print(f"클래스 종류: {full_dataset.classes}")

 

# 데이터 분할
train_size = int(0.8 * len(full_dataset))
test_size = len(full_dataset) - train_size
train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size])

 

# 배치 단위로 불러오는 로더 생성
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

 

# 모델 불러오기 (사전 학습된 가중치 가져오기)
model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT)

 

# 모든 파라미터의 기울이 계산을 오프
for param in model.parameters():
    param.requires_grad = False

 

# 분류기
num_features = model.classifier[3].in_features
model.classifier[3] = nn.Linear(num_features, len(full_dataset.classes)) # 클래스의 개수에 맞춰 셋팅

 

# 모델을 이용가능한 디바이스로 이동
model = model.to(device)

 

# 손실 함수 및 최적화 기법
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

 

# 모델 학습
num_epochs = 5

print("===== 학습 시작 =====")
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad() # 기존에 계산된 기울기 초기화
        outputs = model(images) # 모델에 이미지를 넣고 예측값 도출
        loss = criterion(outputs, labels) # 실제값과 예측값 차이
        loss.backward() # 역전파하여 기울기 계산
        optimizer.step() # 가중치 업데이트

        running_loss += loss.item()

    print(f"epoch [{epoch + 1}], loss: {running_loss / len(train_loader):.4f}")

 

# 평가
model.eval()
correct = 0
total = 0

with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)

        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"테스트 데이터셋 정확도: {accuracy:.2f}%")

 

 

모델 저장

torch.save(model.state_dict(), "tomato_model.pth")
print("모델이 저장되었습니다.")

 

 

가상으로 만든 데이터를 이용하여 예측하기

import pandas as pd
from PIL import Image

csv_file = "./data/real_data.csv"
output_file = "result.csv"
class_names = full_dataset.classes

df = pd.read_csv(csv_file)

 

df.info()

 

df_drop = df.drop(labels=["Unnamed: 0"], axis=1)

 

# 모델을 평가모드로 전환
model.eval()

result = []

with torch.no_grad():
    for idx, row in df_drop.iterrows():
        img_path = row["path"] # 이미지 경로가 담긴 컬럼명

        # 이미지 로드 및 전처리
        image = Image.open(img_path).convert("RGB")
        image = transform(image).unsqueeze(0).to(device) # 배치 차원 추가

        # 예측
        outputs = model(image)
        _, predicted_idx = torch.max(outputs, 1)

        # 클래스 라벨 매칭
        predicted_label = class_names[predicted_idx.item()]
        result.append(predicted_label)
        print(f"{img_path} -> {predicted_label}")

 

# 결과 저장
df["predicted_class"] = result
df.to_csv(output_file, index=False)
print(f"분류 완료. 결과가 {output_file} 에 저장되었습니다.")