Hugging Face's logo

Spaces:
pablovela5620
/
MiniDPVO
Sleeping

App Files Community
1
MiniDPVO
File size: 4,511 Bytes 
899c526
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
 
#include <torch/extension.h>
#include <vector>
#include <unordered_map>
#include <algorithm>
#include <iostream>


std::vector<torch::Tensor> cuda_ba(
    torch::Tensor poses,
    torch::Tensor patches,
    torch::Tensor intrinsics,
    torch::Tensor target,
    torch::Tensor weight,
    torch::Tensor lmbda,
    torch::Tensor ii,
    torch::Tensor jj, 
    torch::Tensor kk,
    int t0, int t1, int iterations);


torch::Tensor cuda_reproject(
    torch::Tensor poses,
    torch::Tensor patches,
    torch::Tensor intrinsics,
    torch::Tensor ii,
    torch::Tensor jj, 
    torch::Tensor kk);

std::vector<torch::Tensor> ba(
    torch::Tensor poses,
    torch::Tensor patches,
    torch::Tensor intrinsics,
    torch::Tensor target,
    torch::Tensor weight,
    torch::Tensor lmbda,
    torch::Tensor ii,
    torch::Tensor jj, 
    torch::Tensor kk,
    int t0, int t1, int iterations) {
  return cuda_ba(poses, patches, intrinsics, target, weight, lmbda, ii, jj, kk, t0, t1, iterations);
}


torch::Tensor reproject(
    torch::Tensor poses,
    torch::Tensor patches,
    torch::Tensor intrinsics,
    torch::Tensor ii,
    torch::Tensor jj, 
    torch::Tensor kk) {
  return cuda_reproject(poses, patches, intrinsics, ii, jj, kk);
}

// std::vector<torch::Tensor> neighbors(torch::Tensor ii, torch::Tensor jj)
// {
//   ii = ii.to(torch::kCPU);
//   jj = jj.to(torch::kCPU);
//   auto ii_data = ii.accessor<long,1>();
//   auto jj_data = jj.accessor<long,1>();

//   std::unordered_map<long, std::vector<long>> graph;
//   std::unordered_map<long, std::vector<long>> index;
//   for (int i=0; i < ii.size(0); i++) {
//     const long ix = ii_data[i];
//     const long jx = jj_data[i];
//     if (graph.find(ix) == graph.end()) {
//       graph[ix] = std::vector<long>();
//       index[ix] = std::vector<long>();
//     }
//     graph[ix].push_back(jx);
//     index[ix].push_back( i);
//   }

//   auto opts = torch::TensorOptions().dtype(torch::kInt64);
//   torch::Tensor ix = torch::empty({ii.size(0)}, opts);
//   torch::Tensor jx = torch::empty({jj.size(0)}, opts); 

//   auto ix_data = ix.accessor<long,1>();
//   auto jx_data = jx.accessor<long,1>();

//   for (std::pair<long, std::vector<long>> element : graph) {
//     std::vector<long>& v = graph[element.first];
//     std::vector<long>& idx = index[element.first];

//     std::stable_sort(idx.begin(), idx.end(),
//        [&v](size_t i, size_t j) {return v[i] < v[j];});

//     ix_data[idx.front()] = -1;
//     jx_data[idx.back()]  = -1;

//     for (int i=0; i < idx.size(); i++) {
//       ix_data[idx[i]] = (i > 0) ? idx[i-1] : -1;
//       jx_data[idx[i]] = (i < idx.size() - 1) ? idx[i+1] : -1;
//     }
//   }

//   ix = ix.to(torch::kCUDA);
//   jx = jx.to(torch::kCUDA);

//   return {ix, jx};
// }


std::vector<torch::Tensor> neighbors(torch::Tensor ii, torch::Tensor jj)
{

  auto tup = torch::_unique(ii, true, true);
  torch::Tensor uniq = std::get<0>(tup).to(torch::kCPU);
  torch::Tensor perm = std::get<1>(tup).to(torch::kCPU);

  jj = jj.to(torch::kCPU);
  auto jj_accessor = jj.accessor<long,1>();

  auto perm_accessor = perm.accessor<long,1>();
  std::vector<std::vector<long>> index(uniq.size(0));
  for (int i=0; i < ii.size(0); i++) {
    index[perm_accessor[i]].push_back(i);
  }

  auto opts = torch::TensorOptions().dtype(torch::kInt64);
  torch::Tensor ix = torch::empty({ii.size(0)}, opts);
  torch::Tensor jx = torch::empty({ii.size(0)}, opts); 

  auto ix_accessor = ix.accessor<long,1>();
  auto jx_accessor = jx.accessor<long,1>();

  for (int i=0; i<uniq.size(0); i++) {
    std::vector<long>& idx = index[i];
    std::stable_sort(idx.begin(), idx.end(),
       [&jj_accessor](size_t i, size_t j) {return jj_accessor[i] < jj_accessor[j];});

    for (int i=0; i < idx.size(); i++) {
      ix_accessor[idx[i]] = (i > 0) ? idx[i-1] : -1;
      jx_accessor[idx[i]] = (i < idx.size() - 1) ? idx[i+1] : -1;
    }
  }

  // for (int i=0; i<ii.size(0); i++) {
  //   std::cout << jj_accessor[i] << " ";
  //   if (ix_accessor[i] >= 0) std::cout << jj_accessor[ix_accessor[i]] << " ";
  //   if (jx_accessor[i] >= 0) std::cout << jj_accessor[jx_accessor[i]] << " ";
  //   std::cout << std::endl;
  // }

  ix = ix.to(torch::kCUDA);
  jx = jx.to(torch::kCUDA);

  return {ix, jx};
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("forward", &ba, "BA forward operator");
  m.def("neighbors", &neighbors, "temporal neighboor indicies");
  m.def("reproject", &reproject, "temporal neighboor indicies");

}