import sys, os
import numpy as np
import cv2

inpath = sys.argv[1] if sys.argv[1:] else 'tNzoqGH.jpg'
im = cv2.imread(inpath)
h,w = im.shape[:2]

# where's the map?
if 1:
    grayim = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    mask = (grayim <= 254) * np.uint8(255)

    kernel = cv2.getStructuringElement(shape=cv2.MORPH_RECT, ksize=(9,9))

    mask = cv2.morphologyEx(src=mask, op=cv2.MORPH_CLOSE, kernel=kernel, iterations=1)
    mask = cv2.morphologyEx(src=mask, op=cv2.MORPH_OPEN, kernel=kernel, iterations=4)

    contourimg, contours, hierarchy = \
        cv2.findContours(image=mask, mode=cv2.RETR_LIST, method=cv2.CHAIN_APPROX_TC89_KCOS)

    rectangles = []
    for c in contours:
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.04*peri, True)
        if len(approx) == 4:
            rectangles.append(approx)

    rectangles.sort(key=lambda c: cv2.contourArea(c))
    #print [cv2.arcLength(c, True) for c in rectangles]
    largestrect = rectangles[-1]
    print largestrect[:,0] # eh...

    outim = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)

    cv2.drawContours(outim, [largestrect], -1, (0, 0, 255), 5)

    cv2.imwrite(os.path.splitext(inpath)[0] + "-mask.png", outim)



if 0:
    K = 4

    im = cv2.cvtColor(im, cv2.COLOR_BGR2YUV)

    pixels = im.reshape((-1, 3)).astype(np.float32)
    pixels = pixels[:,1:3]

    crit = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_MAX_ITER, 10, 0.1)
    rv, labels, centers = cv2.kmeans(pixels, K=K, bestLabels=None, criteria=crit, attempts=1, flags=0)

    labelcols = np.hstack([ [[0]]*K, centers ])
    labelcols[:,0] = 128
    labelcols[labels[0][0]] = (0, 128, 128)


    out = labelcols[labels[:,0]]
    out = out.reshape((h,w,3)).clip(0, 255).astype(np.uint8)
    #out[:,:,0] = im[:,:,0]
    out = cv2.cvtColor(out, cv2.COLOR_YUV2BGR)

    cv2.imwrite(os.path.splitext(inpath)[0] + "-clustered.png", out)