Upload 5 files

image_stitching/__init__.py

@@ -0,0 +1,8 @@
+from .stitcher import ImageStitcher
+from .helpers import display
+from .helpers import load_frames
+
+__doc__ = '''
+    image_stitching is based around the ImageStitcher class which handles all 
+    feature extraction and image warping
+'''

image_stitching/combine.py

@@ -0,0 +1,65 @@
+import logging
+
+import cv2
+import numpy
+
+__doc__ = '''helper functions for combining images, only to be used in the stitcher class'''
+
+
+def compute_matches(features0, features1, matcher, knn=5, lowe=0.7):
+    '''
+    this applies lowe-ratio feature matching between feature0 an dfeature 1 using flann
+    '''
+    keypoints0, descriptors0 = features0
+    keypoints1, descriptors1 = features1
+
+    logging.debug('finding correspondence')
+
+    matches = matcher.knnMatch(descriptors0, descriptors1, k=knn)
+
+    logging.debug("filtering matches with lowe test")
+
+    positive = []
+    for match0, match1 in matches:
+        if match0.distance < lowe * match1.distance:
+            positive.append(match0)
+
+    src_pts = numpy.array([keypoints0[good_match.queryIdx].pt for good_match in positive],
+                          dtype=numpy.float32)
+    src_pts = src_pts.reshape((-1, 1, 2))
+    dst_pts = numpy.array([keypoints1[good_match.trainIdx].pt for good_match in positive],
+                          dtype=numpy.float32)
+    dst_pts = dst_pts.reshape((-1, 1, 2))
+
+    return src_pts, dst_pts, len(positive)
+
+
+def combine_images(img0, img1, h_matrix):
+    '''
+    this takes two images and the homography matrix from 0 to 1 and combines the images together!
+    the logic is convoluted here and needs to be simplified!
+    '''
+    logging.debug('combining images... ')
+
+    points0 = numpy.array(
+        [[0, 0], [0, img0.shape[0]], [img0.shape[1], img0.shape[0]], [img0.shape[1], 0]],
+        dtype=numpy.float32)
+    points0 = points0.reshape((-1, 1, 2))
+    points1 = numpy.array(
+        [[0, 0], [0, img1.shape[0]], [img1.shape[1], img1.shape[0]], [img1.shape[1], 0]],
+        dtype=numpy.float32)
+    points1 = points1.reshape((-1, 1, 2))
+
+    points2 = cv2.perspectiveTransform(points1, h_matrix)
+    points = numpy.concatenate((points0, points2), axis=0)
+
+    [x_min, y_min] = (points.min(axis=0).ravel() - 0.5).astype(numpy.int32)
+    [x_max, y_max] = (points.max(axis=0).ravel() + 0.5).astype(numpy.int32)
+
+    h_translation = numpy.array([[1, 0, -x_min], [0, 1, -y_min], [0, 0, 1]])
+
+    logging.debug('warping previous image...')
+    output_img = cv2.warpPerspective(img1, h_translation.dot(h_matrix),
+                                     (x_max - x_min, y_max - y_min))
+    output_img[-y_min:img0.shape[0] - y_min, -x_min:img0.shape[1] - x_min] = img0
+    return output_img

image_stitching/helpers.py

@@ -0,0 +1,53 @@
+import logging
+import pathlib
+
+from typing import List
+from typing import Generator
+
+import cv2
+import numpy
+
+__doc__ = '''helper functions for loading frames and displaying them'''
+
+
+def display(title, img, max_size=500000):
+    '''
+    resizes the image before it displays it,
+    this stops large stitches from going over the screen!
+    '''
+    assert isinstance(img, numpy.ndarray), 'img must be a numpy array'
+    assert isinstance(title, str), 'title must be a string'
+    scale = numpy.sqrt(min(1.0, float(max_size) / (img.shape[0] * img.shape[1])))
+    shape = (int(scale * img.shape[1]), int(scale * img.shape[0]))
+    img = cv2.resize(img, shape)
+    cv2.imshow(title, img)
+
+
+def read_video(video_path: pathlib.Path):
+    '''read video is a generator class yielding frames'''
+    cap = cv2.VideoCapture(str(video_path))
+
+    while True:
+        ret, frame = cap.read()
+        if not ret or frame is None:
+            break
+
+        yield frame
+
+
+def load_frames(paths: List[str]) -> Generator[numpy.ndarray, None, None]:
+    '''
+    load_frames takes in a list of paths to image,
+    video files, or directories and yields them
+    '''
+    for path in paths:
+        path = pathlib.Path(path)
+
+        if path.is_dir():
+            yield from load_frames(path.rglob('*'))
+        elif path.suffix.lower() in ['.jpg', '.jpeg', '.png']:
+            yield cv2.imread(str(path))
+        elif path.suffix.lower() in ['.avi', '.mp4', '.mov']:
+            yield from read_video(path)
+        else:
+            logging.warning(f'skipping {path.name}...')

image_stitching/stitcher.py

@@ -0,0 +1,66 @@
+import logging
+
+import cv2
+import numpy
+
+from .combine import combine_images
+from .combine import compute_matches
+
+__doc__ = '''ImageStitcher class for combining all images together'''
+
+
+class ImageStitcher:
+    __doc__ = '''ImageStitcher class for combining all images together'''
+
+    def __init__(self, min_num: int = 10, lowe: float = 0.7, knn_clusters: int = 2):
+        '''constructor that initialises the SIFT class and Flann matcher'''
+        self.min_num = min_num
+        self.lowe = lowe
+        self.knn_clusters = knn_clusters
+
+        self.flann = cv2.FlannBasedMatcher({'algorithm': 0, 'trees': 5}, {'checks': 50})
+        self.sift = cv2.SIFT_create()
+
+        self.result_image = None
+        self.result_image_gray = None
+
+    def add_image(self, image: numpy.ndarray):
+        '''
+        this adds a new image to the stitched image by
+        running feature extraction and matching them
+        '''
+        assert image.ndim == 3, 'must be an image!'
+        assert image.shape[-1] == 3, 'must be BGR!'
+        assert image.dtype == numpy.uint8, 'must be a uint8'
+
+        image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+
+        if self.result_image is None:
+            self.result_image = image
+            self.result_image_gray = image_gray
+            return
+
+        # todo(will.brennan) - stop computing features on the results image each time!
+        result_features = self.sift.detectAndCompute(self.result_image_gray, None)
+        image_features = self.sift.detectAndCompute(image_gray, None)
+
+        matches_src, matches_dst, n_matches = compute_matches(result_features,
+                                                              image_features,
+                                                              matcher=self.flann,
+                                                              knn=self.knn_clusters,
+                                                              lowe=self.lowe)
+
+        if n_matches < self.min_num:
+            logging.warning('too few correspondences to add image to stitched image')
+            return
+
+        logging.debug('computing homography between accumulated and new images')
+        homography, _ = cv2.findHomography(matches_src, matches_dst, cv2.RANSAC, 5.0)
+
+        logging.debug('stitching images together')
+        self.result_image = combine_images(image, self.result_image, homography)
+        self.result_image_gray = cv2.cvtColor(self.result_image, cv2.COLOR_RGB2GRAY)
+
+    def image(self):
+        '''class for fetching the stitched image'''
+        return self.result_image

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+opencv-python>=4.5
+opencv-contrib-python>=4.5
+yapf
+pylint