Merge branch 'fix/simplify-realign-image' into 'develop'

Simplified realigning the image

See merge request !41

tests/test_images.py

-import cv2
-import os
-import numpy as np
-import pytest
-
-from zesje.images import get_delta, get_corner_marker_sides, fix_corner_markers, add_tup, sub_tup
-from zesje.scans import find_corner_marker_keypoints
-
-
-@pytest.mark.parametrize(
-    'shape,corners,expected',
-    [((240, 200, 3), [(120, 50), (50, 200), (120, 200)], (50, 50)),
-     ((240, 200, 3), [(50, 50), (50, 200), (120, 200)], (120, 50)),
-     ((240, 200, 3), [(50, 50), (120, 50), (120, 200)], (50, 200)),
-     ((240, 200, 3), [(50, 50), (120, 50), (50, 200)], (120, 200))],
-    ids=["missing top left", "missing top right", "missing bottom left", "missing bottom right"])
-def test_three_straight_corners(shape, corners, expected):
-    corner_markers = fix_corner_markers(corners, shape)
-    assert expected in corner_markers
-
-
-def test_pdf(datadir):
-    # Max deviation of inferred corner marker and actual location
-    epsilon = 2
-
-    # Scan rotated image with 4 corner markers
-    image_filename1 = 'a4-rotated.png'
-    image_path = os.path.join(datadir, 'cornermarkers', image_filename1)
-    page_img = cv2.imread(image_path)
-
-    corners1 = find_corner_marker_keypoints(page_img)
-
-    # Scan the same image with 3 corner markers
-    image_filename2 = 'a4-rotated-3-markers.png'
-    image_path = os.path.join(datadir, 'cornermarkers', image_filename2)
-    page_img = cv2.imread(image_path)
-
-    corners2 = find_corner_marker_keypoints(page_img)
-
-    # Get marker that was removed
-    diff = [corner for corner in corners1 if corner not in corners2]
-    diff_marker = min(diff)
-
-    fixed_corners2 = fix_corner_markers(corners2, page_img.shape)
-    added_marker = [corner for corner in fixed_corners2 if corner not in corners1][0]
-
-    # Check if 'inferred' corner marker is not too far away
-    dist = np.linalg.norm(np.subtract(added_marker, diff_marker))
-
-    assert dist < epsilon
-
-
-@pytest.mark.parametrize(
-    'inputs,expected',
-    [
-        (((0, 1), (1, 1), (0, 0), None), (0, 1)),
-        (((0, 1), None, (0, 0), (0, 1)), (0, 1)),
-        (((1, 1), (2, 1), None, (2, 2)), (0, -1)),
-        ((None, (1, 1), (1, 2), (2, 2)), (-1, -1))
-    ],
-    ids=["missing bottom right", "missing top right", "missing bottom left", "missing top left"]
-)
-def test_get_delta(inputs, expected):
-    # unpack inputs so that the individual elements are paramaters
-    delta = get_delta(*inputs)
-
-    assert delta == expected
-
-
-def test_get_corner_marker_sides_all_four():
-    shape = (100, 100)
-    corner_markers = [(0, 0), (100, 0), (0, 100), (100, 100)]
-
-    assert tuple(corner_markers) == get_corner_marker_sides(corner_markers, shape)
-
-
-def test_get_corner_markers_three():
-    shape = (100, 100)
-    corner_markers = [(0, 0), (0, 100), (100, 0)]
-
-    top_left, top_right, bottom_left, bottom_right = get_corner_marker_sides(corner_markers, shape)
-
-    assert not bottom_right
-
-
-def test_add_tup():
-    tup1 = tup2 = (1, 1)
-
-    assert add_tup(tup1, tup2) == (2, 2)
-
-
-def test_sub_tup():
-    tup1 = tup2 = (1, 1)
-
-    assert sub_tup(tup1, tup2) == (0, 0)

tests/test_scans.py

