Creation of the database (labelling.py)
To create your database, you need to take pictures of the objects you want to use. Take at least 20 pictures of each object to get good results.
The aim is to take pictures of each object under a multitude of angles and different lighting conditions. The pictures will be stored in a folder named with the name of the concerned object, inside the “data” folder.
Tracking of the objects (utils.py)
- Image shooting (“take_workspace_img()” function)
Use the TCP API to ask the robot to send an image, to crop it and to compensate the lens’ distortion.
- Calculation of the mask (“objs_mask()” function)
Use the cv2.cvtColor() to modify the image colorimetry from RGB to HLS, then use cv2.inRange() to get a mask which approximately delineates the objects to detect.
In order to only keep objects with a sufficient surface, combine the cv2.dilate() and cv2.erode() functions to remove images’ impurities.
We finally obtain a black and white picture corresponding to the shape of the objects placed on the workspace.
Use cv2.findContours() to obtain the list of the outline of the objects being on the previously calculated mask.
Then, calculate the center as well as the angle of the objects with the help of the vision functions of the TCP/IP API get_contour_barycenter() and get_contour_angle().
With cv2.minAreaRect() we obtain a square containing the smallest object and use this information to extract the object from the image and put it vertically (giving the same orientation to these images makes the recognition easier for TensorFlow).
Launch training.py or click on the “Train” button on the graphic interface. This creates a TensorFlow model (neural network). Then create a list which contains all the images from the “data” folder and a list which contains the label corresponding to each image.
Use [model].fit([images], [labels]) to train the model with the database.
When the training is over, test the model’s performances and save it in the “model” folder.
Prediction (robot.py / robot_gui.py)
Launch robot.py or the graphic interface’s “Play” menu:
Use the previously trained model to recognize the different objects on the workspace.