fix.py

#! /usr/bin/env python3
import numpy as np
import cv2
from matplotlib import pyplot as plt
import os

def order_points(pts):
    rect = np.zeros((4,2), dtype = "float32")

    s = pts.sum(axis = 1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]

    diff = np.diff(pts, axis = 1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]

    return rect

def four_point_transform(image, pts):
    rect = order_points(pts)
    (tl, tr, br, bl) = rect

    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[0] - bl[0]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[0] - tl[0]) ** 2))
    maxWidth = max(int(widthA), int(widthB))

    heightA = np.sqrt(((tr[1] - br[1]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[1] - bl[1]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))

    dst = np.array([
        [0, 0],
        [maxWidth - 1, 0],
        [maxWidth - 1, maxHeight - 1],
        [0, maxHeight - 1]], dtype = "float32")

    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))

    return warped

a = np.array([0,0], dtype="float32")
def getXY(img_path):
    global a
    a = np.array([0,0], dtype="float32")

    img = cv2.imread(img_path)

    screen_res = 1280.0, 720.0
    scale_width = screen_res[0]/img.shape[1]
    scale_height = screen_res[1] / img.shape[0]
    scale = min(scale_width, scale_height)
    win_width = int(img.shape[1]*scale)
    win_height = int(img.shape[0]*scale)

    def getxy(event, x, y, flags, param):
        global a
        if event == cv2.EVENT_LBUTTONDOWN :
            a = np.vstack([a, np.hstack([int(x/scale),int(y/scale)])])

    newimg = cv2.resize(img, (win_width, win_height))

    cv2.namedWindow('Original', cv2.WINDOW_NORMAL)
    cv2.setMouseCallback('Original', getxy)
    cv2.imshow('Original', newimg)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    b = a[1:,:]
    return b




counter = 1

while counter < 5:
    filename = "data/%d.jpg" % (counter)
    points = getXY(filename)
    points2 = getXY(filename)

    image = cv2.imread(filename)

    warped = four_point_transform(image, points)
    warped2 = four_point_transform(image, points2)

    filename1 = "%d.jpg" % (counter)
    counter = counter + 1
    filename2 = "%d.jpg" % (counter)
    counter = counter + 1
    cv2.imwrite(filename1, warped)
    cv2.imwrite(filename2, warped2)