Skip to content

Instantly share code, notes, and snippets.

@spedy
Last active January 13, 2019 09:27
Show Gist options
  • Save spedy/747a7f6d74df718713cc to your computer and use it in GitHub Desktop.
Save spedy/747a7f6d74df718713cc to your computer and use it in GitHub Desktop.
from freenect import sync_get_depth as get_depth, sync_get_video as get_video
import cv2
import time
import numpy as np
import pickle
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
def save_file(frame, idx):
cv2.imwrite("data/rgb{}.jpeg".format(idx), frame["rgb"])
cv2.imwrite("data/depth{}.jpeg".format(idx), frame["depth"])
objp = np.zeros((6*7,3), np.float32)
objp[:,:2] = np.mgrid[0:7,0:6].T.reshape(-1,2)
objpoints = [] # 3d point in real world space
imgpoints = [] # 2d points in image plane.
calibrate = True
num_images = 30
_corners2 = []
if calibrate:
frames = []
idx = 0
while len(frames) < num_images:
print "Capturing frame {}".format(idx)
frame = {"depth": get_depth()[0], "rgb": get_video()[0]}
save_file(frame, idx)
frames.append(frames)
idx += 1
time.sleep(1)
else:
imgpoints = []
objpoints = []
for i in range(num_images):
image_idx = i
img = cv2.imread("data/rgb{}.jpeg".format(image_idx))
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
ret, corners = cv2.findChessboardCorners(gray, (7,6),None)
if ret:
objpoints.append(objp)
corners2 = cv2.cornerSubPix(gray,corners, (11,11), (-1,-1), criteria)
_corners2.append(corners2)
imgpoints.append(corners2)
img = cv2.drawChessboardCorners(img, (7,6), corners2, ret)
"""
cv2.imshow('img',img)
cv2.waitKey(500)
"""
if imgpoints:
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
calibrate_camera_data = {
"mtx": mtx,
"dist": dist,
"rvecs": rvecs,
"tvecs": tvecs,
"corners2": _corners2
}
pickle.dump(calibrate_camera_data, open("calibrate_camera_data.p", "wb"))
import cv2
import numpy as np
import pickle
focalLength = 525.0
centerX = 319.5
centerY = 239.5
scalingFactor = 5000.0
def write_ply(points):
file = open("out.ply", "w")
file.write('''ply
format ascii 1.0
element vertex %d
property float x
property float y
property float z
end_header
%s
'''%(len(points),"".join(points)))
file.close()
def drawMatches(img1, kp1, img2, kp2, matches):
# Create a new output image that concatenates the two images together
# (a.k.a) a montage
rows1 = img1.shape[0]
cols1 = img1.shape[1]
rows2 = img2.shape[0]
cols2 = img2.shape[1]
out = np.zeros((max([rows1,rows2]),cols1+cols2,3), dtype='uint8')
# Place the first image to the left
out[:rows1,:cols1] = np.dstack([img1, img1, img1])
# Place the next image to the right of it
out[:rows2,cols1:] = np.dstack([img2, img2, img2])
# For each pair of points we have between both images
# draw circles, then connect a line between them
for mat in matches:
# Get the matching keypoints for each of the images
img1_idx = mat.queryIdx
img2_idx = mat.trainIdx
# x - columns
# y - rows
(x1,y1) = kp1[img1_idx].pt
(x2,y2) = kp2[img2_idx].pt
# Draw a small circle at both co-ordinates
# radius 4
# colour blue
# thickness = 1
cv2.circle(out, (int(x1),int(y1)), 4, (255, 0, 0), 1)
cv2.circle(out, (int(x2)+cols1,int(y2)), 4, (255, 0, 0), 1)
# Draw a line in between the two points
# thickness = 1
# colour blue
cv2.line(out, (int(x1),int(y1)), (int(x2)+cols1,int(y2)), (255, 0, 0), 1)
# Show the image
cv2.imshow('Matched Features', out)
cv2.waitKey(0)
cv2.destroyWindow('Matched Features')
# Also return the image if you'd like a copy
return out
def get_rgbd_data_samples(num_data_samples, samples="rabbit2"):
rgb_path = "data/rgb{}.jpeg"
depth_path = "data/depth{}.jpeg"
if samples == "rabbit":
rgb_path = "../data/color_{}.png"
depth_path = "../data/depth_{}.png"
rgbd = []
for idx in range(num_data_samples):
rgbd.append({
"rgb": cv2.imread(rgb_path.format(idx), 0),
"depth": cv2.imread(depth_path.format(idx), 0)
})
return rgbd
def compute_sift(rgb1, rgb2):
sift = cv2.xfeatures2d.SIFT_create()
(kp1, des1) = sift.detectAndCompute(rgb1, None)
(kp2, des2) = sift.detectAndCompute(rgb2, None)
bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
matches = bf.match(des1,des2)
return {
"matches": sorted(matches, key=lambda val: val.distance),
"kp1": kp1,
"kp2": kp2
}
def generate_pointcloud(points3d):
_points = []
for points in points3d:
for point in points:
_points.append("%f %f %f 0\n"%(point[0], point[1], point[2]))
return _points
def transform3d(x, y, depth):
Z = depth / scalingFactor
X = (x - centerX) * Z / focalLength
Y = (y - centerY) * Z / focalLength
return (X, Y, Z)
def transform(sift_data, depth1, depth2):
objpoints, imgpoints = [], []
kp1 = sift_data["kp1"]
kp2 = sift_data["kp2"]
matches = sift_data["matches"]
for mat in matches:
(x, y) = kp1[mat.queryIdx].pt
d = depth1[y][x]
if d == 0:
continue
objpoints.append(transform3d(x,y,d))
imgpoints.append(kp2[mat.trainIdx].pt)
return {
"objpoints": objpoints,
"imgpoints": imgpoints
}
def transform_points(rvecs, tvecs, objpoints):
N = objpoints.shape[1]
objpoints = objpoints.reshape(N, 3)
R, _ = cv2.Rodrigues(rvecs)
T = tvecs
return (np.dot(R , objpoints.T.reshape(3, N)) + np.dot(T, np.ones([1, N]))).T
rgbd_data_samples = get_rgbd_data_samples(num_data_samples=19)
calibrate_camera_data = pickle.load(open("calibrate_camera_data.p", "rb"))
final_points_data = []
axis = np.float32([[3,0,0], [0,3,0], [0,0,-3]]).reshape(-1,3)
points3d = []
transformed_data = {"objpoints": [], "imgpoints": []}
for idx in range(len(rgbd_data_samples)-1):
# retrieve samples
rgb1, depth1 = rgbd_data_samples[idx].values()
rgb2, depth2 = rgbd_data_samples[idx].values()
# compute SIFT (scale invariant feature transform)
sift_data = compute_sift(rgb1, rgb2)
#drawMatches(rgb1, sift_data["kp1"], rgb2, sift_data["kp2"], sift_data["matches"])
# transform sift data and depth into obj and image points
_transformed_data = transform(sift_data, depth1, depth2)
transformed_data["objpoints"] += _transformed_data["objpoints"]
transformed_data["imgpoints"] += _transformed_data["imgpoints"]
transformed_data["objpoints"] = np.array([np.array(transformed_data["objpoints"], np.float32)])
transformed_data["imgpoints"] = np.array([np.array(transformed_data["imgpoints"], np.float32)])
_, rvecs, tvecs, inlier = cv2.solvePnPRansac(transformed_data["objpoints"], transformed_data["imgpoints"],
calibrate_camera_data["mtx"], calibrate_camera_data["dist"]
)
final_points_data.append(transform_points(rvecs, tvecs, transformed_data["objpoints"]))
write_ply(generate_pointcloud(final_points_data))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment