I would like to change the background of this image to a white background in python. I tried canny edge detection, but it has a hard time finding the edges at the top of the product as you can see in the second picture. I tried different thresholds, but it would result in more background being not white. This is probably due to the fact that the top of the product in the image has almost the same color as the background.
Is there a method to detect a small difference like that? I also tried it with a green screen behind the product, but because of the reflective state of the product, the product turns green.
Here is my code:
import cv2
import numpy as np
from skimage import filters
import matplotlib.pyplot as plt
from os import listdir
from os.path import isfile, join
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'
plt.rcParams['figure.dpi'] = 200
#== Parameters =======================================================================
BLUR = 15
CANNY_THRESH_1 = 10
CANNY_THRESH_2 = 255
MASK_DILATE_ITER = 10
MASK_ERODE_ITER = 10
MASK_COLOR = (1.0,1.0,1.0) # In BGR format
#== Processing =======================================================================
mypath = "./images"
images = [f for f in listdir(mypath) if isfile(join(mypath, f))]
#-- Read image -----------------------------------------------------------------------
for image in images:
img_loc = mypath + "/" + image
img = cv2.imread(img_loc)
# threshold
img_thresh = img
thresh = 180
img_thresh[ img_thresh >= thresh ] = 255
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
#-- Edge detection -------------------------------------------------------------------
edges = cv2.Canny(gray, CANNY_THRESH_1, CANNY_THRESH_2)
edges = cv2.dilate(edges, None)
edges = cv2.erode(edges, None)
#-- Find contours in edges, sort by area ---------------------------------------------
contour_info = []
contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
for c in contours:
contour_info.append((
c,
cv2.isContourConvex(c),
cv2.contourArea(c),
))
contour_info = sorted(contour_info, key=lambda c: c[2], reverse=True)
max_contour = contour_info[0]
mask = np.zeros(edges.shape)
cv2.fillConvexPoly(mask, max_contour[0], (255))
#-- Smooth mask, then blur it --------------------------------------------------------
mask = cv2.dilate(mask, None, iterations=MASK_DILATE_ITER)
mask = cv2.erode(mask, None, iterations=MASK_ERODE_ITER)
mask = cv2.GaussianBlur(mask, (BLUR, BLUR), 0)
mask_stack = np.dstack([mask]*3) # Create 3-channel alpha mask
#-- Blend masked img into MASK_COLOR background --------------------------------------
mask_stack = mask_stack.astype('float32') / 255.0 # Use float matrices,
img = img.astype('float32') / 255.0 # for easy blending
masked = (mask_stack * img) + ((1-mask_stack) * MASK_COLOR) # Blend
masked = (masked * 255).astype('uint8') # Convert back to 8-bit
result_dir = "./results/" + image
cv2.imwrite(result_dir, masked) # Save