Training YOLO Models for the RoboSub Competition Missions

Training YOLO Models for Enter the Pacific (Gate), Hydrothermal Vent (Buoy), Ocean Temperatures (Bins), Path, and Mapping (Torpedoes).

05-20-2024 - 06-20-2024

Deciding that specialized YOLO models were our best option to detect complicated objects of interest, rather than other forms of ML learning or conventional CV (computer vision) methods, we focused on training accurate YOLO models, each one individualized for detecting objects inside a certain mission. We used RoboFlow to split our training videos into frames and to annotate data. We also used the “augmentation” feature before exporting our dataset, allowing us to double or triple the number of images and perform manipulations such as shearing, flipping, or Gaussian Blur, to improve the performance of the YOLO model. We then used a Jupyter notebook provided by the Ultralytics community, which allowed us to quickly train a customizable YOLOv8 detection model. This notebook yielded both the model’s weights in a “.pt” file, as well as allowed us to validate the model by running the model with the best weights on the validation images, already set aside by RoboFlow automatically. Once we had the weights, we then used a conversion site in order to convert our “.pt” file into a “.blob” file, as well as generating a “.json” file among others that specified the settings for the model. This “.json” (settings) and “.blob” (weights) combination allows us to run YOLOv8 models on our OAK-D cameras present on Graey and Onyx.

Images: 

Fig. 1: Example of Data Augmentation on RoboFlow. Note that the top and bottom images are from the same image; the top image is the original, while the bottom image is a shear/rotation augmentation performed on the original. 

Fig. 2: Validating our “Bins model” – this step allowed us to test the performance of our model and let us know whether we needed to perform more pre-train augmentations as well as whether we needed to use more iterations (epochs).

Fig. 3: Changing our “.pt” file into a folder containing a “.bin”, “.onnx”, “.xml”, “.json”, and “.blob” files which combined contain all of the information from the “.pt” file necessary to run a YOLOv8 model on an OAK-D camera. Only the latter two, “.json” and “.blob”, are necessary for our purposes – those are the only files that need to be uploaded to our GitHub.