add: torchvisionのsegmentationで背景マスキング

app.py

@@ -6,83 +6,62 @@ import random
 import time
 import seaborn as sns
 
-def get_concat_h(im1, im2):
-    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
-    dst.paste(im1, (0, 0))
-    dst.paste(im2, (im1.width, 0))
-    return dst
-
-def get_concat_v(im1, im2):
-    dst = Image.new('RGB', (im1.width, im1.height + im2.height))
-    dst.paste(im1, (0, 0))
-    dst.paste(im2, (0, im1.height))
-    return dst
-
-def hsv_to_rgb(h, s, v):
-    bgr = cv2.cvtColor(np.array([[[h, s, v]]], dtype=np.uint8), cv2.COLOR_HSV2BGR)[0][0]
-    return [bgr[2]/255, bgr[1]/255, bgr[0]/255]
+from cv_funcs import *
+from torchvision_funcs import *
 
 @st.cache
 def show_generated_image(image):
     st.image(image)
-
+    
 @st.cache(suppress_st_warning=True)
-def randomize_palette_colors(n_rows, n_cols, palette1="Set1", palette2="Set2", seed=0):
+def randomize_palette_colors(n_rows, n_cols, palettes=['Set1', 'Set3', 'Spectral'], seed=time.time(), n_times=10):
     random.seed(seed)
-    colors1 = sns.color_palette(palette1, n_rows*n_cols)
-    colors2 = sns.color_palette(palette2, n_rows*n_cols)    
-    colors1, colors2 = random.sample(colors1, len(colors1)), random.sample(colors2, len(colors2))
-    return colors1, colors2
+    colors = [sns.color_palette(palette, n_rows*n_cols*n_times) for palette in palettes]
+    _ = [random.shuffle(color) for color in colors]
+    return colors
 
 @st.cache(suppress_st_warning=True)
-def randomize_rgb_colors(n_rows, n_cols, seed=0):
-    random.seed(seed)    
-    colors1 = [[random.random() for j in range(3)] for i in range(n_rows*n_cols)] 
-    colors2 = [[random.random() for j in range(3)] for i in range(n_rows*n_cols)]         
-    return colors1, colors2
-
-@st.cache(suppress_st_warning=True)
-def randomize_hsv_colors(n_rows, n_cols, s=255, v=255, seed=0):
-    random.seed(seed)    
-    colors1 = [hsv_to_rgb(random.random()*180, s, v) for i in range(n_rows*n_cols)] 
-    colors2 = [hsv_to_rgb(random.random()*180, s, v) for i in range(n_rows*n_cols)]         
-    return colors1, colors2
+def remove_image_background(image):
+    return deeplabv3_remove_bg(image)
 
 title = 'Andy Warhol like Image Generator'
-st.set_page_config(page_title=title, layout='centered')
+st.set_page_config(page_title=title, page_icon='favicon.jpeg', layout='centered')
 st.title(title)
 uploaded_file = st.file_uploader('Choose an image file')
 if uploaded_file is None: uploaded_file = './sample.jpg'
 
 if uploaded_file is not None:
     im = Image.open(uploaded_file)
-    im.thumbnail((1000, 1000),resample=Image.BICUBIC) # resize    
-    st.image(im, caption='Original')
-
-    im_gray = np.array(im.convert('L'))
+    im.thumbnail((1000, 1000),resample=Image.BICUBIC) # resize
+    
+    is_masked = st.checkbox('With background masking? (3 colors)')
+    if is_masked:
+       im_masked, index_masked = remove_image_background(im)
+       st.image(im_masked, caption='Masked image')        
+    else: st.image(im, caption='Original')
+    
+    im_gray =  np.array(im.convert('L'))
     thresh, _img = cv2.threshold(im_gray, 0, 255, cv2.THRESH_OTSU)
     
     n_rows, n_cols = st.number_input('Rows', value=3), st.number_input('Columns', value=3)
-    
-    # s = st.slider('Saturation', value=125.0, min_value=0.0, max_value=255.0)    
-    # v = st.slider('Brightness', value=255.0, min_value=0.0, max_value=255.0)
-    colors1, colors2 = randomize_palette_colors(n_rows, n_cols)
-    thresh = st.slider('Threshold', value=thresh, min_value=0.0, max_value=255.0)
+
+    thresh = st.slider('Threshold', value=thresh, min_value=0.0, max_value=255.0)        
+    colors = randomize_palette_colors(n_rows, n_cols, seed=0)
     
     if st.button('Shuffle colors'):
