用pytorch的多標註的圖像分類

什麼是多標註圖像分類

在圖像分類領域,你可能需要確定一個物體的若干個屬性。如顏色,大小等類別。

數據集:使用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等。


用pytorch的多標註的圖像分類


此次的目標是創建並訓練一個神經網絡,來預測這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>

這樣就可以訓練了。

此次使用的數據集是低分辨率的,為了得到更好的模型,可以採用更好的數據。



分享到:


相關文章: