# a digital cake knife in the spirit of Hannah Höch
# cuts out composable contours from the herbarium scans
import glob
import os.path
import sqlite3

import cv2
import numpy as np
from PIL import Image

BLUR = 3
CANNY_THRESH_1 = 200
CANNY_THRESH_2 = 250
MASK_DILATE_ITER = 40
MASK_ERODE_ITER = 40
MASK_COLOR = (0.0,0.0,1.0) # In BGR format

src_imgs = glob.glob("specimen_img_raw/*")

db = sqlite3.connect("ratios.db")
dbc = db.cursor()

for src in src_imgs:
    img_color = cv2.imread(src)
    scalar = float(1000.0 / img_color.shape[1])
    new_height = int(scalar * img_color.shape[0])
    img_color = cv2.resize(img_color, (1000, new_height))
    img_gray = cv2.imread(src, flags=cv2.IMREAD_GRAYSCALE)
    img_gray = cv2.resize(img_gray, (1000, new_height))

    # add some white around the edges so we have some space to rotate
    img_color = cv2.copyMakeBorder(img_color,10,10,10,10,cv2.BORDER_CONSTANT, value=[255, 255, 255])
    img_gray = cv2.copyMakeBorder(img_gray,10,10,10,10,cv2.BORDER_CONSTANT, value=[255])
    blurred = cv2.GaussianBlur(img_gray, (9, 9), 0)

    aperture = 5  # default is 3 but 5 or 7 are more sensitive

    # most of this is ganked directly from this https://stackoverflow.com/questions/29313667/how-do-i-remove-the-background-from-this-kind-of-image

    #-- Edge detection -------------------------------------------------------------------
    edges = cv2.Canny(blurred, CANNY_THRESH_1, CANNY_THRESH_2, apertureSize=aperture)
    edges = cv2.dilate(edges, None)
    edges = cv2.erode(edges, None)

    #-- Find contours in edges, sort by area ---------------------------------------------
    contour_info = []
    contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    for c in contours:
        contour_info.append((
            c,
            cv2.isContourConvex(c),
            cv2.contourArea(c),
        ))
    contour_info = sorted(contour_info, key=lambda c: c[2], reverse=True)

    # filter out the accidental box contours - should be less than 85% of pixels
    img_area = img_color.shape[0] * img_color.shape[1]
    for cont in contour_info:
        max_contour = cont
        pixel_ratio = max_contour[2] / img_area
        if pixel_ratio < 0.85:
            test_img = img_color.copy()
            c_rect = cv2.minAreaRect(max_contour[0])
            box = cv2.boxPoints(c_rect)
            box = np.int0(box)
            cv2.drawContours(test_img, [max_contour[0]], -1, (255, 0, 0, 20), 3)
            cv2.drawContours(test_img, [box], -1, (0, 255, 0), 3)
            cv2.imwrite("test.png", test_img)
            test_viewer = Image.open("test.png")
            test_viewer.show()
            if input("contour ok (y/n)?") == "y":
                test_viewer.close()
                break
            test_viewer.close()

    #-- Create empty mask, draw filled polygon on it corresponding to largest contour ----
    # Mask is black, polygon is white
    mask = np.zeros(edges.shape)
    cv2.fillConvexPoly(mask, max_contour[0], (255))

    #-- Smooth mask, then blur it --------------------------------------------------------
    mask = cv2.dilate(mask, None, iterations=MASK_DILATE_ITER)
    mask = cv2.erode(mask, None, iterations=MASK_ERODE_ITER)
    mask = cv2.GaussianBlur(mask, (BLUR, BLUR), 0)
    mask_stack = np.dstack([mask]*3)    # Create 3-channel alpha mask

    #-- Blend masked img into MASK_COLOR background --------------------------------------
    mask_stack = mask_stack.astype('float32') / 255.0          # Use float matrices,
    img_color = img_color.astype('float32') / 255.0                 #  for easy blending

    # split image into channels
    try:
        c_red, c_green, c_blue = cv2.split(img_color)
    except ValueError:
        print("seems to be greyscale already...")
        img_color = cv2.cvtColor(img_color, cv2.COLOR_GRAY2BGR)
        c_red, c_green, c_blue = cv2.split(img_color)

    # merge with mask got on one of a previous steps
    img_a = cv2.merge((c_red, c_green, c_blue, mask.astype('float32') / 255.0))

    # find bounding minimum bounding rect as that's what we want to rotate & save
    rect = cv2.minAreaRect(max_contour[0])
    box = cv2.boxPoints(rect)
    box = np.int0(box)

    width = int(rect[1][0])
    height = int(rect[1][1])

    # set up a new destination exactly the size of our ROI
    src_pts = box.astype("float32")
    dst_pts = np.array([[0, height-1],
                        [0, 0],
                        [width-1, 0],
                        [width-1, height-1]], dtype="float32")

    # now rotate using warp which is more efficient and preserves pixels
    M = cv2.getPerspectiveTransform(src_pts, dst_pts)
    warped = cv2.warpPerspective(img_a, M, (width, height))

    # save to disk
    basename, ext = os.path.splitext(os.path.basename(src))
    out_path = os.path.join("specimen_cutout", basename + ".png")
    print("saving cropped cutout", out_path)
    cv2.imwrite(out_path, warped*255)

    # add to database
    if height > width:
        ratio = float(height) / width
    else:
        ratio = float(width) / height

    try:
        dbc.execute("INSERT INTO images VALUES (?, ?, ?, ?)", (ratio, width, height, out_path))
    except sqlite3.IntegrityError as err:
        print(err)
        print("Trying db update instead.")
        dbc.execute("UPDATE images SET (ratio, width, height) = (?, ?, ?) WHERE img_path = ?", 
                    (ratio, width, height, out_path))


db.commit()
db.close()