@@ -313,7 +313,10 @@ def test_image_extraction(datadir, filename):
 
 @pytest.mark.parametrize('file_name, markers', [("a4-rotated.png", [(59, 59), (1181, 59), (59, 1695), (1181, 1695)]),
                                                 ("a4-3-markers.png", [(1181, 59), (59, 1695), (1181, 1695)]),
-                                                ("a4-rotated-3-markers.png", [(1181, 59), (59, 1695), (1181, 1695)])
+                                                ("a4-rotated-3-markers.png", [(1181, 59), (59, 1695), (1181, 1695)]),
+                                                ("a4-rotated-2-markers.png", [(1181, 59), (59, 1695)]),
+                                                ("a4-rotated-2-bottom-markers.png", [(59, 1695), (1181, 1695)]),
+                                                ("a4-shifted-1-marker.png", [(59, 1695)])
                                                 ])
 def test_realign_image(datadir, file_name, markers):
     dir_name = "cornermarkers"
@@ -323,6 +326,7 @@ def test_realign_image(datadir, file_name, markers):
     test_image = np.array(PIL.Image.open(test_file))
 
     result_image = scans.realign_image(test_image)
+
     result_corner_markers = scans.find_corner_marker_keypoints(result_image)
     assert result_corner_markers is not None
     for i in range(len(markers)):
@@ -337,11 +341,9 @@ def test_incomplete_reference_realign_image(datadir):
     test_file = os.path.join(datadir, dir_name, "a4-rotated.png")
     test_image = cv2.imread(test_file)
 
-    reference_markers = [(59, 59), (59, 1695), (1181, 1695)]
-
     correct_corner_markers = [(59, 59), (1181, 59), (59, 1695), (1181, 1695)]
 
-    result_image = scans.realign_image(test_image, reference_keypoints=reference_markers)
+    result_image = scans.realign_image(test_image)
     result_corner_markers = scans.find_corner_marker_keypoints(result_image)
 
     assert result_corner_markers is not None
@@ -349,3 +351,22 @@ def test_incomplete_reference_realign_image(datadir):
         diff = np.absolute(np.subtract(correct_corner_markers[i], result_corner_markers[i]))
         assert diff[0] <= epsilon
         assert diff[1] <= epsilon
+
+
+def test_shift_image(datadir):
+    dir_name = "cornermarkers"
+    epsilon = 1
+    test_file = os.path.join(datadir, dir_name, "a4-1-marker.png")
+    test_image = cv2.imread(test_file)
+
+    bottom_left = (59, 1695)
+    shift_x, shift_y = 20, -30
+    shift_keypoint = (bottom_left[0] + shift_x, bottom_left[1] + shift_y)
+
+    test_image = scans.shift_image(test_image, shift_x, shift_y)
+    # test_image = scans.shift_image(test_image, bottom_left, shift_keypoint)
+    keypoints = scans.find_corner_marker_keypoints(test_image)
+
+    diff = np.absolute(np.subtract(shift_keypoint, keypoints[0]))
+    assert diff[0] <= epsilon
+    assert diff[1] <= epsilon

zesje/images.py

@@ -3,43 +3,6 @@
 import numpy as np
 
 
-def add_tup(tup1, tup2):
-    """
-    Adds two tuples
-
-    Parameters
-    ----------
-    tup1 : tuple
-        Tuple 1
-    tup2 : tuple
-        Tuple 2
-
-    Returns
-    -------
-    tup : tuple
-        The tuple with the sum of the values in tup1 and tup2.
-    """
-    return tup1[0] + tup2[0], tup1[1] + tup2[1]
-
-
-def sub_tup(tup1, tup2):
-    """Subtracts two tuples
-
-    Parameters
-    ----------
-    tup1 : tuple
-        Tuple 1
-    tup2 : tuple
-        Tuple 2
-
-    Returns
-    -------
-    tup : tuple
-        The tuple with the difference between the values in tup1 and tup2.
-    """
-    return tup1[0] - tup2[0], tup1[1] - tup2[1]
-
-
 def guess_dpi(image_array):
     h, *_ = image_array.shape
     resolutions = np.array([1200, 600, 400, 300, 200, 150, 144, 120, 100, 75, 72, 60, 50, 40])
@@ -71,111 +34,3 @@ def get_box(image_array, box, padding=0.3):
     top, bottom = max(0, min(box[0], h)), max(1, min(box[1], h))
     left, right = max(0, min(box[2], w)), max(1, min(box[3], w))
     return image_array[top:bottom, left:right]
