简体   繁体   中英

Multi-label, multi-class image classifier (ConvNet) with PyTorch

I am trying to implement an image classifier (CNN/ConvNet) with PyTorch where I want to read my labels from a csv-file. I have 4 different classes and an image may belong to more than one class.

I have read through the PyTorch Tutorial and this Stanford tutorial and this one , but none of them cover my specific case. I have managed to build a custom function of the torch.utils.data.Dataset class which works fine for reading the labels from a csv-file for a binary classifier only though.

This is the code for the torch.utils.data.Dataset class I have so far (slightly modified from the third tutorial linked above):

import torch
import torchvision.transforms as transforms
import torch.utils.data as data
from PIL import Image
import numpy as np
import pandas as pd


class MyCustomDataset(data.Dataset):
# __init__ function is where the initial logic happens like reading a csv,
# assigning transforms etc.
def __init__(self, csv_path):
    # Transforms
    self.random_crop = transforms.RandomCrop(800)
    self.to_tensor = transforms.ToTensor()
    # Read the csv file
    self.data_info = pd.read_csv(csv_path, header=None)
    # First column contains the image paths
    self.image_arr = np.asarray(self.data_info.iloc[:, 0])
    # Second column is the labels
    self.label_arr = np.asarray(self.data_info.iloc[:, 1])
    # Calculate len
    self.data_len = len(self.data_info.index)


# __getitem__ function returns the data and labels. This function is
# called from dataloader like this
def __getitem__(self, index):
    # Get image name from the pandas df
    single_image_name = self.image_arr[index]
    # Open image
    img_as_img = Image.open(single_image_name)
    img_cropped = self.random_crop(img_as_img)
    img_as_tensor = self.to_tensor(img_cropped)

    # Get label(class) of the image based on the cropped pandas column
    single_image_label = self.label_arr[index]

    return (img_as_tensor, single_image_label)

def __len__(self):
    return self.data_len

Specifically, I am trying to read my labels from a file with the following structure:

CSV数据

And my specific problem is, that I can't figure out how to implement this into my Dataset class. I think I am missing the link between the (manual) assignment of the labels in the csv and how they are read by PyTorch, as I am rather new to the framework.
I'd appreciate any help on how to get this to work, or if there are actually examples covering this, a link would be highly appreciated as well!

Maybe I am missing something, but if you want to convert your columns 1..N ( N = 4 here) into a label vector or shape (N,) (eg given your example data, label(img1) = [0, 0, 0, 1] , label(img3) = [1, 0, 1, 0] , ...), why not:

  1. Read all the label columns into self.label_arr :

     self.label_arr = np.asarray(self.data_info.iloc[:, 1:]) # columns 1 to N 
  2. Return accordingly the labels in __getitem__() (no change here):

     single_image_label = self.label_arr[index] 

To train your classifier, you could then compute eg the cross-entropy between your (N,) predictions and the target labels.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM