First commit

Author: saisrivatsan

aligner.py

@@ -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)
+

conf.yaml

@@ -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

demo.ipynb

@@ -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)

detector.py

@@ -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)
+

mask.py

@@ -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
+