Commit 1d31ca79 authored by Mathieu Reymond's avatar Mathieu Reymond
Browse files

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

parents 3169daaa 9b60212a
import numpy as np
import cv2
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 segment (img):
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 )
# Noise reduction trick, from http://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l182
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;
return img
# fname = path.split('/')[-1]
# cv2.imshow('output', img)
# cv2.key
# cv2.imwrite('output/' + fname, img)
# print (path)
#segment('original-small.jpg')
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment