How to build a graph using yolo output

[samahwaleed]

I have done object detection using yolo v5 on my custom dataset and I saved the results in txt files. The txt file has label, bbox coordinator and the score and I want to know how to build a graph using those results.

Thanks.

[rusty1s]

You can represent each bounding box as a node, and assign features to it, e.g.,

1. The (normalized) width and height of the bounding box
2. The scores
3. Some image features of each bounding box as given by your pre-trained model,

and save them to a Data object data = Data(x=x, pos=pos) where x denotes a matrix of shape [num_nodes, num_features], and pos denotes the center of each bounding box as a [num_nodes, 2] matrix.
You can then create a graph from it using k-nearest neighbors, e.g.:

from torch_geometric.transforms import KNNGraph
transform = KNNGraph(6)
data = transform(data)