✨博客主页：王乐予🎈
✨年轻人要：Living for the moment（活在当下）！💪
🏆推荐专栏：【图像处理】【千锤百炼Python】【深度学习】【排序算法】

在图像分类领域，可能会遇到需要确定对象的多个属性的场景。例如，这些可以是类别、颜色、大小等。与通常的图像分类相比，此任务的输出将包含 2 个或更多属性。

在本教程中，我们将重点讨论一个问题，即我们事先知道属性的数量。此类任务称为多输出分类。事实上，这是多标签分类的一种特例，还可以预测多个属性，但它们的数量可能因样本而异。

本文程序已解耦，可当做通用型多标签图像分类框架使用。

数据集下载地址：Fashion-Product-Images

😺一、数据集介绍

我们将使用时尚产品图片数据集。它包含超过 44 000 张衣服和配饰图片，每张图片有 9 个标签。

从 kaggle 上下载到数据集后解压可以一个文件夹和一个csv表格，分别是images和styles.csv。

其中images里存放了数据集中所有的图片。

styles.csv中写入了图片的相关信息，包括 id（图片名称）、gender（性别）、masterCategory（主要类别）、subCategory（二级类别）、articleType（服装类型）、baseColour（描述性颜色）、season（季节）、year（年份）、usage（使用说明）、productDisplayName（品牌名称）。

😺二、工程文件夹目录

工程文件夹目录如下，每个py文件具有不同的功能，这么写的好处是未来修改程序更加方便，而且每个py程序都没有很长。如果全部写到一个py程序里，则会显得很臃肿，修改起来也不轻松。

对每个文件的解释如下：

• checkpoints：存放训练的模型权重；
• datasets：存放数据集。并对数据集划分；
• logs：存放训练日志。包括训练、验证时候的损失与精度情况；
• option.py：存放整个工程下需要用到的所有参数；
• utils.py：存放各种函数。包括模型保存、模型加载和损失函数等；
• split_data.py：划分数据集；
• model.py：构建神经网络模型；
• train.py：训练模型；
• predict.py：评估训练模型。

😺三、option.py

import argparse


def get_args():
    parser = argparse.ArgumentParser(description='ALL ARGS')
    parser.add_argument('--device', type=str, default='cuda', help='cuda or cpu')
    parser.add_argument('--start_epoch', type=int, default=0, help='start epoch')
    parser.add_argument('--epochs', type=int, default=100, help='Total Training Times')
    parser.add_argument('--batch_size', type=int, default=32, help='input batch size')
    parser.add_argument('--num_workers', type=int, default=0, help='number of processes to handle dataset loading')
    parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate for adam')
    parser.add_argument('--datasets_path', type=str, default='./datasets/', help='Path to the dataset')
    parser.add_argument('--image_path', type=str, default='./datasets/images', help='Path to the style image')
    parser.add_argument('--original_csv_path', type=str, default='./datasets/styles.csv', help='Original csv file dir')
    parser.add_argument('--train_csv_path', type=str, default='./datasets/train.csv', help='train csv file dir')
    parser.add_argument('--val_csv_path', type=str, default='./datasets/val.csv', help='val csv file dir')
    parser.add_argument('--log_dir', type=str, default='./logs/', help='log dir')
    parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints/', help='checkpoints dir')
    parser.add_argument('--checkpoint', type=str, default='./checkpoints/2024-05-24_13-50/checkpoint-000002.pth', help='choose a checkpoint to predict')
    parser.add_argument('--predict_image_path', type=str, default='./datasets/images/1163.jpg', help='show ground truth')

    return parser.parse_args()

😺四、split_data.py

由于数据集的各个属性严重不均衡，为简单起见，在本教程中仅使用三个标签：gender、articleType 和 baseColour

import csv
import os
import numpy as np
from PIL import Image
from tqdm import tqdm
from option import get_args


