1. Thiết kế Dataset tùy chỉnh

Để huấn luyện mô hình phân loại hình ảnh, ta cần tạo một dataset có cấu trúc phù hợp với PyTorch. Dataset cần tuân theo cấu trúc thư mục với mỗi thư mục con đại diện cho một lớp.

import os
from PIL import Image
from torch.utils.data import Dataset
import cv2

class ImageFolderDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        """
        Khởi tạo dataset tùy chỉnh
        :param root_dir: Thư mục gốc chứa dữ liệu
        :param transform: Các phép biến đổi dữ liệu
        """
        self.root_dir = root_dir
        self.transform = transform
        self.samples = []

        # Quét tất cả các thư mục con
        for class_idx, class_name in enumerate(os.listdir(root_dir)):
            class_path = os.path.join(root_dir, class_name)
            if os.path.isdir(class_path):
                for img_name in os.listdir(class_path):
                    img_path = os.path.join(class_path, img_name)
                    self.samples.append((img_path, class_idx))

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

    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        image = cv2.imread(img_path)
        if image is None:
            raise ValueError(f"Không thể tải ảnh từ: {img_path}")
        
        # Chuyển sang ảnh xám
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        image = Image.fromarray(image)
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

2. Huấn luyện và kiểm tra mô hình

Tiếp theo, ta xây dựng mô hình mạng nơ-ron fully connected và huấn luyện nó.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader

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

# Tiền xử lý dữ liệu
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
])

# Tải dataset
train_data = ImageFolderDataset(root_dir="data_set/training_set", transform=transform)
test_data = ImageFolderDataset(root_dir="data_set/test_set", transform=transform)

train_loader = DataLoader(train_data, batch_size=400, shuffle=True)
test_loader = DataLoader(test_data, batch_size=400, shuffle=False)


class FullyConnectedNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(28 * 28 * 1, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.network(x)
        return x


model = FullyConnectedNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)


def train_model(epochs=5):
    for epoch in range(epochs):
        total_loss = 0
        for images, labels in train_loader:
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        print(f"Epoch [{epoch + 1}/{epochs}], Loss: {total_loss:.4f}")


train_model(5)

3. Các vấn đề phát sinh

Trong quá trình triển khai, tôi nhận thấy một số hạn chế:

• Phân bố dữ liệu không cân bằng: Số lượng ảnh giữa các loại không đồng đều, ảnh hưởng đến quá trình huấn luyện.
• Xử lý ảnh không tối ưu: Sử dụng trực tiếp resize làm mất đi nhiều đặc điểm quan trọng của hình ảnh.
• Không lưu mô hình tốt nhất: Chỉ kiểm tra kết quả của epoch cuối cùng, không lưu lại trọng số tốt nhất.
• Số epoch chưa phù hợp: Cần theo dõi đường cong loss để xác định số epoch tối ưu.

4. Giải pháp cải tiến

4.1. Xử lý ảnh với padding giữ tỷ lệ

def maintain_aspect_ratio(image, target_size):
    """
    Điều chỉnh kích thước ảnh với padding giữ nguyên tỷ lệ
    :param image: Ảnh đầu vào (numpy array)
    :param target_size: Kích thước mục tiêu (vuông)
    :return: Ảnh đã xử lý
    """
    h, w = image.shape[:2]
    
    # Tính tỷ lệ scale
    scale = target_size / max(h, w)
    new_w, new_h = int(w * scale), int(h * scale)
    
    # Resize ảnh
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
    
    # Tính toán padding
    delta_w = target_size - new_w
    delta_h = target_size - new_h
    top, bottom = delta_h // 2, delta_h - delta_h // 2
    left, right = delta_w // 2, delta_w - delta_w // 2
    
    # Thêm padding màu đen
    color = [0, 0, 0]
    padded = cv2.copyMakeBorder(resized, top, bottom, left, right, 
                                 cv2.BORDER_CONSTANT, value=color)
    
    return padded

4.2. Lưu mô hình có loss thấp nhất

def train_model(epochs=100):
    best_loss = float('inf')
    
    for epoch in range(epochs):
        total_loss = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_loader.sampler)
        
        # Lưu mô hình có loss thấp nhất
        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save(model.state_dict(), "BestModel.pt")
            if epoch % 20 == 0:
                print(f"Lưu mô hình, loss = {avg_loss:.4f}")

5. Code hoàn chỉnh sau cải tiến

import os
from PIL import Image
from torch.utils.data import Dataset
import cv2
import numpy as np
import random
from torch.utils.data import DataLoader, SubsetRandomSampler
import torch.nn as nn
import torch
import torch.optim as optim
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter


def maintain_aspect_ratio(image, target_size):
    h, w = image.shape[:2]
    scale = target_size / max(h, w)
    new_w, new_h = int(w * scale), int(h * scale)
    resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
    
    delta_w = target_size - new_w
    delta_h = target_size - new_h
    top, bottom = delta_h // 2, delta_h - delta_h // 2
    left, right = delta_w // 2, delta_w - delta_w // 2
    
    color = [0, 0, 0]
    padded = cv2.copyMakeBorder(resized, top, bottom, left, right, 
                                 cv2.BORDER_CONSTANT, value=color)
    return padded


class ImageFolderDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.samples = []

        for class_idx, class_name in enumerate(os.listdir(root_dir)):
            class_path = os.path.join(root_dir, class_name)
            if os.path.isdir(class_path):
                for img_name in os.listdir(class_path):
                    img_path = os.path.join(class_path, img_name)
                    self.samples.append((img_path, class_idx))

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

    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        image = cv2.imread(img_path)
        if image is None:
            raise ValueError(f"Không thể tải ảnh: {img_path}")
        
        image = maintain_aspect_ratio(image, target_size=56)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        image = Image.fromarray(image)
        
        if self.transform:
            image = self.transform(image)
        
        return image, label


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
writer = SummaryWriter("logs/")

transform = transforms.Compose([
    transforms.ToTensor(),
])

data_set = ImageFolderDataset(root_dir="data_set", transform=transform)
data_index = list(range(len(data_set)))
random.shuffle(data_index)
split_point = int(0.9 * len(data_index))

train_sampler = SubsetRandomSampler(data_index[0:split_point])
test_sampler = SubsetRandomSampler(data_index[split_point:])

train_loader = DataLoader(data_set, batch_size=100, sampler=train_sampler)
test_loader = DataLoader(data_set, batch_size=100, sampler=test_sampler)


class FullyConnectedNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(56 * 56 * 1, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 6),
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.network(x)
        return x


model = FullyConnectedNet().to(device)
criterion = nn.CrossEntropyLoss(reduction="sum")
optimizer = optim.Adam(model.parameters(), lr=0.01)


def train_model(epochs=100):
    best_loss = float('inf')
    for epoch in range(epochs):
        total_loss = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_sampler)
        writer.add_scalar("Loss", avg_loss, epoch)
        
        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save(model.state_dict(), "BestModel.pt")
            if epoch % 20 == 0:
                print(f"Lưu mô hình, loss = {avg_loss:.4f}")


def evaluate():
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            model.load_state_dict(torch.load("BestModel.pt"))
            outputs = model(images)
            _, predictions = torch.max(outputs, dim=1)
            total += labels.size(0)
            correct += (predictions == labels).sum().item()
    
    print(f"Độ chính xác: {100 * correct / total:.2f}%")


train_model(100)
evaluate()

6. Dự đoán một hình ảnh

import torch
from PIL import Image
import torch.nn as nn
from torchvision import transforms

transform = transforms.Compose([
    transforms.Resize((56, 56)),
    transforms.ToTensor(),
])

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


class FullyConnectedNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(56 * 56 * 1, 1024),
            nn.ReLU(),
            nn.Linear(1024, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 6),
        )

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = self.network(x)
        return x


model = FullyConnectedNet().to(device)


def predict(image_path):
    image = Image.open(image_path).convert("L")
    image = transform(image)
    image = image.unsqueeze(0).to(device)
    
    model.load_state_dict(torch.load("BestModel.pt"))
    outputs = model(image)
    
    print(f"Điểm dự đoán: {outputs}")
    _, prediction = torch.max(outputs, dim=1)
    
    class_names = ["Sẻ", "Đại bàng", "Gõ kiến", "Hồng hạc", "Chim cánh cụt", "Chim lửa"]
    print(f"Dự đoán: {class_names[prediction.item()]}")


predict("test_image.jpg")

7. Kết quả và nhận xét

Sau khi cải tiến, mô hình đạt độ chính xác khoảng 51% trên tập test. Các loại chim có đặc điểm nổi bật như đại bàng, chim cánh cụt và chim lửa được dự đoán chính xác hơn. Tuy nhiên, các loại chim có đặc điểm ít rõ ràng như gõ kiến và hồng hạc vẫn còn khó phân loại chính xác.

Việc sử dụng phương pháp padding để giữ tỷ lệ ảnh và lưu lại mô hình có loss thấp nhất đã cải thiện đáng kể hiệu suất của mô hình so với cách tiếp cận ban đầu.