-
-
-def get_corner_marker_sides(corner_markers, shape):
-    """Divides a list of corner markers in the right sides:
-
-    Parameters
-    ----------
-    corner_markers : list of tuples
-        The list of corner marker points
-    shape: tuple
-        The shape of an image
-
-    Returns
-    -------
-    tuples : tuple
-        The corner markers divided into sides
-    """
-
-    def get_val(tup_list):
-        """
-        Returns a tuple if present in the list.
-
-        Parameters
-        ----------
-        tup_list : list of tuples
-            List with one tuple
-
-        Returns
-        -------
-        tup : tuple or None
-            Tuple in list or empty list
-        """
-        return tup_list[0] if tup_list else None
-
-    x_sep = shape[1] / 2
-    y_sep = shape[0] / 2
-
-    top_left = get_val([(x, y) for x, y in corner_markers if x < x_sep and y < y_sep])
-    top_right = get_val([(x, y) for x, y in corner_markers if x > x_sep and y < y_sep])
-    bottom_left = get_val([(x, y) for x, y in corner_markers if x < x_sep and y > y_sep])
-    bottom_right = get_val([(x, y) for x, y in corner_markers if x > x_sep and y > y_sep])
-
-    return top_left, top_right, bottom_left, bottom_right
-
-
-def get_delta(top_left, top_right, bottom_left, bottom_right):
-    """Returns the absolute difference between the left or right points
-
-    Parameters
-    top_left : tuple
-        Top left point
-    top_right : tuple
-        Top right point
-    bottom_left : tuple
-        Bottom left point
-    bottom_right : tuple
-        Bottom right point
-
-    Returns
-    -------
-    delta : tuple
-        The absolute difference as an (x, y) tuple
-    """
-
-    if not top_left or not bottom_left:
-        return sub_tup(top_right, bottom_right)
-
-    return sub_tup(top_left, bottom_left)
-
-
-def fix_corner_markers(corner_keypoints, shape):
-    """Corrects the list of corner markers if three corner markers are found.
-    This function raises if less than three corner markers are found.
-
-    Parameters
-    ----------
-    corner_keypoints : list of tuples
-        List of corner marker locations as tuples
-    shape : (float, float, int)
-        Shape of the image in (x, y, dim)
-
-    Returns
-    -------
-    fixed_corners : (float, float)
-        A list of four corner markers.
-    """
-    if len(corner_keypoints) == 4:
-        return corner_keypoints
-
-    if len(corner_keypoints) < 3:
-        raise RuntimeError("Fewer than 3 corner markers found while trying to fix corners")
-
-    top_left, top_right, bottom_left, bottom_right = get_corner_marker_sides(corner_keypoints, shape)
-    delta = get_delta(top_left, top_right, bottom_left, bottom_right)
-
-    if not top_left:
-        top_left = add_tup(bottom_left, delta)
-
-    if not top_right:
-        top_right = add_tup(bottom_right, delta)
-
-    if not bottom_left:
-        bottom_left = sub_tup(top_left, delta)
-
-    if not bottom_right:
-        bottom_right = sub_tup(top_right, delta)
-
-    return [top_left, top_right, bottom_left, bottom_right]

zesje/scans.py

@@ -9,6 +9,8 @@ import signal
 
 import cv2
 import numpy as np
+from scipy import spatial
+
 from pikepdf import Pdf, PdfImage
 from PIL import Image
 from wand.image import Image as WandImage
@@ -19,7 +21,6 @@ from .datamatrix import decode_raw_datamatrix
 from .images import guess_dpi, get_box
 from .factory import make_celery
 from .pregrader import grade_mcq
-from .images import fix_corner_markers
 
 from .pdf_generation import MARKER_FORMAT, PAGE_FORMATS
 
@@ -316,11 +317,10 @@ def process_page(image_data, exam_config, output_dir=None, strict=False):
     corner_keypoints = find_corner_marker_keypoints(image_array)
     try:
         check_corner_keypoints(image_array, corner_keypoints)
+        image_array = realign_image(image_array, corner_keypoints)
     except RuntimeError as e:
         if strict:
             return False, str(e)
