A Tutorial On Creating Data Pipeline For Object Detection Using PyTorch And Fastai | by Varun Dutt | Towards Data Science

Most of the time you start out in machine learning using pre-processed datasets like Sklearn’s datasets to implement machine learning algorithms but as you upshift to more novel machine learning projects, the datasets are usually in raw formats and may have a very loose structure.
I am sharing a general approach here that I have developed for starting with these raw datasets and building a dataset pipeline in PyTorch.
The Dataset we are going to use is the Chess Dataset which is an object detection dataset, you can download the dataset using the link https://public.roboflow.com/object-detection/chess-full/23, for purposes of this tutorial we will use the Yolo v5 PyTorch version of the dataset.
Download the notebook for the tutorial using https://github.com/Varun9213/Blog_machine_learing/blob/main/Dataset_pipeline_blog_1.ipynb
You can download the dataset with the link generated by the website directly from your notebook. The rest of the article will be structured in subsections which will build a stepwise approach for creating PyTorch dataset pipelines.
Raw datasets often look like complicated puzzles. I too got overwhelmed the first time I opened a very raw dataset but then I looked it over for some time then broke the problem down into smaller pieces and it started making sense. It is always a good idea to break the puzzle into smaller pieces.
Generally, you will come across two forms of structuring schemes.
Commoner structure has the images divided into separate folders based on their labels. The name of the parent folder is the label those images belong to. The other common structure is where all the images are put in one folder and a separate file (.csv, .txt, .json, etc.) holds the labels with index or file name being the name of the image and the corresponding label.
The Chess Dataset contains three folders train, valid and test with each folder having an images folder and labels folder. There are two readme files and one YAML file named data.yaml. Our dataset is of the second type where the classes and bounding boxes are in a separate .txt file which is the structure you will generally find in object detection datasets.
Now that we have explored the structure and looked at the extensions of the files we are dealing with let's open these files to have a better understanding of what steps we need to take next.
When starting out you will frequently come across file types you have never dealt with and that might seem a complicated thing to deal with, however, most of the time it's just a five minutes google search and is not as complicated as it might appear at first.
This dataset has 3 types of files .jpg, .txt, .yaml. Most of the readers will probably be familiar with the first two file types though YAML might be new to some readers but YAML is just a serialization standard similar to JSON and is used mostly to store configurations. In most cases you need not know much about the protocols the file type follows, you only need to know how to open these files in python and extract data from them.
We will start by exploring the YAML file. You can read about how to deal with YAML files by following the link https://stackabuse.com/reading-and-writing-yaml-to-a-file-in-python/.It covers almost all the tools you might need to work with YAML files in python.
Let's start by installing and importing all the libraries we will need for this tutorial.
!pip install fastai --upgrade!pip install pyyamlfrom fastai import *from fastai.vision.all import *from fastai.imports import *import yaml, cv2, osfrom torchvision import transforms as Troot = Path("/content/data")
A Path is an object in fastai that makes your life a lot easier, you can read about Path and other very handy tools that fastai library provides at docs.fast.ai, I would also highly recommend Jeremy’s free fastai course for deep learning on the fastai website https://course19.fast.ai/part2.
Now let's have a peek at the data.yaml file.
with open(root/"data.yaml") as file:data = yaml.load(file)data
It's a dictionary with label names, number of classes and paths of training and validation folder. For the purpose of building a pipeline, we can ignore this file.
Next, let's look at the images
path_train = Path("/content/data/train")path_val = Path("/content/data/valid")path_test = Path("/content/data/test")train_images = get_image_files(path_train/"images")img = cv2.imread(str(train_images[-1]), cv2.IMREAD_UNCHANGED)img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)plt.imshow(img)
get_image_files() is another very handy fastai method that gets all the images with common formats like jpg, png, etc. from all the subdirectories. The returned list does not contain the images themselves but the Paths to all the images.
Note that the resolution of the image is 1371 x 2048 which is quite high, most object detection algorithms like yolo allow rectangular images however the training will be naturally slower when using high resolution images. Hence we might want to resize these images to a lower resolution. However contrary to classification tasks, in object detection the target is not invariant to the resolution as the bounding boxes have to also be resized if the image is resized.
Let's finally open the corresponding label.
f = train_images[-1].name[:-4] + ".txt"with open(root/"train"/"labels"/f, 'r') as file:  targets = file.read()targets = targets.split("\n")targets
The first entry of the samples clearly is the class, the next 4 entries define the bounding box. The format of the bounding box is different for different algorithms, Yolo’s format for bounding box is [ center_x, center_y, width, height] https://towardsdatascience.com/image-data-labelling-and-annotation-everything-you-need-to-know-86ede6c684b1 is an excellent resource if you want to explore the topic in more depth.
This step is very subjective and everyone has their own preference, I like to compile my data into a Data Frame before creating a PyTorch dataset object.
Note this step is not necessary and can be skipped. However, in my opinion, with a small amount of effort and few lines of code you can significantly simplify the dataset building step. Having said that, the most important thing in any approach is familiarity which comes from practice, so I would highly recommend you to figure out a general approach you are comfortable with and try to stick to it as much as you can.
Let's start by choosing just one sample and making a data frame for that one sample. The data frame will have the name of the image, the class and the bounding box coordinates.
col = ["name", "class", "center_x", "center_y", "width", "height"]train_df = pd.DataFrame(columns= col)train_dfrow = {}row[col[0]] = f[:-4]for i,t in enumerate(targets[0].split(" ")):  row[col[i+1]] = float(t)rowtrain_df.append(row, ignore_index=True)
Note it is generally a good idea to take a single sample from the entire dataset and work with that sample. This has two advantages, one its a lot quicker especially if the dataset is very large, second its a lot easier to handle a single sample than the entire dataset. This practice has brought the greatest improvement to my ability to prepare dataset pipelines efficiently.
Finally, we will convert these individual commands into functions which will convert the entire dataset to a Data Frame
def get_targets(image_path):  f = image_path.name[:-4] + ".txt"  with open((image_path.parent).parent/"labels"/f, 'r') as file:    targets = file.read()  return targets.split("\n")def get_row(image_path, target, col):   row = {}  row[col[0]] = image_path.name[:-4]  for i,t in enumerate(target.split(" ")):    try:      row[col[i+1]] = float(t)    except:      print("Could not convert {} to float".format(t))  return rowdef add_targets_to_df(df, image_path):  targets = get_targets(image_path)  df_new = df  for target in targets:    row = get_row(image_path, target, df.columns)    df_new = df_new.append(row, ignore_index = True)  return df_newdef create_df(cols, images_paths):  df = pd.DataFrame(columns=cols)  for path in images_paths:    df = add_targets_to_df(df, path)  return df
It is always a good practice to break down a complex task into smaller parts and wrap each part into a function, it simplifies the problem and produces a more readable code which makes troubleshooting a lot easier compared to jamming everything in a single function.
final_df = create_df(col, train_images)final_df.head()
Now the final step is to create the PyTorch dataset object which will be the final section.
Pytorch and Tensorflow are two of the most popular libraries for deep learning, PyTorch recently has become more popular among researchers because of the flexibility the library provides.
Let's begin by constructing a Dataset pipeline for the chess dataset.
def get_target_ds(name, df):  rows = df[df["name"] == name[:-4]]  return rows["class"].values, rows[bboxes_cols].valuesclass ChessDataset(torch.utils.data.Dataset):    def __init__(self, images_path, df):    super(ChessDataset, self).__init__()    self.images_path = images_path    self.df = df  def __len__(self):    return len(self.images_path)  def __getitem__(self,idx):    img_path = self.images_path[idx]    img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)    target = {}    labels, boxes = get_target_ds(img_path.name, self.df)    areas = boxes[:,2] * boxes[:,3]    iscrowd = torch.zeros((boxes.shape[0],))    image_id = torch.tensor([idx])    labels = torch.as_tensor(labels, dtype=torch.float32)    boxes = torch.as_tensor(boxes, dtype=torch.float32)    areas = torch.as_tensor(areas, dtype=torch.float32)    target["boxes"] = boxes    target["labels"] = labels    target["areas"] = areas    target["iscrowd"] = iscrowd    target["image_id"] = image_id    img = F.to_tensor(img)    return img, target
The ChessDataset class inherits torch.utils.data.Dataset class from PyTorch for which we need to define __init__(), __len__() and __getitem__() functions, to get a more in depth understanding of how to create PyTorch custom datasets check out https://towardsdatascience.com/building-efficient-custom-datasets-in-pytorch-2563b946fd9f blog post.
Finally, let's test our class by creating an instance and getting a sample by calling __getitem__() to see if it is working properly.
dataset = ChessDataset(train_images, final_df)img , target = dataset.__getitem__(1)size = img.shapeprint(img.shape, target["boxes"], target["labels"])plt.imshow(img.permute(1,2,0))
Your dataset is ready to be used with PyTorch's dataloaders !!
The idea here is to have a sequence of steps that are broad enough for most datasets. Preparing datasets for training is also a skill and as with any skill practice is the key to becoming competent at it, I would highly recommend you to go on Kaggle and open closed competitions with complicated-looking dataset and try to build just a dataset pipeline for those competitions and then use the kernels people have published for those competitions to troubleshoot, the more the dataset makes you uncomfortable the more you will learn from the exercise.
Finally, I hope this post provided some insight to all the readers, I plan to post a continuation of this post working on the same dataset, covering various plotting functions and strategy, picking and training a pre-trained model in PyTorch and anything else that might be relevant to the topic, so stay tuned and Thank You.