Spaces:
Running on Zero
Running on Zero
File size: 4,690 Bytes
7266be4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | /*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <torch/extension.h>
#include <torch/torch.h>
#include <list>
#include <numeric>
#include <queue>
#include <tuple>
#include "iou_utils.h"
std::tuple<at::Tensor, at::Tensor> IoUBox3DCpu(
const at::Tensor& boxes1,
const at::Tensor& boxes2) {
const int N = boxes1.size(0);
const int M = boxes2.size(0);
auto float_opts = boxes1.options().dtype(torch::kFloat32);
torch::Tensor vols = torch::zeros({N, M}, float_opts);
torch::Tensor ious = torch::zeros({N, M}, float_opts);
// Create tensor accessors
auto boxes1_a = boxes1.accessor<float, 3>();
auto boxes2_a = boxes2.accessor<float, 3>();
auto vols_a = vols.accessor<float, 2>();
auto ious_a = ious.accessor<float, 2>();
// Iterate through the N boxes in boxes1
for (int n = 0; n < N; ++n) {
const auto& box1 = boxes1_a[n];
// Convert to vector of face vertices i.e. effectively (F, 3, 3)
// face_verts is a data type defined in iou_utils.h
const face_verts box1_tris = GetBoxTris(box1);
// Calculate the position of the center of the box which is used in
// several calculations. This requires a tensor as input.
const vec3<float> box1_center = BoxCenter(boxes1[n]);
// Convert to vector of face vertices i.e. effectively (P, 4, 3)
const face_verts box1_planes = GetBoxPlanes(box1);
// Get Box Volumes
const float box1_vol = BoxVolume(box1_tris, box1_center);
// Iterate through the M boxes in boxes2
for (int m = 0; m < M; ++m) {
// Repeat above steps for box2
// TODO: check if caching these value helps performance.
const auto& box2 = boxes2_a[m];
const face_verts box2_tris = GetBoxTris(box2);
const vec3<float> box2_center = BoxCenter(boxes2[m]);
const face_verts box2_planes = GetBoxPlanes(box2);
const float box2_vol = BoxVolume(box2_tris, box2_center);
// Every triangle in one box will be compared to each plane in the other
// box. There are 3 possible outcomes:
// 1. If the triangle is fully inside, then it will
// remain as is.
// 2. If the triagnle it is fully outside, it will be removed.
// 3. If the triangle intersects with the (infinite) plane, it
// will be broken into subtriangles such that each subtriangle is full
// inside the plane and part of the intersecting tetrahedron.
// Tris in Box1 -> Planes in Box2
face_verts box1_intersect =
BoxIntersections(box1_tris, box2_planes, box2_center);
// Tris in Box2 -> Planes in Box1
face_verts box2_intersect =
BoxIntersections(box2_tris, box1_planes, box1_center);
// If there are overlapping regions in Box2, remove any coplanar faces
if (box2_intersect.size() > 0) {
// Identify if any triangles in Box2 are coplanar with Box1
std::vector<int> tri2_keep(box2_intersect.size());
std::fill(tri2_keep.begin(), tri2_keep.end(), 1);
for (int b1 = 0; b1 < box1_intersect.size(); ++b1) {
for (int b2 = 0; b2 < box2_intersect.size(); ++b2) {
const bool is_coplanar =
IsCoplanarTriTri(box1_intersect[b1], box2_intersect[b2]);
const float area = FaceArea(box1_intersect[b1]);
if ((is_coplanar) && (area > aEpsilon)) {
tri2_keep[b2] = 0;
}
}
}
// Keep only the non coplanar triangles in Box2 - add them to the
// Box1 triangles.
for (int b2 = 0; b2 < box2_intersect.size(); ++b2) {
if (tri2_keep[b2] == 1) {
box1_intersect.push_back((box2_intersect[b2]));
}
}
}
// Initialize the vol and iou to 0.0 in case there are no triangles
// in the intersecting shape.
float vol = 0.0;
float iou = 0.0;
// If there are triangles in the intersecting shape
if (box1_intersect.size() > 0) {
// The intersecting shape is a polyhedron made up of the
// triangular faces that are all now in box1_intersect.
// Calculate the polyhedron center
const vec3<float> polyhedron_center = PolyhedronCenter(box1_intersect);
// Compute intersecting polyhedron volume
vol = BoxVolume(box1_intersect, polyhedron_center);
// Compute IoU
iou = vol / (box1_vol + box2_vol - vol);
}
// Save out volume and IoU
vols_a[n][m] = vol;
ious_a[n][m] = iou;
}
}
return std::make_tuple(vols, ious);
}
|