Merge branch 'master' of https://github.com/Ardillen66/GarbageBot

src/features-detection/colorDetection.py

+import numpy as np
+import cv2
+import argparse
+
+def getSobel (channel):
+
+    sobelx = cv2.Sobel(channel, cv2.CV_16S, 1, 0, borderType=cv2.BORDER_REPLICATE)
+    sobely = cv2.Sobel(channel, cv2.CV_16S, 0, 1, borderType=cv2.BORDER_REPLICATE)
+    sobel = np.hypot(sobelx, sobely)
+
+    return sobel;
+
+def findSignificantContours (img, sobel_8u):
+    image, contours, heirarchy = cv2.findContours(sobel_8u, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+    # Find level 1 contours
+    level1 = []
+    for i, tupl in enumerate(heirarchy[0]):
+        # Each array is in format (Next, Prev, First child, Parent)
+        # Filter the ones without parent
+        if tupl[3] == -1:
+            tupl = np.insert(tupl, 0, [i])
+            level1.append(tupl)
+
+    # From among them, find the contours with large surface area.
+    significant = []
+    tooSmall = sobel_8u.size * 5 / 100 # If contour isn't covering 5% of total area of image then it probably is too small
+    for tupl in level1:
+        contour = contours[tupl[0]];
+        area = cv2.contourArea(contour)
+        if area > tooSmall:
+            cv2.drawContours(img, [contour], 0, (0,255,0),2, cv2.LINE_AA, maxLevel=1)
+            significant.append([contour, area])
+
+    significant.sort(key=lambda x: x[1])
+    return [x[0] for x in significant];
+
+def getColor(colorValues, colorNames, maxValue):
+    for i, item in enumerate(colorValues):
+        if(item == maxValue):
+            return colorNames[i]
+
+def segment (path):
+    img = cv2.imread(path)
+
+    blurred = cv2.GaussianBlur(img, (5, 5), 0) # Remove noise
+
+    # Edge operator
+    sobel = np.max( np.array([ getSobel(blurred[:,:, 0]), getSobel(blurred[:,:, 1]), getSobel(blurred[:,:, 2]) ]), axis=0 )
+
+
+    mean = np.median(sobel)
+
+    # Zero any values less than mean. This reduces a lot of noise.
+    sobel[sobel <= mean] = 0;
+    sobel[sobel > 255] = 255;
+
+    cv2.imwrite('output/edge.png', sobel);
+
+    sobel_8u = np.asarray(sobel, np.uint8)
+
+    # Find contours
+    significant = findSignificantContours(img, sobel_8u)
+
+    # Mask
+    mask = sobel.copy()
+    mask[mask > 0] = 0
+    cv2.fillPoly(mask, significant, 255)
+    # Invert mask
+    mask = np.logical_not(mask)
+
+    #Finally remove the background
+    img[mask] = 0;
+
+    fname = path.split('/')[-1]
+    cv2.imwrite('output/' + fname, img);
+    print (path)
+    
+    #cv2.imshow("Image", img)
+    #cv2.waitKey(0)
+
+    # construct the argument parse and parse the arguments
+    ap = argparse.ArgumentParser()
+    ap.add_argument("-i", "--image", help = "path to the image")
+    args = vars(ap.parse_args())
+
+    # load the image
+    image =  img # cv2.imread(args["image"])
+
+    # define the list of boundaries
+    boundaries = [
+	([17, 15, 100], [255, 56, 200]), #red
+	([86, 31, 4], [220, 88, 50]),  #blue
+	([25, 146, 190], [62, 174, 250]), #yellowish
+	([103, 86, 65], [145, 133, 128]) #grey
+    ]
+
+    colorValues = []
+
+    # loop over the boundaries
+    for (lower, upper) in boundaries:
+        # create NumPy arrays from the boundaries
+	lower = np.array(lower, dtype = "uint8")
+	upper = np.array(upper, dtype = "uint8")
+
+	# find the colors within the specified boundaries and apply
+	# the mask
+	mask = cv2.inRange(image, lower, upper)
+        #print(cl.label(image, mask))
+
+        # save the values
+        colorValues.append(np.count_nonzero(mask))
+
+        #print(np.count_nonzero(mask))
+	output = cv2.bitwise_and(image, image, mask = mask)
+        
+	# show the images
+	#cv2.imshow("images", np.hstack([image, output]))
+	#cv2.waitKey(0)
+
+    print(getColor(colorValues, ['RED', 'BLUE', 'YELLOW(ish)', 'GREY'], max(colorValues)))
+
+    
+
+segment('bal1.jpg')
+
+
+
+
+