def save_csv(data, path, fieldnames=['image_path', 'gender', 'articleType', 'baseColour']):
    with open(path, 'w', newline='') as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        for row in data:
            writer.writerow(dict(zip(fieldnames, row)))


if __name__ == '__main__':

    args = get_args()
    input_folder = args.datasets_path
    output_folder = args.datasets_path

    annotation = args.original_csv_path

    all_data = []
    with open(annotation) as csv_file:
        reader = csv.DictReader(csv_file)
        for row in tqdm(reader, total=reader.line_num):
            img_id = row['id']

            # only three attributes are used: gender articleType、baseColour
            gender = row['gender']
            articleType = row['articleType']
            baseColour = row['baseColour']
            img_name = os.path.join(input_folder, 'images', str(img_id) + '.jpg')

            # Determine if the image exists
            if os.path.exists(img_name):

                # Check if the image is 80 * 60 size and if it is in RGB format
                img = Image.open(img_name)
                if img.size == (60, 80) and img.mode == "RGB":
                    all_data.append([img_name, gender, articleType, baseColour])

    np.random.seed(42)

    all_data = np.asarray(all_data)

    # Randomly select 40000 data points
    inds = np.random.choice(40000, 40000, replace=False)
    
    # Divide training and validation sets
    save_csv(all_data[inds][:32000], args.train_csv_path)
    save_csv(all_data[inds][32000:40000], args.val_csv_path)

😺五、dataset.py

该代码实现了两个类，AttributesDataset用于处理属性标签，FashionDataset类继承自Dataset类，用于处理带有图片路径和属性标签的数据集。关键地方的解释在代码中已经进行了注释。

get_mean_and_std函数用于获取数据集图像的均值与标准差

import csv
import numpy as np
from PIL import Image
import os
from torch.utils.data import Dataset
from torchvision import transforms
from option import get_args

args = get_args()


mean = [0.85418772, 0.83673165, 0.83065592]
std = [0.25331535, 0.26539705, 0.26877365]


class AttributesDataset():
    def __init__(self, annotation_path):
        color_labels = []
        gender_labels = []
        article_labels = []

        with open(annotation_path) as f:
            reader = csv.DictReader(f)
            for row in reader:
                color_labels.append(row['baseColour'])
                gender_labels.append(row['gender'])
                article_labels.append(row['articleType'])

        # Remove duplicate values to obtain a unique label set
        self.color_labels = np.unique(color_labels)
        self.gender_labels = np.unique(gender_labels)
        self.article_labels = np.unique(article_labels)

        # Calculate the number of categories for each label
        self.num_colors = len(self.color_labels)
        self.num_genders = len(self.gender_labels)
        self.num_articles = len(self.article_labels)

        # Create label mapping: Create two dictionaries: one from label ID to label name, and the other from label name to label ID.
        # Mapping results：self.gender_name_to_id：{'Boys': 0, 'Girls': 1, 'Men': 2, 'Unisex': 3, 'Women': 4}
        # Mapping results.gender_id_to_name：{0: 'Boys', 1: 'Girls', 2: 'Men', 3: 'Unisex', 4: 'Women'}
        self.color_id_to_name = dict(zip(range(len(self.color_labels)), self.color_labels))
        self.color_name_to_id = dict(zip(self.color_labels, range(len(self.color_labels))))

        self.gender_id_to_name = dict(zip(range(len(self.gender_labels)), self.gender_labels))
        self.gender_name_to_id = dict(zip(self.gender_labels, range(len(self.gender_labels))))

        self.article_id_to_name = dict(zip(range(len(self.article_labels)), self.article_labels))
        self.article_name_to_id = dict(zip(self.article_labels, range(len(self.article_labels))))


