I need two Python Pytorch tensor functions to be translated/rewritten in Java.

The goal is to complete Yolov5 image segmentation inferences in Java. So far, I am using OnnxRuntime to send the image tensor to the model and receive the inference tensors.

However, the functions used for converting the output tensors to masks are more challenging and need to be completed.

We need to convert the following python functions to Java:

non_max_suppression https://github.com/ultralytics/yolov5/blob/064365d8683fd002e9ad789c1e91fa3d021b44f0/utils/general.py#L884

process_mask
https://github.com/ultralytics/yolov5/blob/064365d8683fd002e9ad789c1e91fa3d021b44f0/utils/segment/general.py#L43

The job is successful once we can apply non_max_suppression and process_mask to our arrays/tensors in Java. However, there should not be any significant differences in speed between the Python and Java Implementations.

Please send a proposal if you have strong experience with numpy/pytorch, tensors, arrays, ONNX and feel you can complete the job AFTER looking at the two Python functions displayed above.