Skip to content
Permalink
Browse files

Releasing code for the network (heatmap prediction) and the 9DoF alig…

…nment (LM optimization).
  • Loading branch information...
skanti committed Jul 10, 2019
1 parent 2ab7fa6 commit 11a57a66842493182b464c03ec05c8fa0ddd5b84
Showing with 12,426 additions and 28 deletions.
  1. +29 −0 Network/base/CSVHelper.py
  2. +90 −0 Network/base/ExtendBox/ExtendBox.cpp
  3. +7 −0 Network/base/ExtendBox/makefile
  4. +10 −0 Network/base/JSONHelper.py
  5. +124 −0 Network/base/Vox.py
  6. +105 −0 Network/base/error_metrics.py
  7. +22 −0 Network/base/kernels.py
  8. +29 −0 Network/base/loss_container.py
  9. +9 −0 Network/base/nms.py
  10. +136 −0 Network/base/pytorch_dataloader.py
  11. +111 −0 Network/base/sample_loader.py
  12. +7 −0 Network/base/test.py
  13. +121 −0 Network/base/user_query.py
  14. +183 −0 Network/base/util.py
  15. +100 −0 Network/pytorch/HeatmapStatistics.py
  16. +66 −0 Network/pytorch/IsMatchStatistics.py
  17. +7 −0 Network/pytorch/README.md
  18. +67 −0 Network/pytorch/SaveOutput.py
  19. +38 −0 Network/pytorch/custom_dataloader.py
  20. +51 −0 Network/pytorch/dataloader.py
  21. +77 −0 Network/pytorch/losses.py
  22. +303 −0 Network/pytorch/main.py
  23. +283 −0 Network/pytorch/model.py
  24. +31 −0 Network/pytorch/run.sh
  25. +18 −0 Routines/AlignmentHeatmap/CMakeLists.txt
  26. +155 −0 Routines/AlignmentHeatmap/Grid.h
  27. +45 −0 Routines/AlignmentHeatmap/MatrixLib/Base.props
  28. +136 −0 Routines/AlignmentHeatmap/MatrixLib/MatrixLib.vcxproj
  29. +141 −0 Routines/AlignmentHeatmap/MatrixLib/MatrixLib.vcxproj.filters
  30. +8 −0 Routines/AlignmentHeatmap/MatrixLib/MatrixLibInclude.h
  31. +2 −0 Routines/AlignmentHeatmap/MatrixLib/MatrixLibSerialization.h
  32. +8 −0 Routines/AlignmentHeatmap/MatrixLib/main.cpp
  33. +457 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/Mat2.h
  34. +547 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/Mat3.h
  35. +757 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/Mat4.h
  36. +896 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/MatrixOperations.h
  37. +11 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/MatrixStructuresInclude.h
  38. +399 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/Vec2.h
  39. +421 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/Vec3.h
  40. +435 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/Vec4.h
  41. +106 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/VecArrays.h
  42. +328 −0 Routines/AlignmentHeatmap/MatrixLib/matrix_structures/VecX.h
  43. +142 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/AffineIncrement.h
  44. +188 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/AffinePose.h
  45. +68 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/Pose.h
  46. +40 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseAliases.h
  47. +68 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseArrays.h
  48. +46 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseCommon.h
  49. +55 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseIncrement.h
  50. +296 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseOperations.h
  51. +22 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseProcessing.h
  52. +14 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/PoseStructuresInclude.h
  53. +257 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/RigidPose.h
  54. +161 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/SE3.h
  55. +229 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/SO3.h
  56. +335 −0 Routines/AlignmentHeatmap/MatrixLib/pose_structures/UnitQuaternion.h
  57. BIN Routines/AlignmentHeatmap/MatrixLib/stdafx.cpp
  58. BIN Routines/AlignmentHeatmap/MatrixLib/stdafx.h
  59. BIN Routines/AlignmentHeatmap/MatrixLib/targetver.h
  60. +14 −0 Routines/AlignmentHeatmap/MatrixLib/utils/Constants.h
  61. +9 −0 Routines/AlignmentHeatmap/MatrixLib/utils/LibIncludeCPU.h
  62. +42 −0 Routines/AlignmentHeatmap/MatrixLib/utils/Math.h
  63. +13 −0 Routines/AlignmentHeatmap/MatrixLib/utils/Promotion.h
  64. +101 −0 Routines/AlignmentHeatmap/MatrixLib/utils/Serialization.h
  65. +56 −0 Routines/AlignmentHeatmap/MatrixLib/utils/SoAConversions.h
  66. +4 −0 Routines/AlignmentHeatmap/MatrixLib/utils/SoloIncludeCPU.cpp
  67. +39 −0 Routines/AlignmentHeatmap/MatrixLib/utils/SoloIncludeCPU.h
  68. +137 −0 Routines/AlignmentHeatmap/MatrixLib/utils/SoloSupport.h
  69. +1,449 −0 Routines/AlignmentHeatmap/MatrixLib/utils/VectorMath.h
  70. +456 −0 Routines/AlignmentHeatmap/main-new.cpp
  71. +374 −0 Routines/AlignmentHeatmap/main-refractor.cpp
  72. +660 −0 Routines/AlignmentHeatmap/main.cpp
  73. +2 −0 Routines/AlignmentHeatmap/run.sh
  74. +102 −6 Routines/Base/SE3.h
  75. +10 −0 Routines/CropCentered/main.cpp
  76. +1 −1 Routines/CropCentered/makefile
  77. +1 −1 Routines/Keypoints2Grid/makefile
  78. +15 −12 Routines/Script/Alignment2Mesh.py
  79. +128 −0 Routines/Script/Alignment9DoF.py
  80. +16 −8 Routines/Script/GenerateCorrespondences.py
@@ -0,0 +1,29 @@
import csv

def read(filename, skip_header=False):
rows = []
with open(filename, 'r') as csvfile:
reader = csv.reader(csvfile, delimiter=',')
if skip_header:
next(reader, None)
for row in reader:
if len(row) == 1:
rows.append(row[0])
else:
rows.append(row)
return rows

def read_as_dict(filename, split=","):
rows = []
with open(filename, 'r') as csvfile:
for row in csv.DictReader(csvfile, delimiter=','):
rows.append(row)
return rows

def write(filename, rows, mode="w"):
with open(filename, mode) as csvfile:
writer = csv.writer(csvfile, delimiter=',')
if type(rows[0]) is tuple:
writer.writerows(rows)
else:
writer.writerow(rows)
@@ -0,0 +1,90 @@
#include <eigen3/Eigen/Dense>
#include <vector>

#include <pybind11/stl.h>
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

namespace py = pybind11;

#define IS_IN_RANGE3(x1, y1, z1) ((x1) >= 0 && (x1) < dimx && (y1) >= 0 && (y1) < dimy && (z1) >= 0 && (z1) < dimz)
#define INDEX3(dz, dy, dx) (IS_IN_RANGE3((x + dx), (y + dy), (z + dz)) ? *ptr.data(i, 0, z + dz, y + dy, x + dx) : 0)
#define CHECK3() \
(INDEX3(-1, -1, -1) || INDEX3(-1, -1, 0) || INDEX3(-1, -1, 1) \
|| INDEX3(-1, 0, -1) || INDEX3(-1, 0, 0) || INDEX3(-1, 0, 1) \
|| INDEX3(-1, 1, -1) || INDEX3(-1, 1, 0) || INDEX3(-1, 1, 1) \
|| INDEX3(0, -1, -1) || INDEX3(0, -1, 0) || INDEX3(0, -1, 1) \
|| INDEX3(0, 0, -1) || INDEX3(0, 0, 0) || INDEX3(0, 0, 1) \
|| INDEX3(0, 1, -1) || INDEX3(0, 1, 0) || INDEX3(0, 1, 1) \
|| INDEX3(1, -1, -1) || INDEX3(1, -1, 0) || INDEX3(1, -1, 1) \
|| INDEX3(1, 0, -1) || INDEX3(1, 0, 0) || INDEX3(1, 0, 1) \
|| INDEX3(1, 1, -1) || INDEX3(1, 1, 0) || INDEX3(1, 1, 1)) \