-    else:
-        image_array = realign_image(image_array, corner_keypoints)
 
     try:
         barcode, upside_down = decode_barcode(image_array, exam_config)
@@ -684,8 +684,7 @@ def check_corner_keypoints(image_array, keypoints):
         raise RuntimeError("Found multiple corner markers in the same corner")
 
 
-def realign_image(image_array, keypoints=None,
-                  reference_keypoints=None, page_format="A4"):
+def realign_image(image_array, keypoints=None, page_format="A4"):
     """
     Transform the image so that the keypoints match the reference.
 
@@ -696,9 +695,6 @@ def realign_image(image_array, keypoints=None,
     keypoints : List[(int,int)]
         tuples of coordinates of the found keypoints, (x,y), in pixels. Can be a set of 3 or 4 tuples.
         if none are provided, they are found by using find_corner_marker_keypoints on the input image.
-    reference_keypoints: List[(int,int)]
-        Similar to keypoints, only these belong to the keypoints found on the original scan.
-        If none are provided, ideal ones are calculated based on the dpi of the input image.
     returns
     -------
     return_array: numpy.array
@@ -707,27 +703,34 @@ def realign_image(image_array, keypoints=None,
 
     if not keypoints:
         keypoints = find_corner_marker_keypoints(image_array)
-        check_corner_keypoints(image_array, keypoints)
 
-    if len(keypoints) != 4:
-        keypoints = fix_corner_markers(keypoints, image_array.shape)
+    keypoints = np.asarray(keypoints)
 
-    # use standard keypoints if no custom ones are provided
-    if not reference_keypoints:
-        dpi = guess_dpi(image_array)
-        reference_keypoints = original_corner_markers(page_format, dpi)
+    if not len(keypoints):
+        raise RuntimeError("No keypoints provided for alignment.")
 
-    if len(reference_keypoints) != 4:
-        # this function assumes that the template has the same dimensions as the input image
-        reference_keypoints = fix_corner_markers(reference_keypoints, image_array.shape)
+    # generate the coordinates where the markers should be
+    dpi = guess_dpi(image_array)
+    reference_keypoints = original_corner_markers(page_format, dpi)
+
+    # create a matrix with the distances between each keypoint and match the keypoint sets
+    dists = spatial.distance.cdist(keypoints, reference_keypoints)
+
+    idxs = np.argmin(dists, 1)  # apply to column 1 so indices for input keypoints
+    adjusted_markers = reference_keypoints[idxs]
+
+    if len(adjusted_markers) == 1:
+        x_shift, y_shift = np.subtract(adjusted_markers[0], keypoints[0])
+        return shift_image(image_array, x_shift, y_shift)
 
     rows, cols, _ = image_array.shape
 
     # get the transformation matrix
-    M = cv2.getPerspectiveTransform(np.float32(keypoints), np.float32(reference_keypoints))
+    M = cv2.estimateAffinePartial2D(keypoints, adjusted_markers)[0]
+
     # apply the transformation matrix and fill in the new empty spaces with white
-    return_image = cv2.warpPerspective(image_array, M, (cols, rows),
-                                       borderValue=(255, 255, 255, 255))
+    return_image = cv2.warpAffine(image_array, M, (cols, rows),
+                                  borderValue=(255, 255, 255, 255))
 
     return return_image
 
@@ -742,3 +745,28 @@ def original_corner_markers(format, dpi):
                      (right_x, top_y),
                      (left_x, bottom_y),
                      (right_x, bottom_y)])
+
+
+def shift_image(image_array, x_shift, y_shift):
+    """
+    take an image and shift it along the x and y axis using opencv
+
+    params:
+    -------
+    image_array: numpy.array
+        an array of the image to be shifted
+    x_shift: int
+        indicates how many pixels it has to be shifted on the x-axis, so from left to right
+    y_shift: int
+        indicates how many pixels it has to be shifted on the y-axis, so from top to bottom
+    returns:
+    --------
+    a numpy array of the image shifted where empty spaces are filled with white
+
+    """
+
+    M = np.float32([[1, 0, x_shift],
+                   [0, 1, y_shift]])
+    h, w, _ = image_array.shape
+    return cv2.warpAffine(image_array, M, (w, h),
+                          borderValue=(255, 255, 255, 255))