summaryrefslogtreecommitdiff
path: root/samples/python/tutorial_code/ImgTrans/distance_transformation/imageSegmentation.py
diff options
context:
space:
mode:
Diffstat (limited to 'samples/python/tutorial_code/ImgTrans/distance_transformation/imageSegmentation.py')
-rw-r--r--samples/python/tutorial_code/ImgTrans/distance_transformation/imageSegmentation.py138
1 files changed, 138 insertions, 0 deletions
diff --git a/samples/python/tutorial_code/ImgTrans/distance_transformation/imageSegmentation.py b/samples/python/tutorial_code/ImgTrans/distance_transformation/imageSegmentation.py
new file mode 100644
index 0000000000..e679001bc1
--- /dev/null
+++ b/samples/python/tutorial_code/ImgTrans/distance_transformation/imageSegmentation.py
@@ -0,0 +1,138 @@
+from __future__ import print_function
+import cv2 as cv
+import numpy as np
+import argparse
+import random as rng
+
+rng.seed(12345)
+
+## [load_image]
+# Load the image
+parser = argparse.ArgumentParser(description='Code for Image Segmentation with Distance Transform and Watershed Algorithm.\
+ Sample code showing how to segment overlapping objects using Laplacian filtering, \
+ in addition to Watershed and Distance Transformation')
+parser.add_argument('--input', help='Path to input image.', default='../data/cards.png')
+args = parser.parse_args()
+
+src = cv.imread(args.input)
+if src is None:
+ print('Could not open or find the image:', args.input)
+ exit(0)
+
+# Show source image
+cv.imshow('Source Image', src)
+## [load_image]
+
+## [black_bg]
+# Change the background from white to black, since that will help later to extract
+# better results during the use of Distance Transform
+src[np.all(src == 255, axis=2)] = 0
+
+# Show output image
+cv.imshow('Black Background Image', src)
+## [black_bg]
+
+## [sharp]
+# Create a kernel that we will use to sharpen our image
+# an approximation of second derivative, a quite strong kernel
+kernel = np.array([[1, 1, 1], [1, -8, 1], [1, 1, 1]], dtype=np.float32)
+
+# do the laplacian filtering as it is
+# well, we need to convert everything in something more deeper then CV_8U
+# because the kernel has some negative values,
+# and we can expect in general to have a Laplacian image with negative values
+# BUT a 8bits unsigned int (the one we are working with) can contain values from 0 to 255
+# so the possible negative number will be truncated
+imgLaplacian = cv.filter2D(src, cv.CV_32F, kernel)
+sharp = np.float32(src)
+imgResult = sharp - imgLaplacian
+
+# convert back to 8bits gray scale
+imgResult = np.clip(imgResult, 0, 255)
+imgResult = imgResult.astype('uint8')
+imgLaplacian = np.clip(imgLaplacian, 0, 255)
+imgLaplacian = np.uint8(imgLaplacian)
+
+#cv.imshow('Laplace Filtered Image', imgLaplacian)
+cv.imshow('New Sharped Image', imgResult)
+## [sharp]
+
+## [bin]
+# Create binary image from source image
+bw = cv.cvtColor(imgResult, cv.COLOR_BGR2GRAY)
+_, bw = cv.threshold(bw, 40, 255, cv.THRESH_BINARY | cv.THRESH_OTSU)
+cv.imshow('Binary Image', bw)
+## [bin]
+
+## [dist]
+# Perform the distance transform algorithm
+dist = cv.distanceTransform(bw, cv.DIST_L2, 3)
+
+# Normalize the distance image for range = {0.0, 1.0}
+# so we can visualize and threshold it
+cv.normalize(dist, dist, 0, 1.0, cv.NORM_MINMAX)
+cv.imshow('Distance Transform Image', dist)
+## [dist]
+
+## [peaks]
+# Threshold to obtain the peaks
+# This will be the markers for the foreground objects
+_, dist = cv.threshold(dist, 0.4, 1.0, cv.THRESH_BINARY)
+
+# Dilate a bit the dist image
+kernel1 = np.ones((3,3), dtype=np.uint8)
+dist = cv.dilate(dist, kernel1)
+cv.imshow('Peaks', dist)
+## [peaks]
+
+## [seeds]
+# Create the CV_8U version of the distance image
+# It is needed for findContours()
+dist_8u = dist.astype('uint8')
+
+# Find total markers
+_, contours, _ = cv.findContours(dist_8u, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
+
+# Create the marker image for the watershed algorithm
+markers = np.zeros(dist.shape, dtype=np.int32)
+
+# Draw the foreground markers
+for i in range(len(contours)):
+ cv.drawContours(markers, contours, i, (i+1), -1)
+
+# Draw the background marker
+cv.circle(markers, (5,5), 3, (255,255,255), -1)
+cv.imshow('Markers', markers*10000)
+## [seeds]
+
+## [watershed]
+# Perform the watershed algorithm
+cv.watershed(imgResult, markers)
+
+#mark = np.zeros(markers.shape, dtype=np.uint8)
+mark = markers.astype('uint8')
+mark = cv.bitwise_not(mark)
+# uncomment this if you want to see how the mark
+# image looks like at that point
+#cv.imshow('Markers_v2', mark)
+
+# Generate random colors
+colors = []
+for contour in contours:
+ colors.append((rng.randint(0,256), rng.randint(0,256), rng.randint(0,256)))
+
+# Create the result image
+dst = np.zeros((markers.shape[0], markers.shape[1], 3), dtype=np.uint8)
+
+# Fill labeled objects with random colors
+for i in range(markers.shape[0]):
+ for j in range(markers.shape[1]):
+ index = markers[i,j]
+ if index > 0 and index <= len(contours):
+ dst[i,j,:] = colors[index-1]
+
+# Visualize the final image
+cv.imshow('Final Result', dst)
+## [watershed]
+
+cv.waitKey()