How To Use Matplotlib For Plotting Samples From An Object Detection Dataset | by Varun Dutt | Towards Data Science

Visualizing samples from the dataset is essential for efficiently exploring machine learning tasks. It provides key insights to help you narrow down your choices out of the several options available in the model architecture, data augmentation, etc. Plotting samples from an object detection dataset is trickier compared to classification or regression tasks.
Here is a short tutorial on how to use the Matplotlib library to create high-quality plots for object detection samples.
We will be using the Yolov5 PyTorch version of the Chess Dataset which you can download from https://public.roboflow.com/object-detection/chess-full/23. We will also use data frames from the previous post on “A general approach for creating a dataset pipeline in PyTorch using an object detection dataset”.
Please give it a read if you haven’t already done so. Here are the links for the data frames.
https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/train_fd.csvhttps://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/valid_fd.csvhttps://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/test_fd.csv
https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/train_fd.csv
https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/valid_fd.csv
https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/test_fd.csv
You can download the full notebook for the tutorial from here.
Now, let’s get coding!!
First of all, import all the libraries and files needed for the tutorial
!pip install fastai --upgradefrom fastai import *from fastai.vision.all import *from fastai.imports import *import cv2, osfrom matplotlib import patches, text, patheffectsroot = Path("/content/data")!curl -L "https://public.roboflow.com/ds/pK9BobmV9A?key=G9IUWXkCZA" > roboflow.zip; unzip roboflow.zip; rm roboflow.ziptrain_df = pd.read_csv("https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/train_fd.csv")valid_df = pd.read_csv("https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/valid_fd.csv")test_df = pd.read_csv("https://raw.githubusercontent.com/Varun9213/Blog_machine_learing/main/Data_blog_1/test_fd.csv")train_df.head()
Let’s start by looking at how we want the final plot to turn out.
Alright, now that we have a clear goal in mind let’s look at the implementation.
If you have read my previous post, I like to break down a problem into smaller parts so, let's start by plotting an image from the dataset and drawing a bounding box for only one of the objects in the image.
train_images = get_image_files(root/"train")fig, ax = plt.subplots()path_img = train_images[0]img = cv2.imread(str(path_img))img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)ax.xaxis.tick_top()ax.imshow(img)
Note get_image_files() is a fastai function that returns a list of Paths for all the images in the subdirectories of the parent, you can visit https://docs.fast.ai/ to further explore the library.
def get_target_ds(name, df):  rows = df[df["name"] == name[:-4]]  return rows["class"].values, rows[bboxes_cols].valuesbboxes_cols = ["center_x", "center_y", "width", "height"]labels, boxes = get_target_ds(path_img.name, train_df)x = (boxes[0][0] - boxes[0][2]*0.5)*img.shape[1]y = (boxes[0][1] - boxes[0][3]*0.5)*img.shape[0]w = boxes[0][2] * img.shape[1]h = boxes[0][3] * img.shape[0]fig, ax = plt.subplots(figsize = (6,9))ax.xaxis.tick_top()ax.imshow(img)ax.add_patch(patches.Rectangle((x,y),w,h, fill=False, edgecolor='red', lw=2))
Okay, now let’s break this code block down to see how exactly that box appeared in the image.
ax.add_patch() is a Matplotlib method to draw a figure or a patch onto a plot which we use here to draw a rectangle given by the bounding box coordinates.
patches.Rectangle() is the argument we pass to the add_patch() method which then draws the rectangle onto the image. Notice, that the format for the bounding box that needs to be passed to the function is ((x_min, y_min), width, height) however, the format that the dataset provides for the bounding box is different, and had to be converted to the desired format.
As we chose the Yolo version of the dataset the format for the bounding boxes was [x_center, y_center, width, height]. Notice here that the coordinates had been normalized which had to be scaled back according to the resolution of the image.
Note that the image from the albumentations docs which you can find here shows you how the coordinate system works in most plotting and augmentation libraries. It is flipped upside down compared to the coordinate system we generally use. It also shows you how the coordinates of a bounding box are represented in various formats.
Finally, let’s add the corresponding label to the box.
fig, ax = plt.subplots(figsize = (6,9))ax.xaxis.tick_top()ax.imshow(img)ax.add_patch(patches.Rectangle((x,y),w,h, fill=False, edgecolor='red', lw=2))ax.text(x,(y-50),str(labels[0]),verticalalignment='top',color='white',fontsize=10,weight='bold')
There seems to be something still missing in this plot, our target plot was a lot cleaner than the one we have just created.
The difference is the missing outlines that our target plot has, let us add the outline to get a cleaner plot.
fig, ax = plt.subplots(figsize = (6,9))ax.xaxis.tick_top()ax.imshow(img)ax.add_patch(patches.Rectangle((x,y),w,h, fill=False, edgecolor='red', lw=2)).set_path_effects([patheffects.Stroke(linewidth=4, foreground='black'), patheffects.Normal()])ax.text(x,(y-50),str(labels[0]),verticalalignment='top',color='white',fontsize=10,weight='bold').set_path_effects([patheffects.Stroke(linewidth=4, foreground='black'), patheffects.Normal()])
Yes! This looks much better. Finally, to wrap it up let's create functions for every step and plot all the boxes with corresponding labels in the image. Also, let’s create a function to generate plots for multiple samples.
def get_bb(bboxes, img):  boxes = bboxes.copy()  boxes[:,0] = (boxes[:,0] - boxes[:,2]*0.5)*img.shape[1]  boxes[:,1] = (boxes[:,1] - boxes[:,3]*0.5)*img.shape[0]  boxes[:,2] = boxes[:,2] * img.shape[1]  boxes[:,3] = boxes[:,3] * img.shape[0]  if boxes.shape[0] == 1 : return boxes  return np.squeeze(boxes)def img_show(img, ax = None, figsize=(7,11)):  if ax is None: fig, ax = plt.subplots(figsize=figsize)  ax.xaxis.tick_top()  ax.imshow(img)  return axdef draw_outline(obj):  obj.set_path_effects([patheffects.Stroke(linewidth=4,  foreground='black'), patheffects.Normal()])def draw_box(img, ax, bb):  patch = ax.add_patch(patches.Rectangle((bb[0],bb[1]), bb[2], bb[3], fill=False, edgecolor='red', lw=2))  draw_outline(patch)def draw_text(ax, bb, txt, disp):  text = ax.text(bb[0],(bb[1]-disp),txt,verticalalignment='top'  ,color='white',fontsize=10,weight='bold')  draw_outline(text)def plot_sample(img, bboxes, labels, ax=None, figsize=(7,11)):  bb = get_bb(bboxes, img)  ax = img_show(img, ax=ax)  for i in range(len(bboxes)):    draw_box(img,ax,bb[i])    draw_text(ax, bb[i], str(labels[i]), img.shape[0]*0.05)def multiplot(dim:tuple, df, images, idxs = None, figsize=(18,10)):  if idxs is None: idxs = np.random.randint(0, len(images)-1,          dim[0]*dim[1])  fig, ax = plt.subplots(dim[0],dim[1], figsize=figsize)  plt.subplots_adjust(wspace=0.1, hspace=0)  fig.tight_layout()  for i in range(dim[0]):    for j in range(dim[1]):      img = images[idxs[(i+1)*j]]      labels, bboxes = get_target_ds(img.name, df)      img = cv2.imread(str(img), cv2.IMREAD_UNCHANGED)      img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)      plot_sample(img, bboxes, labels, ax=ax[i][j])plot_sample(img, boxes, labels)
multiplot((2,4), train_df, train_images, figsize=(20,8))
After finishing a machine learning project you realize that the majority of your coding time goes into pre-processing, EDA, plotting samples, augmentation, and building pipelines. Getting comfortable with plotting libraries saves you a lot of time in the long run.
I hope this post helped you in using these plotting functions more effectively. I will be covering data augmentation in my next post which in my opinion gives you the best bang for your buck in terms of performance if used properly.
So stay tuned and thank you!!