Commit 73156285 authored by Fabio Navarrete's avatar Fabio Navarrete
Browse files

Added second algorithm for background removal

parent e7254bef
import numpy as np
import cv2
import matplotlib.pyplot as plt
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)
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;
return sobel
def findSignificantContours (img, sobel_8u):
image, contours, heirarchy = cv2.findContours(sobel_8u, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
def findSignificantContours(img, sobel_8u):
image, contours, heirarchy = cv2.findContours(sobel_8u, cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
# Find level 1 contours
level1 = []
......@@ -23,32 +29,40 @@ def findSignificantContours (img, sobel_8u):
# 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
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]];
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)
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];
return [x[0] for x in significant]
def segment (img):
blurred = cv2.GaussianBlur(img, (5, 5), 0) # Remove noise
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 )
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;
sobel[sobel <= mean] = 0
sobel[sobel > 255] = 255
cv2.imwrite('output/edge.png', sobel);
cv2.imwrite('output/edge.png', sobel)
sobel_8u = np.asarray(sobel, np.uint8)
......@@ -62,14 +76,63 @@ def segment (img):
# Invert mask
mask = np.logical_not(mask)
#Finally remove the background
img[mask] = 0;
# 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')
def isolate_shape(img, rect=None, only_shape=True):
img = cv2.GaussianBlur(img, (5, 5), 0)
# plt.imshow(img)
# plt.show(block=False)
height, width = img.shape[:2]
# Create a mask holder
mask = np.zeros(img.shape[:2], np.uint8)
# Grab Cut the object
bgdModel = np.zeros((1, 65), np.float64)
fgdModel = np.zeros((1, 65), np.float64)
# "Hard Coding the Rect The object must lie within this rect."
if not rect:
rect = (300, 0, 100, height / 2)
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 2, cv2.GC_INIT_WITH_RECT)
# grabCut
mask = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
img1 = img * mask[:, :, np.newaxis]
# Get the background
background = img - img1
# Change all pixels in the background that are not black to white
background[np.where((background > [0, 0, 0]).all(axis=2))] = [
255, 255, 255
]
# Add the background and the image
final = background + img1
if only_shape:
return background
return final
# Load the Image
# img = cv2.imread(
# '/home/fabio/dev/workspaces/personal/ctai/GarbageBot/src/detecting/dataset/REDTRIANGLE.right.1.jpg'
# )
#segment('original-small.jpg')
# plt.imshow(isolate_shape(img, rect=(300, 0, 100, 300), only_shape=False))
# plt.show()
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