Adapt code for data generation

[DariusCorp]

Hello. Great job with the repo!
If I may, I want to experiment the use of Hungarian Algorithm with DNN in my project to check how the efficiency improves. My question for you is. Let's say my max DOAs are 16, it may be less I do not consider to be more, and lets say I have a 640x640 image with distances between objects. What should I adapt for the data generation part?

Much appreciated