|
|
|
#pragma once
|
|
#include <torch/types.h>
|
|
|
|
namespace detectron2 {
|
|
|
|
at::Tensor box_iou_rotated_cpu(
|
|
const at::Tensor& boxes1,
|
|
const at::Tensor& boxes2);
|
|
|
|
#if defined(WITH_CUDA) || defined(WITH_HIP)
|
|
at::Tensor box_iou_rotated_cuda(
|
|
const at::Tensor& boxes1,
|
|
const at::Tensor& boxes2);
|
|
#endif
|
|
|
|
|
|
|
|
|
|
inline at::Tensor box_iou_rotated(
|
|
const at::Tensor& boxes1,
|
|
const at::Tensor& boxes2) {
|
|
assert(boxes1.device().is_cuda() == boxes2.device().is_cuda());
|
|
if (boxes1.device().is_cuda()) {
|
|
#if defined(WITH_CUDA) || defined(WITH_HIP)
|
|
return box_iou_rotated_cuda(boxes1.contiguous(), boxes2.contiguous());
|
|
#else
|
|
AT_ERROR("Detectron2 is not compiled with GPU support!");
|
|
#endif
|
|
}
|
|
|
|
return box_iou_rotated_cpu(boxes1.contiguous(), boxes2.contiguous());
|
|
}
|
|
|
|
}
|
|
|