class FashionDataset(Dataset):
    def __init__(self, annotation_path, attributes, transform=None):
        super().__init__()

        self.transform = transform
        self.attr = attributes

        # Initialize a list to store the image path and corresponding labels of the dataset
        self.data = []
        self.color_labels = []
        self.gender_labels = []
        self.article_labels = []

        # Read data from a CSV file and store the image path and corresponding labels in a list
        with open(annotation_path) as f:
            reader = csv.DictReader(f)
            for row in reader:
                self.data.append(row['image_path'])
                self.color_labels.append(self.attr.color_name_to_id[row['baseColour']])
                self.gender_labels.append(self.attr.gender_name_to_id[row['gender']])
                self.article_labels.append(self.attr.article_name_to_id[row['articleType']])
        
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):

        img_path = self.data[idx]
        img = Image.open(img_path)

        if self.transform:
            img = self.transform(img)

        dict_data = {
            'img': img,
            'labels': {
                'color_labels': self.color_labels[idx],
                'gender_labels': self.gender_labels[idx],
                'article_labels': self.article_labels[idx]
            }
        }
        return dict_data

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
    ])

val_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
    ])



# Calculate the mean and variance of all images in the dataset
def get_mean_and_std(image_paths, transform):  
    
    # Initialize the accumulator of mean and variance
    means = np.zeros((3,))  
    stds = np.zeros((3,))  
    count = 0  
  
    for image_path in image_paths:   
        image = Image.open(image_path).convert('RGB')   
        image_tensor = transform(image).unsqueeze(0)  
  
        image_array = image_tensor.numpy()  
  
        # Calculate the mean and variance of the image
        batch_mean = np.mean(image_array, axis=(0, 2, 3))  
        batch_var = np.var(image_array, axis=(0, 2, 3))  
  
        # Accumulate to the total
        means += batch_mean  
        stds += batch_var  
        count += 1  
  
    # Calculate the mean and standard deviation of the entire dataset
    means /= count  
    stds = np.sqrt(stds / count)  
  
    return means, stds  


# Calculate the mean and variance of the dataset
if __name__ == '__main__':

    mena_std_transform = transforms.Compose([transforms.ToTensor()])

    image_path = []
    for root, _, files in os.walk(args.image_path):
        for file in files:
            if os.path.splitext(file)[1].lower() in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif'):
                image_path.append(os.path.join(root, file))

    means, stds = get_mean_and_std(image_path, mena_std_transform)
    print("Calculated mean and standard deviation:=========>") 
    print("Mean:", means)  
    print("Std:", stds)

😺六、model.py

该代码用来创建网络模型，需要注意的是最后使用了三个分类头对三个属性进行分类。

import torch
import torch.nn as nn
import torchvision.models as models


class MultiOutputModel(nn.Module):
    def __init__(self, n_color_classes, n_gender_classes, n_article_classes):
        super().__init__()

        self.base_model = models.mobilenet_v2().features
        last_channel = models.mobilenet_v2().last_channel

        self.pool = nn.AdaptiveAvgPool2d((1, 1))

        # Create three independent classifiers for predicting three categories
        self.color = nn.Sequential(nn.Dropout(p=0.2), nn.Linear(in_features=last_channel, out_features=n_color_classes))
        self.gender = nn.Sequential(nn.Dropout(p=0.2), nn.Linear(in_features=last_channel, out_features=n_gender_classes))
        self.article = nn.Sequential(nn.Dropout(p=0.2), nn.Linear(in_features=last_channel, out_features=n_article_classes))

    def forward(self, x):
        x = self.base_model(x)
        x = self.pool(x)

        x = torch.flatten(x, 1)

        return {
            'color': self.color(x),
            'gender': self.gender(x),
            'article': self.article(x)
        }

😺七、utils.py

对utils.py中各函数的解释：

• get_cur_time：获取当前时间。
• checkpoint_save：保存模型。
• checkpoint_load：加载模型。
• get_loss：定义损失函数。
• calculate_metrics：计算精度。

import os
from datetime import datetime
import warnings
from sklearn.metrics import balanced_accuracy_score
import torch
import torch.nn.functional as F


# Get the current date and time and format it as a string
def get_cur_time():
    return datetime.strftime(datetime.now(), '%Y-%m-%d_%H-%M')


def checkpoint_save(model, name, epoch):
    f = os.path.join(name, 'checkpoint-{:06d}.pth'.format(epoch))
    torch.save(model, f)
    print('Saved checkpoint:', f)

# Load Checkpoints
def checkpoint_load(model, name):
    print('Restoring checkpoint: {}'.format(name))
    model = torch.load(name, map_location='cpu')
    epoch = int(os.path.splitext(os.path.basename(name))[0].split('-')[1])
    return model, epoch

def get_loss(net_output, ground_truth):

    color_loss = F.cross_entropy(net_output['color'], ground_truth['color_labels'])
    gender_loss = F.cross_entropy(net_output['gender'], ground_truth['gender_labels'])
    article_loss = F.cross_entropy(net_output['article'], ground_truth['article_labels'])

    loss = color_loss + gender_loss + article_loss
    
    return loss, {'color': color_loss, 'gender': gender_loss, 'article': article_loss}


def calculate_metrics(output, target):
    _, predicted_color = output['color'].cpu().max(1)
    gt_color = target['color_labels'].cpu()

    _, predicted_gender = output['gender'].cpu().max(1)
    gt_gender = target['gender_labels'].cpu()

    _, predicted_article = output['article'].cpu().max(1)
    gt_article = target['article_labels'].cpu()

    with warnings.catch_warnings():  # sklearn may produce a warning when processing zero row in confusion matrix
        warnings.simplefilter("ignore")
        accuracy_color = balanced_accuracy_score(y_true=gt_color.numpy(), y_pred=predicted_color.numpy())
        accuracy_gender = balanced_accuracy_score(y_true=gt_gender.numpy(), y_pred=predicted_gender.numpy())
        accuracy_article = balanced_accuracy_score(y_true=gt_article.numpy(), y_pred=predicted_article.numpy())

    return accuracy_color, accuracy_gender, accuracy_article

😺八、train.py

该程序用于模型训练。

程序记录了训练日志，可以启动tensorboard观察训练过程（需要改成自己的路径）：
tensorboard --logdir=logs/2024-05-24_15-16

程序还添加了学习率衰减的训练策略。

程序使用tqdm库用于在终端可视化训练时间。

# Start Tensorboard：tensorboard --logdir=logs/2024-05-24_15-16
import os
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from dataset import AttributesDataset, FashionDataset, train_transform, val_transform
from model import MultiOutputModel
from utils import get_loss, get_cur_time, checkpoint_save
from predict import calculate_metrics, validate
from option import get_args

args = get_args()

# Initial parameters
start_epoch = args.start_epoch
N_epochs = args.epochs
batch_size = args.batch_size
num_workers = args.num_workers
batch_size = args.batch_size
device = args.device

# Initial paths
original_csv_path = args.original_csv_path
train_csv_path = args.train_csv_path
val_csv_path = args.val_csv_path
log_dir = args.log_dir
checkpoint_dir = args.checkpoint_dir

# Load attribute classes, The attributes contain labels and mappings for three categories
attributes = AttributesDataset(original_csv_path)

# Load Dataset
train_dataset = FashionDataset(train_csv_path, attributes, train_transform)
val_dataset = FashionDataset(val_csv_path, attributes, val_transform)

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

# Load model
model = MultiOutputModel(n_color_classes=attributes.num_colors,
                            n_gender_classes=attributes.num_genders,
                            n_article_classes=attributes.num_articles)
model.to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
sch = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9)    # Add learning rate decay

logdir = os.path.join(log_dir, get_cur_time())
savedir = os.path.join(checkpoint_dir, get_cur_time())

os.makedirs(logdir, exist_ok=True)
os.makedirs(savedir, exist_ok=True)


logger = SummaryWriter(logdir)

n_train_samples = len(train_dataloader)