-        colors1, colors2 = randomize_palette_colors(n_rows, n_cols, seed=time.time())
-
+        colors = randomize_palette_colors(n_rows, n_cols, seed=time.time())            
+    
     if True or st.button('Generate'):
-        im_bool = im_gray > thresh
-
         ims_generated = []
+
         for row in range(n_rows):
             for col in range(n_cols):
                 i_color = n_cols * row + col
-                rgb1, rgb2 = np.array(colors1[i_color])*np.array([255, 255, 255]).tolist(), np.array(colors2[i_color])*np.array([255, 255, 255]).tolist()
+                rgbs = [np.array(color[i_color])*np.array([255, 255, 255]).tolist() for color in colors]
                 ims_col = np.empty((*im_gray.shape, 3))
                 for i in range(3): # RGB
-                     ims_col[:, :, i] = (im_gray > thresh) * rgb1[i] + (im_gray <= thresh) * rgb2[i]
+                     ims_col[:, :, i] = (im_gray <= thresh) * rgbs[0][i] + (im_gray > thresh) * rgbs[1][i]
+                     if is_masked: ims_col[:, :, i][index_masked] = rgbs[2][i]
                 if col == 0:
                     im_col_concat = Image.fromarray(ims_col.astype(np.uint8))
                 else:

cv_funcs.py

@@ -0,0 +1,67 @@
+import cv2
+from PIL import Image
+import numpy as np
+
+def get_concat_h(im1, im2):
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+def get_concat_v(im1, im2):
+    dst = Image.new('RGB', (im1.width, im1.height + im2.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (0, im1.height))
+    return dst
+
+def hsv_to_rgb(h, s, v):
+    bgr = cv2.cvtColor(np.array([[[h, s, v]]], dtype=np.uint8), cv2.COLOR_HSV2BGR)[0][0]
+    return [bgr[2]/255, bgr[1]/255, bgr[0]/255]
+
+# def remove_bg(
+#     path,
+#     BLUR = 21,
+#     CANNY_THRESH_1 = 10,
+#     CANNY_THRESH_2 = 200,
+#     MASK_DILATE_ITER = 10,
+#     MASK_ERODE_ITER = 10,
+#     MASK_COLOR = (0.0,0.0,1.0),
+# ):
+#     img = cv2.imread(path)
+#     gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
+
+#     edges = cv2.Canny(gray, CANNY_THRESH_1, CANNY_THRESH_2)
+#     edges = cv2.dilate(edges, None)
+#     edges = cv2.erode(edges, None)
+
+#     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))
+
+#     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
+
+#     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 
+
+#     c_blue, c_green, c_red = cv2.split(img)
+
+#     img_a = cv2.merge((c_red, c_green, c_blue, mask.astype('float32') / 255.0))
+#     index = np.where(img_a[:, :, 3] == 0)
+#     #img_a[index] = [1.0, 1.0, 1.0, 1.0]
+#     return img_a

requirements.txt

@@ -1,4 +1,6 @@
 streamlit==0.76.0
 Pillow
 opencv-python
-seaborn
+seaborn
+http://download.pytorch.org/whl/cpu/torch-1.4.0%2Bcpu-cp36-cp36m-linux_x86_64.whl
+torchvision==0.5.0

torchvision_funcs.py

@@ -0,0 +1,41 @@
+import numpy as np
+import cv2
+from PIL import Image
+
+import torch
+import torchvision
+from torchvision import transforms
+
+def deeplabv3_remove_bg(img):
+    img = np.array(img, dtype=np.uint8)
+    # img = cv2.imread(image_path)
+    # img = img[...,::-1] #BGR->RGB
+    h,w,_ = img.shape
+    # img = cv2.resize(img,(1000,1000))
+    
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)
+    model = model.to(device)
+    model.eval();
+    
+    preprocess = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    input_tensor = preprocess(img)
+    input_batch = input_tensor.unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        output = model(input_batch)['out'][0]
+    output = output.argmax(0)
+    mask = output.byte().cpu().numpy()
+    # mask = cv2.resize(mask,(w,h))
+    # img = cv2.resize(img,(w,h))
+    mask[mask>0] = 1.0 # NOTE: なぜか3が入っていたので
+    mask = np.dstack([mask, mask, mask])
+    img_masked = Image.fromarray(cv2.multiply(img, mask))
+    index_masked = np.where(np.array(mask)[:,:,2]==0)
+    return img_masked, index_masked
+    
+