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:
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:
Read all the label columns into self.label_arr
:
self.label_arr = np.asarray(self.data_info.iloc[:, 1:]) # columns 1 to N
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.