if __name__ == '__main__':

    for epoch in range(start_epoch, N_epochs):

        # Initialize training loss and accuracy for each category
        total_loss, color_loss, gender_loss, article_loss = 0, 0, 0, 0
        accuracy_color, accuracy_gender, accuracy_article = 0, 0, 0

        # Create a tqdm instance to visualize training progress
        pbar = tqdm(total=len(train_dataset), desc='Training', unit='img')

        for batch in train_dataloader:
            pbar.update(train_dataloader.batch_size)    # Update progress bar

            optimizer.zero_grad()

            img = batch['img']
            target_labels = batch['labels']
            target_labels = {t: target_labels[t].to(device) for t in target_labels}
            output = model(img.to(device))

            # Calculate losses
            loss_train, losses_train = get_loss(output, target_labels)
            total_loss += loss_train.item()
            color_loss += losses_train['color']
            gender_loss += losses_train['gender']
            article_loss += losses_train['article']

            # Calculation accuracy
            batch_accuracy_color, batch_accuracy_gender, batch_accuracy_article = calculate_metrics(output, target_labels)

            accuracy_color += batch_accuracy_color
            accuracy_gender += batch_accuracy_gender
            accuracy_article += batch_accuracy_article

            loss_train.backward()
        sch.step()

        # Print epoch, total loss, loss for each category, accuracy for each category
        print("epoch {:2d}, total_loss: {:.4f}, color_loss: {:.4f}, gender_loss: {:.4f}, article_loss: {:.4f}, color_acc: {:.4f}, gender_acc: {:.4f}, article_acc: {:.4f}".format(
            epoch,
            total_loss / n_train_samples, color_loss / n_train_samples, gender_loss / n_train_samples, article_loss / n_train_samples,
            accuracy_color / n_train_samples, accuracy_gender / n_train_samples, accuracy_article / n_train_samples))

        # Loss and accuracy write to logs
        logger.add_scalar('train_total_loss', total_loss / n_train_samples, epoch)  
        logger.add_scalar('train_color_loss', color_loss / n_train_samples, epoch)  
        logger.add_scalar('train_gender_loss', gender_loss / n_train_samples, epoch)  
        logger.add_scalar('train_article_loss', article_loss / n_train_samples, epoch)  
        logger.add_scalar('train_color_acc', accuracy_color / n_train_samples, epoch)  
        logger.add_scalar('train_gender_acc', accuracy_gender / n_train_samples, epoch)  
        logger.add_scalar('train_article_acc', accuracy_article / n_train_samples, epoch) 

        if epoch % 2 == 0:
            validate(model=model, dataloader=val_dataloader, logger=logger, iteration=epoch, device=device, checkpoint=None)

        if epoch % 2 == 0:
            checkpoint_save(model, savedir, epoch)
        pbar.close() 

😺九、predict.py

该程序中定义了两个函数：

• validate用于在训练过程中启动验证。
• visualize_grid用于对测试集进行评估。

在visualize_grid中，添加了三种属性测试结果的混淆矩阵，以及可视化预测结果。
在main函数中，需要对测试集进行评估就注释掉Single image testing。反之，如果需要对单张图片测试，需要注释掉Dir testing。

from PIL import Image  
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from dataset import FashionDataset, AttributesDataset, mean, std
from model import MultiOutputModel
from utils import get_loss, calculate_metrics, checkpoint_load
from option import get_args


args = get_args()
batch_size = args.batch_size
num_workers = args.num_workers
device = args.device
original_csv_path = args.original_csv_path
val_csv_path = args.val_csv_path
checkpoint=args.checkpoint
predict_image_path = args.predict_image_path


