什么是多标注图像分类
在图像分类领域,你可能需要确定一个物体的若干个属性。如颜色,大小等类别。
数据集:使用Kaggle的Fashion Product Images数据集(低分辨率)。
它包括了4.4万个图像,每个图像有9个标注。它的目录形式如下;
. ├── fashion-product-images │ ├── images │ └── styles.csv ├── dataset.py ├── model.py ├── requirements.txt ├── split_data.py ├── test.py └── train.py
其中文件style.csv是数据标注数据。此次只是用其中的gender,articletype,basecolour标签。
gender标签有5个值: Boys, Girls, Men, Unisex, Women。 color有47种颜色。 article有143种,如Sports Sandals, Wallets ,Sweaters等。
此次的目标是创建并训练一个神经网络,来预测这3种标签。
依赖的库:
matplotlib numpy pillow scikit-learn torch torchvision tqdm
用split_data.py 来拆分数据为训练集,测试集,对应train.csv和val.csv。
数据加载:创建一个基于PyTorch Dataset的类,来解析数据。如下:
class FashionDataset(Dataset):
def __init__(...):
...
?
# initialize the arrays to store the ground truth labels and paths to the images
?
? self.data = []
? self.color_labels = []
? self.gender_labels = []
? self.article_labels = []
?
# read the annotations from the CSV file
?
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']])
模型:可直接使用torchvision.models中的mobilenet_v2网络,但需要修改以适应3个标签的预测。如下:
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 # take the model without classifier
last_channel = models.mobilenet_v2().last_channel # size of the layer before the classifier
?
# the input for the classifier should be two-dimensional, but we will have
# [<batch_size>, <channels>, <width>, <height>]
# so, let's do the spatial averaging: reduce <width> and <height> to 1
self.pool = nn.AdaptiveAvgPool2D((1, 1))
# create separate classifiers for our outputs
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)
)
训练:参数可以调整。
N_epochs = 50
batch_size = 16
?
...
?
model = MultiOutputModel(n_color_classes=attributes.num_colors, n_gender_classes=attributes.num_genders,
n_article_classes=attributes.num_articles).to(device)
?
optimizer = torch.optim.Adam(model.parameters())
这样就可以训练了。
此次使用的数据集是低分辨率的,为了得到更好的模型,可以采用更好的数据。