Skip to content

Commit

Permalink
First commit
Browse files Browse the repository at this point in the history
  • Loading branch information
saisrivatsan committed Apr 5, 2018
1 parent a367979 commit e5fb6ac
Show file tree
Hide file tree
Showing 9 changed files with 1,517 additions and 0 deletions.
153 changes: 153 additions & 0 deletions aligner.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
#from rh_renderer import models
#from rh_aligner.common import ransac
import sys
import os
import glob
import yaml
import cv2
import h5py
import numpy as np
from rh_logger.api import logger
import logging
import rh_logger
import time
from detector import FeaturesDetector
from matcher import FeaturesMatcher
import multiprocessing as mp

class StackAligner(object):

def __init__(self, conf, processes_num=1):
self._conf = conf

# Initialize the detector, amtcher and optimizer objects
detector_params = conf.get('detector_params', {})
matcher_params = conf.get('matcher_params', {})
self._detector = FeaturesDetector(conf['detector_type'], **detector_params)
self._matcher = FeaturesMatcher(self._detector, **matcher_params)

self._processes_num = processes_num



@staticmethod
def read_imgs(folder):
img_fnames = sorted(glob.glob(os.path.join(folder, '*')))
print("Loading {} images from {}.".format(len(img_fnames), folder))
imgs = [cv2.imread(img_fname, 0) for img_fname in img_fnames]
return img_fnames, imgs


@staticmethod
def load_conf_from_file(conf_fname):
'''
Loads a given configuration file from a yaml file
'''
print("Using config file: {}.".format(conf_fname))
with open(conf_fname, 'r') as stream:
conf = yaml.load(stream)
return conf


@staticmethod
def _compute_l2_distance(pts1, pts2):
delta = pts1 - pts2
s = np.sum(delta**2, axis=1)
return np.sqrt(s)

@staticmethod
def _compute_features(detector, img, i):
result = detector.detect(img)
logger.report_event("Img {}, found {} features.".format(i, len(result[0])), log_level=logging.INFO)
return result

@staticmethod
def _match_features(features_result1, features_result2, i, j):
transform_model, filtered_matches = self._matcher.match_and_filter(features_result1[0], features_result1[1], features_result2[0], features_result2[1])
assert(transform_model is not None)
transform_matrix = transform_model.get_matrix()
logger.report_event("Imgs {} -> {}, found the following transformations\n{}\nAnd the average displacement: {} px".format(i, j, transform_matrix, np.mean(StackAligner._compute_l2_distance(transform_model.apply(filtered_matches[1]), filtered_matches[0]))), log_level=logging.INFO)
return transform_matrix


def align_imgs(self, imgs):
'''
Receives a stack of images to align and aligns that stack using the first image as an anchor
'''

#pool = mp.Pool(processes=processes_num)

# Compute features
logger.start_process('align_imgs', 'aligner.py', [len(imgs), self._conf])
logger.report_event("Computing features...", log_level=logging.INFO)
st_time = time.time()
all_features = []
pool_results = []
for i, img in enumerate(imgs):
#res = pool.apply_async(StackAligner._compute_features, (self._detector, img, i))
#pool_results.append(res)
all_features.append(self._detector.detect(img))
logger.report_event("Img {}, found {} features.".format(i, len(all_features[-1][0])), log_level=logging.INFO)
for res in pool_results:
all_features.append(res.get())
logger.report_event("Features computation took {} seconds.".format(time.time() - st_time), log_level=logging.INFO)

# match features of adjacent images
logger.report_event("Pair-wise feature mathcing...", log_level=logging.INFO)
st_time = time.time()
pairwise_transforms = []
for i in range(len(imgs) - 1):
transform_model, filtered_matches = self._matcher.match_and_filter(all_features[i + 1][0], all_features[i+1][1], all_features[i][0], all_features[i][1])
assert(transform_model is not None)
transform_matrix = transform_model.get_matrix()
pairwise_transforms.append(transform_matrix)
logger.report_event("Imgs {} -> {}, found the following transformations\n{}\nAnd the average displacement: {} px".format(i, i+1, transform_matrix, np.mean(StackAligner._compute_l2_distance(transform_model.apply(filtered_matches[1]), filtered_matches[0]))), log_level=logging.INFO)
logger.report_event("Feature matching took {} seconds.".format(time.time() - st_time), log_level=logging.INFO)

# Compute the per-image transformation (all images will be aligned to the first section)
logger.report_event("Computing transformations...", log_level=logging.INFO)
st_time = time.time()
transforms = []
cur_transform = np.eye(3)
transforms.append(cur_transform)

for pair_transform in pairwise_transforms:
cur_transform = np.dot(cur_transform, pair_transform)
transforms.append(cur_transform)
logger.report_event("Transformations computation took {} seconds.".format(time.time() - st_time), log_level=logging.INFO)

assert(len(imgs) == len(transforms))

#pool.close()
#pool.join()

logger.end_process('align_imgs ending', rh_logger.ExitCode(0))
return transforms




