Chroma with YOLOv5

[atroyn]

How to use Chroma with YOLOv5

See the chroma docs and quick-start guide for general information about Chroma.

This PR is intended to illustrate the changes necessary to attach Chroma to a more complex model like the popular YOLO family of object detection models. It is not intended to be merged upstream. We use the YOLOv5 repository from Ultralytics as the base.

⚠️ The PR assumes that you have Chroma installed. If not, follow the quick-start guide to install Chroma.

get_embeddings.py is a fully functional script for extracting embeddings from YOLOv5. It takes parts from the detection and validation scripts, but focuses on logging the necessary information to the Chroma client.

It also contains a function for associating labels to detections when appropriate. In most cases, this will not require a new function to be written - this is an artifact of how the existing YOLOv5 repo is set up.

We also modify the following:

In models/yolo.py, the Detect class is modified so that the model outputs embeddings alongside detections in the forward pass. This could also be accomplished with a forward hook, without the need to modify the model, but requires some additional complexity. Note that this change does not require the model to be re-trained; it's fine to load any existing weight dict.

In utils/general.py, we modify non_max_suppression so that embeddings are filtered alongside predictions when they are passed in, such that each prediction which makes it through is correctly associated with its embedding.

utils/dataloaders.py is modified so that we can work with datasets which have images, but not labels corresponding to each image. This is mostly a convenience.

Try the branch out for yourself!