#define IS_IN_RANGE1(z1) ((z1) >= 0 && (z1) < dimz)
#define INDEX1(dz) (IS_IN_RANGE1((z + dz)) ? *ptr.data(i, z + dz) : 0)
#define CHECK1() \
(INDEX1(-1) || INDEX1(0) || INDEX1(1)) \

int count_tp1(py::array_t<bool> &p, py::array_t<bool> &gt) {
assert(p.ndim() == 2);
assert(gt.ndim() == 2);
const int n_batch = p.shape(0);
const int dimz = p.shape(1);

auto& ptr = p;
int counter = 0;
for (int i = 0; i < n_batch; i++)
for (int z = 0 ; z < dimz ; z++)
if (*gt.data(i, z)) {
counter += CHECK1();
}
return counter;
}

int count_tp3(py::array_t<bool> &p, py::array_t<bool> &gt) {
assert(p.ndim() == 5);
assert(gt.ndim() == 5);
const int n_batch = p.shape(0);
assert(p.shape(1) == 1);
const int dimz = p.shape(2);
const int dimy = p.shape(3);
const int dimx = p.shape(4);

auto& ptr = p;
int counter = 0;
for (int i = 0; i < n_batch; i++)
for (int z = 0 ; z < dimz ; z++)
for (int y = 0; y < dimy; y++)
for (int x = 0; x < dimx; x++) {
if (*gt.data(i, 0, z, y, x)) {
counter += CHECK3();
//printf("i %d x %d y %d z %d in %d\n", i, x, y, z, res);
}
}
return counter;
}

void extend3(py::array_t<bool> &in, py::array_t<bool> &out) {
assert(in.ndim() == 5);
assert(out.ndim() == 5);
int n_batch = in.shape(0);
assert(in.shape(1) == 1);
int dimz = in.shape(2);
int dimy = in.shape(3);
int dimx = in.shape(4);

auto& ptr = in;
for (int i = 0; i < n_batch; i++)
for (int z = 1 ; z < dimz - 1; z++)
for (int y = 1; y < dimy - 1; y++)
for (int x = 1; x < dimx - 1; x++) {
*out.mutable_data(i, 0, z, y, x) = CHECK3();
}
}


PYBIND11_MODULE(ExtendBox, m) {
m.def("count_tp1", &count_tp1, "count tp func");
m.def("count_tp3", &count_tp3, "count tp func");
m.def("extend3", &extend3, "extend3 func");
}
@@ -0,0 +1,7 @@
all : ExtendBox.so

ExtendBox.so : ExtendBox.cpp
c++ -O3 -Wall -shared -std=c++14 -fPIC `python3 -m pybind11 --includes` $^ -o $@

clean:
rm ExtendBox.so
@@ -0,0 +1,10 @@

import json

def write(filename, data):
with open(filename, 'w') as outfile:
json.dump(data, outfile)

def read(filename):
with open(filename, 'r') as infile:
return json.load(infile)
@@ -0,0 +1,124 @@
import sys
import os
import struct
import glob
import numpy as np
import torch

class Vox:
def __init__(self, dims=[0, 0, 0], res=0.0, grid2world=None, sdf=None, pdf=None, noc=None, bbox=None):
self.dims = dims
self.res = res
self.grid2world = grid2world
self.sdf = sdf
self.pdf = pdf
self.noc = noc
self.bbox = bbox

def make_torch(self):
self.sdf = torch.as_tensor(self.sdf)
self.grid2world = torch.as_tensor(self.grid2world)
if self.pdf is not None:
self.pdf = torch.as_tensor(self.pdf)
if self.noc is not None:
self.noc = torch.as_tensor(self.noc)

def make_numpy(self):
self.sdf = np.asarray(self.sdf)
self.grid2world = np.asarray(self.grid2world)
if self.pdf is not None:
self.pdf = np.asarray(self.pdf)
if self.noc is not None:
self.noc = np.asarray(self.noc)

