什麼是多標註圖像分類
在圖像分類領域,你可能需要確定一個物體的若干個屬性。如顏色,大小等類別。
數據集:使用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的類,來解析數據。如下:
<code>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']])/<code>
模型:可直接使用torchvision.models中的mobilenet_v2網絡,但需要修改以適應3個標籤的預測。如下:
<code>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>, <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)
)/<height>/<width>/<height>/<width>/<channels>/<batch>/<code>
訓練:參數可以調整。
<code>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())/<code>
這樣就可以訓練了。
此次使用的數據集是低分辨率的,為了得到更好的模型,可以採用更好的數據。
閱讀更多 AI踐行者 的文章