Training RetinaNet on Google Colab to detect pliers, hammers, and screwdrivers from KTH Handtools Dataset. | by Iustina Ivanova | Towards Data Science

In this article, RetinaNet is trained in Google Colab to detect plier, hammer and screwdriver instruments. The dataset was taken from an opened source called KTH Handtools Dataset. It consists of 3 types of images for the handtools: hammer, plier and screwdriver in different illuminations and different locations.
Jupyter notebook code for the article can be found in my github.
Firstly, the dataset should be downloaded from the source locally (unzip file and rename it to KTH_Handtool_Dataset), then download folder into google drive. After that we should mount google disc to google colab:
from google.colab import drivedrive.mount('/content/drive')
Secondly, the data need to be prepared to train RetinaNet: all the images in the initial dataset should be saved in different folders with correlated bounding boxes for each image. To train RetinaNet we need to create a CSV file: each line in this file should consist of the name of the image file from the dataset, bounding box coordinates for each object and the name of a class for the object. The following script parses all the folders (there are 3 folders for blue, white and brown background for images).
The first folder to parse is ‘Blue_background’. The difference between this folder and ‘Brown_background’ or ‘White_background’ is that the folders with instruments do not contain ‘Kinect’ and ‘webcam’ subfolders.
from bs4 import BeautifulSoupimport osimport csvfolder = '/content/drive/My Drive/KTH_Handtool_Dataset'subfolder = ['Blue_background']in_subfolder = ['Artificial', 'Cloudy', 'Directed']instruments = ['hammer1', 'hammer2', 'hammer3', 'plier1', 'plier2', 'plier3', 'screw1', 'screw2', 'screw3']with open('train.csv', mode='w') as file:  writer = csv.writer(file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)    for folder_name in subfolder:    for in_folder_name in in_subfolder:      for instrument in instruments:        directory_name = os.path.join(folder, folder_name, 'rgb', in_folder_name, instrument)        directory_name_xml = os.path.join(folder, folder_name, 'bboxes', in_folder_name, instrument)        for filename in os.listdir(directory_name_xml):          label = instrument          filename_jpg = filename[:-4] + '.jpg'          filename_str = os.path.join(directory_name, filename_jpg)          handler = open(os.path.join(directory_name_xml, filename)).read()          soup = BeautifulSoup(handler, "xml")          xmin = int(soup.xmin.string)          xmax = int(soup.xmax.string)          ymin = int(soup.ymin.string)          ymax = int(soup.ymax.string)          row = [filename_str, xmin, ymin, xmax, ymax, label]          writer.writerow(row)
parameter folder — the directory where all the data is stored (parent directory for extracted KTH_Handtool_Dataset folder).
For ‘White_background’ and ‘Brown_background’ script should be different as they contain subfolders ‘Kinect’ and ‘webcam’ in the folders for instruments (‘hammer1’, ‘hammer2’, ‘hammer3’, ‘plier1’, ‘plier2’, ‘plier3’, ‘screw1’, ‘screw2’, ‘screw3’). I added this line to include these folders
for sub_instrument in sub_instruments:
Another trick is that ‘rgb’ folder and ‘bbox’ folder have different names for Kinect: in some folders it is ‘Kinect’, in others it is ‘kinect’, therefore I decided to make different ways to parse the directory for images (‘rgb’ folder) and for bounding boxes (‘bbox’ folder). The dictionary is used to change the names:
dict_instr = {'Kinect': 'kinect'}
Next step: if the name ‘Kinect’ is not in the xml path folder, then it changes to lower case ‘kinect’:
dir_name_xml = os.path.join(folder, folder_name, 'bbox', in_folder_name, instrument)          if sub_instrument not in os.listdir(dir_name_xml):            sub_instrument = dict_instr[sub_instrument]
In the result, the script for the “Brown_background” and “White_background” is as following:
subfolder = ['Brown_background', 'White_background']sub_instruments = ['Kinect', 'webcam']dict_instr = {'Kinect': 'kinect'}# open file to write to the end of a filewith open('train.csv', mode='a') as file:  writer = csv.writer(file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)    for folder_name in subfolder:    for in_folder_name in in_subfolder:      for instrument in instruments:        for sub_instrument in sub_instruments:          directory_name = os.path.join(folder, folder_name, 'rgb', in_folder_name, instrument, sub_instrument)          dir_name_xml = os.path.join(folder, folder_name, 'bbox', in_folder_name, instrument)          if sub_instrument not in os.listdir(dir_name_xml):            sub_instrument = dict_instr[sub_instrument]          directory_name_xml = os.path.join(dir_name_xml, sub_instrument)          for filename in os.listdir(directory_name_xml):            label = instrument            filename_jpg = filename[:-4] + '.jpg'            filename_str = os.path.join(directory_name, filename_jpg)            handler = open(os.path.join(directory_name_xml, filename)).read()            soup = BeautifulSoup(handler, "xml")            xmin = int(soup.xmin.string)            xmax = int(soup.xmax.string)            ymin = int(soup.ymin.string)            ymax = int(soup.ymax.string)            row = [filename_str, xmin, ymin, xmax, ymax, label]            writer.writerow(row)
After running the code above we should do preprocessing, as some of the xml files do not have correlated image:
import pandas as pdimport os.pathlist_indexes_to_drop = []data = pd.read_csv("train.csv", header=None)with open('train_new.csv', mode='a') as file:  for i in range(len(data)):    fname = data.iloc[i, 0]    if not os.path.isfile(fname):      list_indexes_to_drop.append(i)data = data.drop(data.index[list_indexes_to_drop])data.to_csv(path_or_buf='train.csv', index=False, header=None)
After preprocessing the data should be split. All the elements in CSV file were randomly shuffled
data = pd.read_csv("train.csv", header=None)data = data.sample(frac=1).reset_index(drop=True)
then the data was split to a training set (80%) and a testing set (20%) by the next code:
amount_80 = int(0.8*len(data))train_data = data[:amount_80]test_data = data[amount_80:]print(len(train_data))print(len(test_data))
We should save train_data as train_annotations.csv:
train_data.to_csv(path_or_buf='train_annotations', index=False)
Next step is to download weights for training neural network from the website and put it into a folder ‘weights’:
!mkdir weights!wget -O /content/weights/resnet50_coco_best_v2.h5 https://github.com/fizyr/keras-retinanet/releases/download/0.5.1/resnet50_coco_best_v2.1.0.h5
We should also create folders ‘snapshots’ where RetinaNet will save weights during training, and ‘tensorboard’, where information about training will be solved:
!mkdir /content/drive/My\ Drive/kth_article/snapshots!mkdir /content/drive/My\ Drive/kth_article/tensorboard
Next step is to create csv file ‘classes.csv’ (which contains names of the instruments):
dict_classes = {    'hammer1': 0,    'hammer2': 1,    'hammer3': 2,    'plier1': 3,    'plier2': 4,    'plier3': 5,    'screw1': 6,    'screw2': 7,    'screw3': 8}with open('classes.csv', mode='w') as file:  writer = csv.writer(file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)  for key, val in dict_classes.items():      row = [key, val]    print(row)    writer.writerow(row)
We also need to install RetinaNet:
!cd ~!git clone https://github.com/fizyr/keras-retinanet%cd keras-retinanet!git checkout 42068ef9e406602d92a1afe2ee7d470f7e9860df!python setup.py install!python setup.py build_ext --inplace
We need to return to the parent directory:
%cd ..
Next script is for training neural network:
!retinanet-train --weights /weights/resnet50_coco_best_v2.h5 \--batch-size 4 --steps 4001 --epochs 20 \--snapshot-path snapshots --tensorboard-dir tensorboard \csv train_annotations.csv classes.csv
train_annotations.csv: file with information about all the images from training data, the format for each image should be made in the following template: <path_to_image>,<xmin>,<ymin>,<xmax>,<ymax>,<label>, where xmin, ymin, xmax, ymax are bounding boxes given in pixels, label is a name of class name;
labels.csv: file with labels for classes, should be followed by the template: <class_name>, <id>
After training 3 of epochs, we will get some accuracy. We can validate the model on the training data. Another option is to test accuracy on some real images:
To test the model on a validation set, we need to convert weights to a testing format:
!retinanet-convert-model snapshots/resnet50_csv_03.h5 weights/resnet50_csv_03.h5
To check results on a testing set:
!retinanet-evaluate csv val_annotations.csv classes.csv weights/resnet50_csv_03.h5
We can see that results after epochs of training are already good on a testing set, as the Mean Average Precision is 66%:
116 instances of class hammer1 with average precision: 0.7571113 instances of class hammer2 with average precision: 0.6968110 instances of class hammer3 with average precision: 0.8040123 instances of class plier1 with average precision: 0.5229119 instances of class plier2 with average precision: 0.5567122 instances of class plier3 with average precision: 0.8953126 instances of class screw1 with average precision: 0.4729152 instances of class screw2 with average precision: 0.5130131 instances of class screw3 with average precision: 0.7651mAP: 0.6649
at the same time not so good on the real environment: to check it, we can run the prediction of Retinanet on our image.
# show images inline%matplotlib inline# automatically reload modules when they have changed%load_ext autoreload%autoreload 2# import kerasimport keras# import keras_retinanetfrom keras_retinanet import modelsfrom keras_retinanet.utils.image import read_image_bgr, preprocess_image, resize_imagefrom keras_retinanet.utils.visualization import draw_box, draw_captionfrom keras_retinanet.utils.colors import label_color#from keras_retinanet.keras_retinanet.utils.gpu import setup_gpu# import miscellaneous modulesimport matplotlib.pyplot as pltimport cv2import osimport numpy as npimport time# use this to change which GPU to usegpu = 0# set the modified tf session as backend in keras#setup_gpu(gpu)# adjust this to point to your downloaded/trained model# models can be downloaded here: https://github.com/fizyr/keras-retinanet/releasesmodel_path = os.path.join('..', 'weights', 'resnet50_csv_04.h5')model = models.load_model(model_path, backbone_name='resnet50')# if the model is not converted to an inference model, use the line below# see: https://github.com/fizyr/keras-retinanet#converting-a-training-model-to-inference-model#model = models.convert_model(model)#print(model.summary())# load label to names mapping for visualization purposeslabels_to_names = {    0: 'hammer1',     1: 'hammer2',    3: 'hammer3',    4: 'plier1',    5: 'plier2',    6: 'plier3',    7: 'screw1',    8: 'screw2',    9: 'screw3'}# load imageimage = read_image_bgr('/content/img.jpg')# copy to draw ondraw = image.copy()draw = cv2.cvtColor(draw, cv2.COLOR_BGR2RGB)# preprocess image for networkimage = preprocess_image(image)image, scale = resize_image(image)# process imagestart = time.time()boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))print("processing time: ", time.time() - start)# correct for image scaleboxes /= scale# visualize detectionsfor box, score, label in zip(boxes[0], scores[0], labels[0]):    # scores are sorted so we can break    if score < 0.21:        break            color = label_color(label)        b = box.astype(int)    draw_box(draw, b, color=color)        caption = "{} {:.3f}".format(labels_to_names[label], score)    draw_caption(draw, b, caption)    plt.figure(figsize=(15, 15))plt.axis('off')plt.imshow(draw)plt.show()
I took 21% of the threshold, so the model predicts the instrument with 22% accuracy.
We can train the model further to improve the model, and train model for more than 4 epochs.
Conclusion.
This article is a part of the research, where we want to create a model, which would recognize the instruments in the video. We used KTH Handtool Dataset to improve the accuracy of a model. Experiments showed that additional images improve object detection. There are some improvements to the model that can be made (extra dataset can be used to improve instrument detection, more epochs for training can be done).
References
Focal Loss for Dense Object DetectionKeras-retinanetAn Introduction to Implementing Retinanet in Keras fro Multi Object Detection on Custom DatasetKTH Handtool Dataset
Focal Loss for Dense Object Detection
Keras-retinanet
An Introduction to Implementing Retinanet in Keras fro Multi Object Detection on Custom Dataset
KTH Handtool Dataset