def load_vox_header(filename):
assert os.path.isfile(filename), "file not found: %s" % filename
if filename.endswith(".df"):
f_or_c = "C"
elif filename.endswith(".sdf"):
f_or_c = "C"
else:
f_or_c = "F"

fin = open(filename, 'rb')

s = Vox()
s.dims[0] = struct.unpack('I', fin.read(4))[0]
s.dims[1] = struct.unpack('I', fin.read(4))[0]
s.dims[2] = struct.unpack('I', fin.read(4))[0]
s.res = struct.unpack('f', fin.read(4))[0]
n_elems = s.dims[0]*s.dims[1]*s.dims[2]

s.grid2world = struct.unpack('f'*16, fin.read(16*4))
s.grid2world = np.asarray(s.grid2world, dtype=np.float32).reshape([4, 4], order=f_or_c)

return s


def load_vox(filename):
assert os.path.isfile(filename), "file not found: %s" % filename
if filename.endswith(".df"):
f_or_c = "C"
elif filename.endswith(".sdf"):
f_or_c = "C"
else:
f_or_c = "F"

fin = open(filename, 'rb')

s = Vox()
s.dims[0] = struct.unpack('I', fin.read(4))[0]
s.dims[1] = struct.unpack('I', fin.read(4))[0]
s.dims[2] = struct.unpack('I', fin.read(4))[0]
s.res = struct.unpack('f', fin.read(4))[0]
n_elems = s.dims[0]*s.dims[1]*s.dims[2]

s.grid2world = struct.unpack('f'*16, fin.read(16*4))
s.grid2world = np.asarray(s.grid2world, dtype=np.float32).reshape([4, 4], order=f_or_c)

# -> sdf 1-channel
s.sdf = struct.unpack('f'*n_elems, fin.read(n_elems*4))
s.sdf = np.asarray(s.sdf, dtype=np.float32).reshape([1, s.dims[2], s.dims[1], s.dims[0]])
# <-

# -> pdf 1-channel
pdf_bytes = fin.read(n_elems*4)
if pdf_bytes:
s.pdf = struct.unpack('f'*n_elems, pdf_bytes)
s.pdf = np.asarray(s.pdf, dtype=np.float32).reshape([1, s.dims[2], s.dims[1], s.dims[0]])
# <-

# -> noc 3-channels
noc_bytes = fin.read(3*n_elems*4)
if noc_bytes:
s.noc = struct.unpack('f'*n_elems*3, noc_bytes)
s.noc = np.asarray(s.noc, dtype=np.float32).reshape([3, s.dims[2], s.dims[1], s.dims[0]])
# <-

# -> bbox 1-channel
bbox_bytes = fin.read(n_elems*4)
if bbox_bytes:
s.bbox = struct.unpack('f'*n_elems, bbox_bytes)
s.bbox = np.asarray(s.bbox, dtype=np.float32).reshape([1, s.dims[2], s.dims[1], s.dims[0]])
# <-
fin.close()

return s

def write_vox(filename, s):
fout = open(filename, 'wb')
fout.write(struct.pack('I', s.dims[0]))
fout.write(struct.pack('I', s.dims[1]))
fout.write(struct.pack('I', s.dims[2]))
fout.write(struct.pack('f', s.res))
n_elems = np.prod(s.dims)
fout.write(struct.pack('f'*16, *s.grid2world.flatten('F')))
fout.write(struct.pack('f'*n_elems, *s.sdf.flatten('C')))
if s.pdf is not None:
fout.write(struct.pack('f'*n_elems, *s.pdf.flatten('C')))
if s.noc is not None:
fout.write(struct.pack('f'*n_elems*3, *s.noc.flatten('C')))
if s.bbox is not None:
fout.write(struct.pack('f'*n_elems, *s.bbox.flatten('C')))
fout.close()

@@ -0,0 +1,105 @@
import numpy as np
import sys
sys.path.append("../base/ExtendBox")
import ExtendBox

