initial commit

watershed.py

@@ -0,0 +1,85 @@
+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()