@staticmethod
def align_img_files(imgs_dir, conf, processes_num):
# Read the files
img_names, imgs = StackAligner.read_imgs(imgs_dir)

aligner = StackAligner(conf, processes_num)
return img_names, imgs, aligner.align_imgs(imgs)


def get_transforms(imgs_dir='gt-4x6x6_image.h5', conf_fname='conf.yaml', processes_num = 8):

conf = StackAligner.load_conf_from_file(conf_fname)
if imgs_dir.endswith('.h5'):
with h5py.File(imgs_dir, "r") as fd:
imgs = fd[fd.keys()[0]][:].astype(np.uint8)

img_names = None
aligner = StackAligner(conf, processes_num)
transforms = aligner.align_imgs(imgs)

else:
img_names, imgs, transforms = StackAligner.align_img_files(imgs_dir, conf, processes_num)
return img_names, imgs, np.asarray(transforms)

31 changes: 31 additions & 0 deletions conf.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
#detector_type: SIFT
#detector_params:
# sigma: 1.8
detector_type: ORB
detector_params:
nfeatures: 10000
#detector_type: SURF
#detector_type: BRISK
#detector_type: AKAZE

#descriptor:
# type: SIFT

matcher_params:
ROD_cutoff: 0.9
model_index: 3 # 0 - Translation, 1 - Rigid, 3 - Affine
#num_filtered_percent: 0.1
#filter_rate_cutoff: 0.1
min_features_num: 10
iterations: 5000
max_epsilon: 20
min_inlier_ratio: 0.01
min_num_inlier: 0.01
max_trust: 3
det_delta: 0.95
max_stretch: 0.95
#use_regularizer: True
#regularizer_model_index: 1
#regularizer_lambda: 0.1

out_path: ./analysis/AKAZE
301 changes: 301 additions & 0 deletions demo.ipynb

Large diffs are not rendered by default.

40 changes: 40 additions & 0 deletions detector.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
import cv2
from enum import Enum
import numpy as np

class FeaturesDetector(object):

class Type(Enum):
SIFT = 1
ORB = 2
SURF = 3
BRISK = 4
AKAZE = 5

def __init__(self, detector_type_name, **kwargs):
detector_type = FeaturesDetector.Type[detector_type_name]
if detector_type == FeaturesDetector.Type.SIFT:
self._detector = cv2.xfeatures2d.SIFT_create(**kwargs)
self._matcher = cv2.BFMatcher(cv2.NORM_L2)
elif detector_type == FeaturesDetector.Type.ORB:
cv2.ocl.setUseOpenCL(False) # Avoiding a bug in OpenCV 3.1
self._detector = cv2.ORB_create(**kwargs)
self._matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
elif detector_type == FeaturesDetector.Type.SURF:
self._detector = cv2.xfeatures2d.SURF_create(**kwargs)
self._matcher = cv2.BFMatcher(cv2.NORM_L2)
elif detector_type == FeaturesDetector.Type.BRISK:
self._detector = cv2.BRISK_create(**kwargs)
self._matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
elif detector_type == FeaturesDetector.Type.AKAZE:
self._detector = cv2.AKAZE_create(**kwargs)
self._matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
else:
raise("Unknown feature detector algorithm given")


def detect(self, img):
return self._detector.detectAndCompute(img, None)

def match(self, features_descs1, features_descs2):
return self._matcher.knnMatch(features_descs1, features_descs2, k=2)
19 changes: 19 additions & 0 deletions mask.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
import os
import cv2
import h5py
import numpy as np

def compute_mask(imgs, n_imgs, nTransforms):

#with h5py.File("output/transforms.h5", "r") as fd:
# nTransforms = fd["nTransforms"][:]
#print(nTransforms.shape)
masks = []
for img, n_img, transform in zip(imgs, n_imgs, nTransforms):
nH, nW = n_img.shape
tmp = np.ones_like(img)
img_transformed = cv2.warpAffine(tmp, transform[:2,:], (nW, nH), flags=cv2.INTER_NEAREST)
masks.append(img_transformed)

return np.asarray(masks)

130 changes: 130 additions & 0 deletions matcher.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
import numpy as np
import ransac
#from scipy.spatial import KDTree
from collections import defaultdict
import models

#NEIGHBORHOOD_RADIUS = 70
GRID_SIZE = 50

class FeaturesMatcher(object):

def __init__(self, detector, **kwargs):
self._detector = detector

self._params = {}
# get default values if no value is present in kwargs
#self._params["num_filtered_percent"] = kwargs.get("num_filtered_percent", 0.25)
#self._params["filter_rate_cutoff"] = kwargs.get("filter_rate_cutoff", 0.25)
self._params["ROD_cutoff"] = kwargs.get("ROD_cutoff", 0.92)
self._params["min_features_num"] = kwargs.get("min_features_num", 40)