class Statistics:
def __init__(self):
self.tp = {}
self.fp = {}
self.tn = {}
self.fn = {}
self.t0 = [round(t, 2) for t in np.linspace(0, 1, 21)]

for t in self.t0:
self.tp[t] = 0
self.fp[t] = 0
self.fn[t] = 0
assert(0.5 in self.tp)

def pr_curve(self):
pr = []
for t in self.t0:
pr.append([t, self.precision(t), self.recall(t)])
return pr

def update_column(self, a0, gt0):
for t in self.t0:
a = np.greater_equal(a0, t)
gt = np.greater_equal(gt0, t)
tp = ExtendBox.count_tp1(a, gt)
fp = np.maximum(np.sum(a) - tp, 0)
fn = np.maximum(np.sum(gt) - tp, 0)

self.tp[t] += tp
self.fp[t] += fp
self.fn[t] += fn

def update(self, a0, gt0):
assert(np.size(a0) == np.size(gt0))
for t in self.t0:
a = np.greater_equal(a0, t)
gt = np.greater_equal(gt0, t)
#gte = np.zeros(np.shape(gt), dtype="bool")
tp = ExtendBox.count_tp3(a, gt)
fp = np.maximum(np.sum(a) - tp, 0)
fn = np.maximum(np.sum(gt) - tp, 0)

self.tp[t] += tp
self.fp[t] += fp
self.fn[t] += fn

def extend_column(self, gt):
gte = np.zeros(np.shape(gt))
dimz = np.size(gt, 1)
for z in range(1, dimz - 2):
gte[:, z] = gt[:, z-1:z+2].any()
return gte

def extend_box(self, gt):
gte = np.zeros(np.shape(gt))
dimz = np.size(gt, 2)
dimy = np.size(gt, 3)
dimx = np.size(gt, 4)
for z in range(1, dimz - 2):
for y in range(1, dimy - 2):
for x in range(1, dimx - 2):
gte[:, 0, z, y, x] = gt[:, 0, z-1:z+2, y-1:y+2, x-1:x+2].any()
return gte


def update_default(self, a0, gt0):
for t in self.t0:
a = np.greater_equal(a0, t)
gt = np.greater_equal(gt0, t)

self.tp[t] += np.sum(np.logical_and(a, gt))
self.fp[t] += np.sum(np.logical_and(a, np.logical_not(gt)))
self.fn[t] += np.sum(np.logical_and(np.logical_not(a), gt))

def mAP(self):
p = [self.precision(t) for t in self.t0]
r = [self.recall(t) for t in self.t0]
integral = np.abs(np.trapz(p[1:-1], x=r[1:-1]))
return integral

def precision(self, t=0.5):
p = self.tp[t] + self.fp[t]
if p != 0:
return self.tp[t]/p
else:
return 0

def recall(self, t=0.5):
r = self.tp[t] + self.fn[t]
if r != 0:
return self.tp[t]/r
else:
return 0


def calc_diff(a, b):
d0 = np.sum(a.sdf - b.sdf)
d1 = np.sum(a.pdf - b.pdf)
print(d0, d1)

@@ -0,0 +1,22 @@
import numpy as np

def gaussian3d(l, sigma=1.0):
"""
creates gaussian kernel with side length l and a sigma of sigma
"""

ax = np.arange(-l//2 + 1.0, l//2 + 1.0)
xx, yy, zz = np.meshgrid(ax, ax, ax)

kernel = np.exp(-(xx**2 + yy**2 + zz**2)/(2.0*sigma**2))

return np.asarray(kernel, dtype=np.float32)

def gaussian1d(l, sigma=1.0):

ax = np.arange(-l//2 + 1.0, l//2 + 1.0)
xx = np.meshgrid(ax)

kernel = np.exp(-np.power(xx, 2)/(2.0*sigma**2))

return np.asarray(kernel, dtype=np.float32)

0 comments on commit 11a57a6

Please sign in to comment.
You can’t perform that action at this time.