def validate(model, dataloader, logger, iteration, device, checkpoint):
    if checkpoint is not None:
        checkpoint_load(model, checkpoint)

    model.eval()
    with torch.no_grad():

        # The total loss and accuracy of each category in initializing the validation set
        avg_loss, accuracy_color, accuracy_gender, accuracy_article = 0, 0, 0, 0

        for batch in dataloader:
            img = batch['img']
            target_labels = batch['labels']
            target_labels = {t: target_labels[t].to(device) for t in target_labels}
            output = model(img.to(device))

            val_train, val_train_losses = get_loss(output, target_labels)
            avg_loss += val_train.item()
            batch_accuracy_color, batch_accuracy_gender, batch_accuracy_article = calculate_metrics(output, target_labels)

            accuracy_color += batch_accuracy_color
            accuracy_gender += batch_accuracy_gender
            accuracy_article += batch_accuracy_article

    n_samples = len(dataloader)
    avg_loss /= n_samples
    accuracy_color /= n_samples
    accuracy_gender /= n_samples
    accuracy_article /= n_samples
    print('-' * 80)
    print("Validation ====> loss: {:.4f}, color_acc: {:.4f}, gender_acc: {:.4f}, article_acc: {:.4f}\n".format(
        avg_loss, accuracy_color, accuracy_gender, accuracy_article))

    logger.add_scalar('val_loss', avg_loss, iteration)
    logger.add_scalar('val_color_acc', accuracy_color, iteration)
    logger.add_scalar('val_color_acc', accuracy_gender, iteration)
    logger.add_scalar('val_color_acc', accuracy_article, iteration)

    model.train()


