naver/PUMP
1
1// Copyright 2022-present NAVER Corp.2// CC BY-NC-SA 4.03// Available only for non-commercial use4 5#include <torch/extension.h>6using namespace torch::indexing; // Slice7#include <vector>8 9#define MIN(x, y) ((x) < (y) ? (x) : (y))10#define MAX(x, y) ((x) < (y) ? (y) : (x))11#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")12#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")13#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)14 15inline Slice sl(bool x) {16 if (x)17 return Slice(0, -1);18 else19 return Slice(1, None);20}21 22torch::Tensor forward_agg_cuda( int level, float norm, const torch::Tensor lower, 23 const at::optional<at::Tensor> weights, torch::Tensor upper );24 25std::vector<torch::Tensor> forward_agg( int level, float norm, const torch::Tensor lower, 26 const at::optional<at::Tensor> weights = at::nullopt ) {27 TORCH_CHECK(level >= 1, "level must be >= 1");28 TORCH_CHECK(lower.dim() == 4, "input must have 4 dimensions");29 const auto LH1 = lower.size(0);30 const auto LW1 = lower.size(1);31 const auto LH2 = lower.size(2);32 const auto LW2 = lower.size(3);33 if (weights) TORCH_CHECK(weights->size(0) == LH1 && weights->size(1) == LW1, "weights should have shape == lower.shape[:2]");34 const auto UH1 = (level == 1) ? LH1+1 : LH1;35 const auto UW1 = (level == 1) ? LW1+1 : LW1;36 37 TORCH_CHECK(lower.is_cuda())38 auto upper = torch::zeros({UH1, UW1, LH2, LW2}, lower.options());39 torch::Tensor new_weights = forward_agg_cuda( level, norm, lower, weights, upper );40 return {upper, new_weights};41}42 43 44torch::Tensor forward_pool_agg_cuda( int level, float norm, const torch::Tensor lower,45 const at::optional<at::Tensor> weights, torch::Tensor upper );46 47std::vector<torch::Tensor> forward_pool_agg( int level, float norm, const torch::Tensor lower, 48 const at::optional<at::Tensor> weights = at::nullopt ) {49 TORCH_CHECK(level >= 1, "level must be >= 1");50 TORCH_CHECK(lower.dim() == 4, "input must have 4 dimensions");51 const auto LH1 = lower.size(0);52 const auto LW1 = lower.size(1);53 const auto LH2 = lower.size(2);54 const auto LW2 = lower.size(3);55 if (weights) TORCH_CHECK(weights->size(0) == LH1 && weights->size(1) == LW1, "weights should have shape == lower.shape[:2]");56 const auto UH1 = (level == 1) ? LH1+1 : LH1;57 const auto UW1 = (level == 1) ? LW1+1 : LW1;58 59 TORCH_CHECK(lower.is_cuda())60 auto upper = torch::zeros({UH1, UW1, 1+(LH2-1)/2, 1+(LW2-1)/2}, lower.options());61 torch::Tensor new_weights = forward_pool_agg_cuda( level, norm, lower, weights, upper );62 return {upper, new_weights};63}64 65// forward declaration66void backward_agg_unpool_cuda( int level, const torch::Tensor upper, torch::Tensor lower, bool exclude_borders );67 68void backward_agg_unpool( int level, const torch::Tensor upper, torch::Tensor lower, bool exclude_borders = true ) {69 TORCH_CHECK(level >= 1, "level must be >= 1");70 TORCH_CHECK( upper.dim() == 4 && lower.dim() == 4, "inputs should be 4-dimensional" );71 72 TORCH_CHECK(upper.is_cuda() && lower.is_cuda())73 backward_agg_unpool_cuda(level, upper, lower, exclude_borders);74}75 76 77void max_pool3d_cuda( const torch::Tensor tensor, const int kernel_size, const int stride,78 torch::Tensor maxima, torch::Tensor indices );79 80std::vector<torch::Tensor> max_pool3d( const torch::Tensor tensor, const int kernel_size, const int stride ) {81 TORCH_CHECK(tensor.dim() == 4, "tensor should be 4-dimensional: BxCxHxW");82 TORCH_CHECK( 1 <= kernel_size, "bad kernel size %d", kernel_size );83 TORCH_CHECK( 1 <= stride, "bad stride %d", stride );84 const int IB = tensor.size(0);85 const int IH = tensor.size(2); // input height86 const int IW = tensor.size(3); // input width87 88 // output size89 const int OH = 1 + (IH - kernel_size) / stride;90 const int OW = 1 + (IW - kernel_size) / stride;91 92 torch::Tensor maxima = torch::empty({IB, OH, OW}, tensor.options());93 torch::Tensor indices = torch::empty({IB, OH, OW}, tensor.options().dtype(torch::kInt64));94 95 if (tensor.is_cuda())96 max_pool3d_cuda( tensor, kernel_size, stride, maxima, indices );97 else98 TORCH_CHECK(false, "CPU max_pool3d not implemented yet");99 return {maxima, indices};100}101 102static inline float ptdot( const float* m, float x, float y ) {103 return x*m[0] + y*m[1] + m[2];104}105 106static inline float pow2(float v) {107 return v*v;108}109 110void merge_corres_cpu( const torch::Tensor corres, int offset, const torch::Tensor _inv_rot, 111 float dmax, torch::Tensor all_corres, const int all_step ) {112 const int H = corres.size(0);113 const int W = corres.size(1);114 const float tol = 2*2; // squared115 dmax *= dmax; // squared116 117 TORCH_CHECK( _inv_rot.is_contiguous() );118 const float* inv_rot = _inv_rot.data_ptr<float>();119 120 auto corres_a = corres.accessor<float,3>();121 auto all_corres_a = all_corres.accessor<float,3>();122 123 // for each bin of the final histograms, we get the nearest-neighbour bin in corres0 and corres1124 for (int j=0; j<all_corres.size(0); j++) 125 for (int i=0; i<all_corres.size(1); i++) {126 // printf("accessing all_corres[%d,%d]", j, i);127 auto all_cor = all_corres_a[j][i];128 129 // center of the bin in the reference frame130 float x = i*all_step + all_step/2;131 float y = j*all_step + all_step/2;132 // printf(" -> (%g,%g) in ref img", x, y);133 134 // center of the bin on the rescaled+rotated image135 float xr = ptdot( inv_rot + 0, x, y ); 136 float yr = ptdot( inv_rot + 3, x, y );137 // printf(" -> (%g,%g) in rescaled", xr, yr);138 139 // iterate on the nearby bins140 int xb = (int)(0.5+ xr/4); // rescaled+rotated desc always has step 4141 int yb = (int)(0.5+ yr/4);142 // printf(" -> (%d,%d) in bins\n", xb, yb);143 144 float best = dmax;145 for (int v = MAX(0,yb-1); v <= MIN(H,yb+1); v++)146 for (int u = MAX(0,xb-1); u <= MIN(W,xb+1); u++) {147 // assert( v >= 0 && v < corres_a.size(0) );148 // assert( u >= 0 && u < corres_a.size(1) );149 auto cor = corres_a[v][u];150 float d = pow2(cor[offset]-x) + pow2(cor[offset+1]-y);151 if( d < best ) best = d;152 }153 154 for (int v = MAX(0,yb-1); v <= MIN(H,yb+1); v++)155 for (int u = MAX(0,xb-1); u <= MIN(W,xb+1); u++) {156 // assert( v >= 0 && v < corres_a.size(0) );157 // assert( u >= 0 && u < corres_a.size(1) );158 auto cor = corres_a[v][u];159 float d = pow2(cor[offset]-x) + pow2(cor[offset+1]-y);160 if( d <= tol*best ) { // spatially close161 // merge correspondence if score is better than actual162 // printf("update all_corres[%d,%d]\n", v,u);163 if( cor[4] > all_cor[4] )164 for (int k = 0; k < all_corres.size(2); k++) 165 all_cor[k] = cor[k];166 }167 }168 }169}170 171void merge_corres_cuda( const torch::Tensor corres, int offset, const torch::Tensor inv_rot, 172 float dmax, torch::Tensor all_corres, const int all_step );173 174void merge_corres( const torch::Tensor corres, int offset, const torch::Tensor rot, 175 torch::Tensor all_corres, const int all_step ) {176 TORCH_CHECK( corres.dim() == 3 && corres.size(2) == 6, "corres.shape should be (H,W,6)" );177 TORCH_CHECK( all_corres.dim() == 3 && all_corres.size(2) == 6, "all_corres.shape should be (H,W,6)" );178 179 float dmax = 8 * torch::sqrt(torch::det(rot)).item<float>();180 torch::Tensor inv_rot = torch::inverse(rot).contiguous();181 182 if (all_corres.is_cuda()) 183 merge_corres_cuda( corres, offset, inv_rot, dmax, all_corres, all_step );184 else185 merge_corres_cpu( corres, offset, inv_rot, dmax, all_corres, all_step );186}187 188 189void mask_correlations_radial_cuda( torch::Tensor corr, const torch::Tensor targets, 190 const float radius, const float alpha);191 192void mask_correlations_radial( torch::Tensor corr, const torch::Tensor targets, 193 const float radius, const float alpha) {194 // radius: protected area in pixels around each target center195 // alpha: in [0,1]. If alpha = 0: no effect. If alpha = 1: full effect.196 TORCH_CHECK( corr.dim() == 4 );197 TORCH_CHECK( targets.dim() == 3 );198 TORCH_CHECK( targets.size(0) == corr.size(0) && targets.size(1) == corr.size(1) && targets.size(2) == 2, 199 "correlations and targets should have the same shape[:2]" );200 201 if (corr.is_cuda()) 202 mask_correlations_radial_cuda( corr, targets, radius, alpha );203 else204 TORCH_CHECK(false, "TODO");205}206 207 208PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {209 m.def("forward_agg", &forward_agg, "forward aggregation (CUDA)");210 m.def("forward_pool_agg", &forward_pool_agg, "forward pooling and aggregation (CUDA)");211 m.def("backward_agg_unpool", &backward_agg_unpool, "backward sparse-conv and max-unpooling (C++ & CUDA)");212 m.def("max_pool3d", &max_pool3d, "max_pool3d that can handle big inputs (CUDA)");213 m.def("merge_corres_one_side", &merge_corres, "merge correspondences on CPU or GPU" );214 m.def("mask_correlations_radial", &mask_correlations_radial, "mask correlations radially (CUDA)" );215}216