import numpy as np
import cv2
from matplotlib import pyplot as plt
import os

print('-----')
dirname = os.path.dirname(__file__)
filename = os.path.join(dirname,'data/legos3.jpg')
print(os.path.exists(filename))
img = cv2.imread(filename)
b,g,r = cv2.split(img)
rgb_img = cv2.merge([r,g,b])

gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

# smooth = cv2.bilateralFilter(gray,5,150,150)
# smooth = cv2.medianBlur(gray,5)
smooth = cv2.GaussianBlur(gray,(5,5),0)

# sobelx = cv2.Sobel(smooth,cv2.CV_8UC1,1,0,ksize=5)
# sobely = cv2.Sobel(smooth,cv2.CV_8UC1,0,1,ksize=5)

# ret, thresh = cv2.threshold(smooth,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
# thresh = cv2.adaptiveThreshold(smooth,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,21,5)
threshLevel = np.mean(smooth)*1.6
ret, thresh = cv2.threshold(smooth,threshLevel,255,cv2.THRESH_BINARY)

#thresh = cv2.Canny(smooth,200,200)

# noise removal
kernel = np.ones((3,3),np.uint8)
#thresh = cv2.dilate(thresh,kernel,iterations=1)
opening = cv2.morphologyEx(thresh,cv2.MORPH_OPEN,kernel, iterations = 5)
# closing = cv2.morphologyEx(thresh,cv2.MORPH_CLOSE,kernel, iterations = 2)

# sure background area
sure_bg = cv2.dilate(opening,kernel,iterations=3)

# Finding sure foreground area
dist_transform = cv2.distanceTransform(sure_bg,cv2.DIST_L2,3)

# Threshold
ret, sure_fg = cv2.threshold(dist_transform,0.1*dist_transform.max(),255,0)

# Finding unknown region
sure_fg = np.uint8(sure_fg)
unknown = cv2.subtract(sure_bg,sure_fg)

# Marker labelling
ret, markers = cv2.connectedComponents(sure_fg)

# Add one to all labels so that sure background is not 0, but 1
markers = markers+1

# Now, mark the region of unknown with zero
markers[unknown==255] = 0

markers = cv2.watershed(img,markers)
img[markers == -1] = [255,0,0]

plt.subplot(421),plt.imshow(rgb_img)
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
# plt.subplot(422),plt.imshow(thresh, 'gray')
# plt.title("Otsu's binary threshold"), plt.xticks([]), plt.yticks([])
plt.subplot(422),plt.imshow(thresh, 'gray')
plt.title("smoothing"), plt.xticks([]), plt.yticks([])

plt.subplot(423),plt.imshow(opening, 'gray')
plt.title("morphologyEx:Closing:2x2"), plt.xticks([]), plt.yticks([])
plt.subplot(424),plt.imshow(sure_bg, 'gray')
plt.title("Dilation"), plt.xticks([]), plt.yticks([])

plt.subplot(425),plt.imshow(dist_transform, 'gray')
plt.title("Distance Transform"), plt.xticks([]), plt.yticks([])
plt.subplot(426),plt.imshow(sure_fg, 'gray')
plt.title("Thresholding"), plt.xticks([]), plt.yticks([])

plt.subplot(427),plt.imshow(unknown, 'gray')
plt.title("Unknown"), plt.xticks([]), plt.yticks([])

plt.subplot(428),plt.imshow(img, 'gray')
plt.title("Result from Watershed"), plt.xticks([]), plt.yticks([])

plt.tight_layout()
plt.show()