# Parameters for the RANSAC
self._params["model_index"] = kwargs.get("model_index", 3)
self._params["iterations"] = kwargs.get("iterations", 5000)
self._params["max_epsilon"] = kwargs.get("max_epsilon", 30.0)
self._params["min_inlier_ratio"] = kwargs.get("min_inlier_ratio", 0.01)
self._params["min_num_inlier"] = kwargs.get("min_num_inliers", 7)
self._params["max_trust"] = kwargs.get("max_trust", 3)
self._params["det_delta"] = kwargs.get("det_delta", 0.9)
self._params["max_stretch"] = kwargs.get("max_stretch", 0.25)

self._params["use_regularizer"] = True if "use_regularizer" in kwargs.keys() else False
self._params["regularizer_lambda"] = kwargs.get("regularizer_lambda", 0.1)
self._params["regularizer_model_index"] = kwargs.get("regularizer_model_index", 1)


def match(self, features_kps1, features_descs1, features_kps2, features_descs2):
features_kps2 = np.asarray(features_kps2)

# because the sections were already pre-aligned, we only match a feature in section1 to its neighborhood in section2 (according to a grid)
grid = defaultdict(set)

# build the grid of the feature locations in the second section
for i, kp2 in enumerate(features_kps2):
pt_grid = (np.array(kp2.pt) / GRID_SIZE).astype(np.int)
grid[tuple(pt_grid)].add(i)

match_points = [[], [], []]
for kp1, desc1 in zip(features_kps1, features_descs1):
# For each kp1 find the closest points in section2
pt_grid = (np.array(kp1.pt) / GRID_SIZE).astype(np.int)
close_kps2_idxs = set()
# search in a [-1, -1] -> [1, 1] delta windows (3*3)
for delta_y in range(-1, 2):
for delta_x in range(-1, 2):
delta = np.array([delta_x, delta_y], dtype=np.int)
delta_grid_loc = tuple(pt_grid + delta)
if delta_grid_loc in grid.keys():
close_kps2_idxs |= grid[delta_grid_loc]

close_kps2_indices = list(close_kps2_idxs)
close_descs2 = features_descs2[close_kps2_indices]
matches = self._detector.match(desc1.reshape(1, len(desc1)), close_descs2)
if len(matches[0]) == 2:
if matches[0][0].distance < self._params["ROD_cutoff"] * matches[0][1].distance:
match_points[0].append(kp1.pt)
match_points[1].append(features_kps2[close_kps2_indices][matches[0][0].trainIdx].pt)
match_points[2].append(matches[0][0].distance)




# # because the sections were already pre-aligned, we only match a feature in section1 to its neighborhood in section2
# features_kps2_pts = [kp.pt for kp in features_kps2]
# kps2_pts_tree = KDTree(features_kps2_pts)
#
# match_points = [[], [], []]
# for kp1, desc1 in zip(features_kps1, features_descs1):
# # For each kp1 find the closest points in section2
# close_kps2_indices = kps2_pts_tree.query_ball_point(kp1.pt, NEIGHBORHOOD_RADIUS)
# close_descs2 = features_descs2[close_kps2_indices]
# matches = self._detector.match(desc1.reshape(1, len(desc1)), close_descs2)
# if len(matches[0]) == 2:
# if matches[0][0].distance < self._params["ROD_cutoff"] * matches[0][1].distance:
# match_points[0].append(kp1.pt)
# match_points[1].append(features_kps2[close_kps2_indices][matches[0][0].trainIdx].pt)
# match_points[2].append(matches[0][0].distance)

match_points = (np.array(match_points[0]), np.array(match_points[1]), np.array(match_points[2]))


# matches = self._detector.match(features_descs1, features_descs2)
#
# good_matches = []
# for m, n in matches:
# #if (n.distance == 0 and m.distance == 0) or (m.distance / n.distance < actual_params["ROD_cutoff"]):
# if m.distance < self._params["ROD_cutoff"] * n.distance:
# good_matches.append(m)
#
# match_points = (
# np.array([features_kps1[m.queryIdx].pt for m in good_matches]),
# np.array([features_kps2[m.trainIdx].pt for m in good_matches]),
# np.array([m.distance for m in good_matches])
# )

return match_points

def match_and_filter(self, features_kps1, features_descs1, features_kps2, features_descs2):
match_points = self.match(features_kps1, features_descs1, features_kps2, features_descs2)

model, filtered_matches = ransac.filter_matches(match_points, match_points, self._params['model_index'],
self._params['iterations'], self._params['max_epsilon'], self._params['min_inlier_ratio'],
self._params['min_num_inlier'], self._params['max_trust'], self._params['det_delta'], self._params['max_stretch'])

if model is None:
return None, None

if self._params["use_regularizer"]:
regularizer_model, _ = ransac.filter_matches(match_points, match_points, self._params['regularizer_model_index'],
self._params['iterations'], self._params['max_epsilon'], self._params['min_inlier_ratio'],
self._params['min_num_inlier'], self._params['max_trust'], self._params['det_delta'], self._params['max_stretch'])

if regularizer_model is None:
return None, None

result = model.get_matrix() * (1 - self._params["regularizer_lambda"]) + regularizer_model.get_matrix() * self._params["regularizer_lambda"]
model = models.AffineModel(result)

return model, filtered_matches

Loading

0 comments on commit e5fb6ac

Please sign in to comment.