v2.2.1
EXAMPLE CODE
Creating Complex Datasets
How to create datasets with multiple annotation types.
Datasets often have multiple labels such as classifications, bounding boxes, segmentations, and others. In order to create an intuitive layout of tensors, it's advisable to create a dataset hierarchy that captures the relationship between the different label types. This can be done with hub tensor
groups.This example show to to use groups to create a dataset containing image classifications of "indoor" and "outdoor", as well as bounding boxes of objects such as "dog" and "cat".
Create the Hub Dataset
The first step is to download the small dataset below called animals complex.
animals_complex.zip
278KB
Binary
animals complex dataset
The images and their classes are stored in a
classification folder where the subfolders correspond to the class names. Bounding boxes for object detection are stored in a separate boxes subfolder, which also contains a list of class names for object detection in the file box_names.txt. In YOLO format, images and annotations are typically matched using a common filename such as image -> filename.jpeg and annotation -> filename.txt . The data structure for the dataset is shown below:1
data_dir
2
|_classification
3
|_indoor
4
|_image1.png
5
|_image2.png
6
|_outdoor
7
|_image3.png
8
|_image4.png
9
|_boxes
10
|_image1.txt
11
|_image3.txt
12
|_image3.txt
13
|_image4.txt
14
|_classes.txt
Copied!
Now that you have the data, let's create a Hub Dataset in the
./animals_complex_hub folder by running: 1
import hub
2
from PIL import Image, ImageDraw
3
import numpy as np
4
import os
5
6
ds = hub.empty('./animals_complex_hub') # Create the dataset locally
Copied!
Next, let's specify the folder paths containing the classification and object detection data. It's also helpful to create a list of all of the image files and class names for classification and object detection tasks.
1
classification_folder = './animals_complex/classification'
2
boxes_folder = './animals_complex/boxes'
3
4
# List of all class names for classification
5
class_names = os.listdir(classification_folder)
6
7
fn_imgs = []
8
for dirpath, dirnames, filenames in os.walk(classification_folder):
9
for filename in filenames:
10
fn_imgs.append(os.path.join(dirpath, filename))
11
12
# List of all class names for object detection
13
with open(os.path.join(boxes_folder, 'classes.txt'), 'r') as f:
14
class_names_boxes = f.read().splitlines()
Copied!
Since annotations in YOLO are typically stored in text files, it's useful to write a helper function that parses the annotation file and returns numpy arrays with the bounding box coordinates and bounding box classes.
1
def read_yolo_boxes(fn:str):
2
"""
3
Function reads a label.txt YOLO file and returns a numpy array of yolo_boxes
4
for the box geometry and yolo_labels for the corresponding box labels.
5
"""
6
7
box_f = open(fn)
8
lines = box_f.read()
9
box_f.close()
10
11
# Split each box into a separate lines
12
lines_split = lines.splitlines()
13
14
yolo_boxes = np.zeros((len(lines_split),4))
15
yolo_labels = np.zeros(len(lines_split))
16
17
# Go through each line and parse data
18
for l, line in enumerate(lines_split):
19
line_split = line.split()
20
yolo_boxes[l,:]=np.array((float(line_split[1]), float(line_split[2]), float(line_split[3]), float(line_split[4])))
21
yolo_labels[l]=int(line_split[0])
22
23
return yolo_boxes, yolo_labels
Copied!
Next, let's create the groups and tensors for this data. In order to separate the two annotations, a
boxes group is created to wrap around the label and bbox tensors which contains the coordinates and labels for the bounding boxes.1
with ds:
2
# Image
3
ds.create_tensor('images', htype='image', sample_compression='jpeg')
4
5
# Classification
6
ds.create_tensor('labels', htype='class_label', class_names = class_names)
7
8
# Object Detection
9
ds.create_group('boxes')
10
ds.boxes.create_tensor('bbox', htype='bbox')
11
ds.boxes.create_tensor('label', htype='class_label', class_names = class_names_boxes)
12
# An alternate approach is to use '/' notation, which automatically creates the boxes group
13
# ds.create_tensor('boxes/bbox', ...)
14
# ds.create_tensor('boxes/label', ...)
Copied!
In order for Activeloop Platform to correctly visualize the labels,
class_names must be a list of strings, where the numerical labels correspond to the index of the label in the list.Finally, let's iterate through all the images in the dataset in order to populate the data in Hub. The first axis of the
boxes.bbox sample array corresponds to the first-and-only axis of the boxes.label sample array (i.e. if there are 3 boxes in an image, the labels array is 3x1 and the boxes array is 3x4).1
with ds:
2
#Iterate throgh the images
3
for fn_img in fn_imgs:
4
5
img_name = os.path.splitext(os.path.basename(fn_img))[0]
6
fn_box = img_name+'.txt'
7
8
# Get the class number for the classification
9
label_text = os.path.basename(os.path.dirname(fn_img))
10
label_num = class_names.index(label_text)
11
12
# Get the arrays for the bounding boxes and their classes
13
yolo_boxes, yolo_labels = read_yolo_boxes(os.path.join(boxes_folder,fn_box))
14
15
# Append classification data to tensors
16
ds.images.append(hub.read(os.path.join(fn_img)))
17
ds.labels.append(np.uint32(label_num))
18
19
# Append object detection data to tensors
20
ds.boxes.label.append(yolo_labels.astype(np.uint32))
21
ds.boxes.bbox.append(yolo_boxes.astype(np.float32))
Copied!
Inspect the Hub Dataset
Let's check out the second sample from this dataset and visualize the labels.
1
# Draw bounding boxes and the classfication label for the second image
2
3
ind = 1
4
img = Image.fromarray(ds.images[ind].numpy())
5
draw = ImageDraw.Draw(img)
6
(w,h) = img.size
7
boxes = ds.boxes.bbox[ind].numpy()
8
9
for b in range(boxes.shape[0]):
10
(xc,yc) = (int(boxes[b][0]*w), int(boxes[b][1]*h))
11
(x1,y1) = (int(xc-boxes[b][2]*w/2), int(yc-boxes[b][3]*h/2))
12
(x2,y2) = (int(xc+boxes[b][2]*w/2), int(yc+boxes[b][3]*h/2))
13
draw.rectangle([x1,y1,x2,y2], width=2)
14
draw.text((x1,y1), ds.boxes.label.info.class_names[ds.boxes.label[ind].numpy()[b]])
15
draw.text((0,0), ds.labels.info.class_names[ds.labels[ind].numpy()[0]])
Copied!
1
# Display the image and its bounding boxes
2
img
Copied!
Congrats! You just created a dataset with multiple types of annotations! 🎉
Last modified 1d ago