Automatic pick and place with computer vision and RoboDK
The complete pick and place example.
This example shows how to build datasets with the Stray Robots Toolkit for picking and placing cardboard boxes of variable sizes within a RoboDK simulation environment. The souce code is available here: https://github.com/StrayRobots/stray/tree/main/examples/robodk.
The purpose of this simulation is to showcase how computer vision can be used in dynamic pick and place applications. In this case, our robot is tasked to move boxes rolling on one conveyor belt to the other.
The robot waits for boxes to come down the conveyor belt, where it detects the top corners of the box. We then compute a picking location at the center top of the cardboard box and command the robot to pick it up with its suction cup gripper. The robot then moves the box over to the other conveyor belt and starts over.
The main files in this project are:
Contains the main script containing the picking logic.
This script is used to collect datasets for training the object detector.
Contains the object detection logic.
Contains the simulation logic for the conveyor belt, reseting and spawning boxes etc.
Contains the logic for training a detection model.
Contains the logic for converting a PyTorch Lightning checkpoint into a serialized model that can be run in production.
This is the RoboDK simulation environment file.
A pretrained pytorch object detection model.
We recommend using Anaconda for package management. To create a new environment and install the dependencies, run:
conda create -n robodk python=3.8 && conda activate robodk pip install -r requirements.txt
Follow the instructions here to install the Stray Robots Toolkit.
Collecting a dataset and training an object detector
As is common these days, the object detection algorithm used is learning based. To train this algorithm, we need to collect example data from the robot’s workspace and annotate it to teach our robot to recognize the objects we want to pick. The rest of this section will show you how to collect a custom dataset. Alternatively you can download a sample dataset to proceed to model training and the next section. The sample includes scans of 20 boxes.
First open the simulation in the RoboDK UI by opening RoboDK, then select File > Open… and open the
To collect data, we provide the script
scan.py which runs the simulation. It can be triggerred to stop the production line and scan the current state of the line. This is done by running the following command while the picking simulation is open in the RoboDK UI:
python scan.py <out>
To stop the line and scan the conveyor belt, press the
The scans are saved in the path given by
out. For each performed scan, a subdirectory will be created within that directory which will contain the captured color and depth images, along with the camera poses.
After scanning, we need to process the scans to compute the 3D representation from the captured camera and depth images. This is done using the
stray integrate command. Run it with:
stray integrate <out> --skip-mapping --voxel-size 0.005
As in this case, our camera is mounted on our calibrated robot, we use
--skip-mapping parameter to tell the system that we know the camera poses and that these do not have to be inferred.
We then annotate each of the scans by opening the up in the Stray Studio user interface. A scan can be opened with:
stray studio boxes/0001
In this case, we want to detect the top corners of each box. Therefore, we opt to annotate each scanned box with a rectangle annotation type. Here is what an annotated scan looks like:
Training a picking model
The model training can be run with
python train.py <out> --eval-folder <path-to-eval-data> --num-epochs <num-epochs> --primitive rectangle --num-instances 1
The main available paramaters are:
Path to the directory containing the scans
Path to the directory containing evaluation scans, for testing purposes this can be the same as
For how many iterations the model should be trained. 100 is enough for the example dataset of 20 scans.
Which primitive type to use for determining the keypoints on the box. In this example we use
How many instances of
primitiveshould be detected. The instances should be labeled with unique instance ids in Studio, ranging from
num-instances - 1. In this case there is only one instance per scene/image.
--batch-size (default 4)
Batch size to use during training, adjust this as high as possible as long as there are no memory errors
For additional settings refer to the
The training is implemented using PyTorch Lightning. Logs of the training are saved to
./train_logs but this can be adjusted with the
--logs flag. Intermediate versions of the model are saved to
Once the training is completed, we can pick one of the checkpoints from
./train_logs/default/version_x/checkpoints and convert it into a serialized model that can be used in production and the example picking script.
The model can be converted with
python convert_model.py <path-to-checkpoint> and it will be saved as
Running the picking simulation with the trained model
Again, open the simulation in the RoboDK UI by opening RoboDK, then select File > Open… and open the
To run the simulation with the trained model, use the command
python pick.py --model model.pt.
You should now see the simulation running with the robot picking boxes and moving them over to the other conveyor belt.