def visualize_grid(model, dataloader, attributes, device, show_cn_matrices=True, show_images=True, checkpoint=None,
                   show_gt=False):
    if checkpoint is not None:
        model, _ = checkpoint_load(model, checkpoint)
    model.eval()

    # Define image list
    imgs = []       

    # Define a list of predicted results (predicted labels, predicted color labels, predicted gender labels, predicted article labels)
    labels, predicted_color_all, predicted_gender_all, predicted_article_all = [], [], [], []

    # Define a list of real values (real labels, real color labels, real gender labels, real article labels)
    gt_labels, gt_color_all, gt_gender_all, gt_article_all = [], [], [], []

    # Initialize precision for each category
    accuracy_color = 0
    accuracy_gender = 0
    accuracy_article = 0

    with torch.no_grad():
        for batch in dataloader:
            img = batch['img']
            gt_colors = batch['labels']['color_labels']
            gt_genders = batch['labels']['gender_labels']
            gt_articles = batch['labels']['article_labels']
            output = model(img)

            batch_accuracy_color, batch_accuracy_gender, batch_accuracy_article = \
                calculate_metrics(output, batch['labels'])
            accuracy_color += batch_accuracy_color
            accuracy_gender += batch_accuracy_gender
            accuracy_article += batch_accuracy_article

            # Calculate maximum probability prediction label
            _, predicted_colors = output['color'].cpu().max(1)
            _, predicted_genders = output['gender'].cpu().max(1)
            _, predicted_articles = output['article'].cpu().max(1)

            for i in range(img.shape[0]):
                image = np.clip(img[i].permute(1, 2, 0).numpy() * std + mean, 0, 1)

                predicted_color = attributes.color_id_to_name[predicted_colors[i].item()]
                predicted_gender = attributes.gender_id_to_name[predicted_genders[i].item()]
                predicted_article = attributes.article_id_to_name[predicted_articles[i].item()]

                gt_color = attributes.color_id_to_name[gt_colors[i].item()]
                gt_gender = attributes.gender_id_to_name[gt_genders[i].item()]
                gt_article = attributes.article_id_to_name[gt_articles[i].item()]

                gt_color_all.append(gt_color)
                gt_gender_all.append(gt_gender)
                gt_article_all.append(gt_article)

                predicted_color_all.append(predicted_color)
                predicted_gender_all.append(predicted_gender)
                predicted_article_all.append(predicted_article)

                imgs.append(image)
                labels.append("{}\n{}\n{}".format(predicted_gender, predicted_article, predicted_color))
                gt_labels.append("{}\n{}\n{}".format(gt_gender, gt_article, gt_color))

    if not show_gt:
        n_samples = len(dataloader)
        print("Accuracy ====> color: {:.4f}, gender: {:.4f}, article: {:.4f}".format(
            accuracy_color / n_samples,
            accuracy_gender / n_samples,
            accuracy_article / n_samples))

    # Draw confusion matrix
    if show_cn_matrices:
        # Color confusion matrix
        cn_matrix = confusion_matrix(
            y_true=gt_color_all,
            y_pred=predicted_color_all,
            labels=attributes.color_labels,
            normalize='true')
        ConfusionMatrixDisplay(confusion_matrix=cn_matrix, display_labels=attributes.color_labels).plot(include_values=False, xticks_rotation='vertical')
        plt.title("Colors")
        plt.tight_layout()
        plt.savefig("confusion_matrix_color.png")
        # plt.show()

        # Gender confusion matrix
        cn_matrix = confusion_matrix(
            y_true=gt_gender_all,
            y_pred=predicted_gender_all,
            labels=attributes.gender_labels,
            normalize='true')
        ConfusionMatrixDisplay(confusion_matrix=cn_matrix, display_labels=attributes.gender_labels).plot(xticks_rotation='horizontal')
        plt.title("Genders")
        plt.tight_layout()
        plt.savefig("confusion_matrix_gender.png")
        # plt.show()

        # Article confusion matrix (with too many categories, images may be too large to display fully)
        cn_matrix = confusion_matrix(
            y_true=gt_article_all,
            y_pred=predicted_article_all,
            labels=attributes.article_labels,
            normalize='true')
        plt.rcParams.update({'font.size': 1.8})
        plt.rcParams.update({'figure.dpi': 300})
        ConfusionMatrixDisplay(confusion_matrix=cn_matrix, display_labels=attributes.article_labels).plot(
            include_values=False, xticks_rotation='vertical')
        plt.rcParams.update({'figure.dpi': 100})
        plt.rcParams.update({'font.size': 5})
        plt.title("Article types")
        plt.savefig("confusion_matrix_article.png")
        # plt.show()

    if show_images:
        labels = gt_labels if show_gt else labels
        title = "Ground truth labels" if show_gt else "Predicted labels"
        n_cols = 5
        n_rows = 3
        fig, axs = plt.subplots(n_rows, n_cols, figsize=(10, 10))
        axs = axs.flatten()
        for img, ax, label in zip(imgs, axs, labels):
            ax.set_xlabel(label, rotation=0)
            ax.get_xaxis().set_ticks([])
            ax.get_yaxis().set_ticks([])
            ax.imshow(img)
        plt.suptitle(title)
        plt.tight_layout()
        plt.savefig("images.png")
        # plt.show()

    model.train()


if __name__ == '__main__':

    """
    Dir testing
    """
    attributes = AttributesDataset(original_csv_path)
    val_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    test_dataset = FashionDataset(val_csv_path, attributes, val_transform)
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)

    model = MultiOutputModel(n_color_classes=attributes.num_colors, n_gender_classes=attributes.num_genders,
                             n_article_classes=attributes.num_articles).to('cpu')

    visualize_grid(model, test_dataloader, attributes, device, checkpoint)

    """
    Single image testing
    """
    model = torch.load(checkpoint, map_location='cpu')
    img = Image.open(predict_image_path)  
    if img.mode != 'RGB':  
        img = img.convert('RGB')  
    img_tensor = val_transform(img).unsqueeze(0)
    with torch.no_grad():
        outputs = model(img_tensor)
        _, predicted_color = outputs['color'].cpu().max(1)
        _, predicted_gender = outputs['gender'].cpu().max(1)
        _, predicted_article = outputs['article'].cpu().max(1)

        print("Predicted color ====> {}, gender: {}, article: {}".format(
            attributes.color_id_to_name[predicted_color.item()],
            attributes.gender_id_to_name[predicted_gender.item()],
            attributes.article_id_to_name[predicted_article.item()]))