background_removal

!pip install tensorflow --upgrade

!pip install opencv-python

!pip install keras --upgrade

!pip install tensorboardX

!pip install toml torch torchvision

import os import cv2 import keras import numpy as np import matplotlib.pyplot as plt import glob import pandas as pd import toml import numpy as np import torch from torch.nn import functional as F import cv2 import utils from utils import CONFIG import networks

!ls /datasets/model_seg

BACKBONE = 'inceptionv3' BATCH_SIZE = 4 CLASSES = ['product'] LR = 0.0001 EPOCHS = 10 preprocess_input = sm.get_preprocessing(BACKBONE) n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1) # case for binary and multiclass segmentation activation = 'sigmoid' if n_classes == 1 else 'softmax' model = sm.Unet(BACKBONE, classes=n_classes, activation=activation,encoder_weights='imagenet', encoder_freeze=False,decoder_block_type='transpose',decoder_use_batchnorm=True) model.load_weights('/datasets/model_seg/inception_new_categories.h5')

def single_inference(model, image_dict, return_offset=True,): with torch.no_grad(): image, trimap = image_dict['image'], image_dict['trimap'] alpha_shape = image_dict['alpha_shape'] image = image trimap = trimap alpha_pred, info_dict = model(image, trimap) if CONFIG.model.trimap_channel == 3: trimap_argmax = trimap.argmax(dim=1, keepdim=True) alpha_pred[trimap_argmax == 2] = 1 alpha_pred[trimap_argmax == 0] = 0 h, w = alpha_shape test_pred = alpha_pred[0, 0, ...].data.cpu().numpy() * 255 test_pred = test_pred.astype(np.uint8) test_pred = test_pred[32:h+32, 32:w+32] if return_offset: short_side = h if h < w else w ratio = 512 / short_side offset_1 = utils.flow_to_image(info_dict['offset_1'][0][0,...].data.cpu().numpy()).astype(np.uint8) # write softmax_scale to offset image scale = info_dict['offset_1'][1].cpu() offset_1 = cv2.resize(offset_1, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST) text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1,0].item(), scale[-1,1].item()) offset_1 = cv2.putText(offset_1, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) offset_2 = utils.flow_to_image(info_dict['offset_2'][0][0,...].data.cpu().numpy()).astype(np.uint8) # write softmax_scale to offset image scale = info_dict['offset_2'][1].cpu() offset_2 = cv2.resize(offset_2, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_NEAREST) text = 'unknown: {:.2f}, known: {:.2f}'.format(scale[-1,0].item(), scale[-1,1].item()) offset_2 = cv2.putText(offset_2, text, (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.8, 0, thickness=2) return test_pred, (offset_1, offset_2) else: return test_pred, None def generator_tensor_dict(image, trimap): # read images image = image trimap = trimap image = cv2.resize(image,(int(image.shape[1]/4),int(image.shape[0]/4))) trimap = cv2.resize(trimap,(int(trimap.shape[1]/4),int(trimap.shape[0]/4))) sample = {'image': image, 'trimap': trimap, 'alpha_shape': trimap.shape} # reshape h, w = sample["alpha_shape"] if h % 32 == 0 and w % 32 == 0: padded_image = np.pad(sample['image'], ((32,32), (32, 32), (0,0)), mode="reflect") padded_trimap = np.pad(sample['trimap'], ((32,32), (32, 32)), mode="reflect") sample['image'] = padded_image sample['trimap'] = padded_trimap else: target_h = 32 * ((h - 1) // 32 + 1) target_w = 32 * ((w - 1) // 32 + 1) pad_h = target_h - h pad_w = target_w - w padded_image = np.pad(sample['image'], ((32,pad_h+32), (32, pad_w+32), (0,0)), mode="reflect") padded_trimap = np.pad(sample['trimap'], ((32,pad_h+32), (32, pad_w+32)), mode="reflect") sample['image'] = padded_image sample['trimap'] = padded_trimap # ImageNet mean & std mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1) std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1) # convert GBR images to RGB image, trimap = sample['image'][:,:,::-1], sample['trimap'] # swap color axis image = image.transpose((2, 0, 1)).astype(np.float32) trimap[trimap < 85] = 0 trimap[trimap >= 170] = 2 trimap[trimap >= 85] = 1 # normalize image image /= 255. # to tensor sample['image'], sample['trimap'] = torch.from_numpy(image), torch.from_numpy(trimap).to(torch.long) sample['image'] = sample['image'].sub_(mean).div_(std) if CONFIG.model.trimap_channel == 3: sample['trimap'] = F.one_hot(sample['trimap'], num_classes=3).permute(2, 0, 1).float() elif CONFIG.model.trimap_channel == 1: sample['trimap'] = sample['trimap'][None, ...].float() else: raise NotImplementedError("CONFIG.model.trimap_channel can only be 3 or 1") # add first channel sample['image'], sample['trimap'] = sample['image'][None, ...], sample['trimap'][None, ...] return sample

with open('config/gca-dist-all-data.toml') as f: utils.load_config(toml.load(f)) # Check if toml config file is loaded if CONFIG.is_default: raise ValueError("No .toml config loaded.") model_matting = networks.get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder) #model_matting.cuda() # load checkpoint checkpoint = torch.load('gca-dist-all-data.pth',map_location='cpu') model_matting.load_state_dict(utils.remove_prefix_state_dict(checkpoint['state_dict']), strict=True) # inference model_matting = model_matting.eval()

import cv2 import numpy as np def erode_dilate(msk, struc="ELLIPSE", size=(3, 3)): if struc == "RECT": kernel = cv2.getStructuringElement(cv2.MORPH_RECT, size) elif struc == "CORSS": kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, size) else: kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, size) #msk = msk.astype(np.float32) msk = msk / 255 #msk = msk.astype(np.uint8) # val in 0 or 255 dilated = cv2.dilate(msk, kernel, iterations=10) * 255 eroded = cv2.erode(msk, kernel, iterations=10) * 255 cnt1 = len(np.where(msk >= 0)[0]) cnt2 = len(np.where(msk == 0)[0]) cnt3 = len(np.where(msk == 1)[0]) #print("all:{} bg:{} fg:{}".format(cnt1, cnt2, cnt3)) assert(cnt1 == cnt2 + cnt3) cnt1 = len(np.where(dilated >= 0)[0]) cnt2 = len(np.where(dilated == 0)[0]) cnt3 = len(np.where(dilated == 255)[0]) #print("all:{} bg:{} fg:{}".format(cnt1, cnt2, cnt3)) assert(cnt1 == cnt2 + cnt3) cnt1 = len(np.where(eroded >= 0)[0]) cnt2 = len(np.where(eroded == 0)[0]) cnt3 = len(np.where(eroded == 255)[0]) #print("all:{} bg:{} fg:{}".format(cnt1, cnt2, cnt3)) assert(cnt1 == cnt2 + cnt3) res = dilated.copy() #res[((dilated == 255) & (msk == 0))] = 128 res[((dilated == 255) & (eroded == 0))] = 128 return res

def get_bg_removed_image(url,model,model_matting): ''' image= cv2.imread(path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) ''' image = download_image(url) height,width = image.shape[:2] #actual_image = denormalize(image.squeeze()) #actual_image = cv2.cvtColor(actual_image,cv2.COLOR_BGR2RGB) image_resized = image.copy() image_resized = cv2.resize(image,(224,224)) image_resized = np.expand_dims(image_resized, axis=0) pr_mask = model.predict(preprocess_input(image_resized).reshape(1,224,224,3)).round() alpha = pr_mask.reshape(224,224,1) alpha_seg = alpha.reshape(224,224) alpha_seg = cv2.resize(alpha_seg,(width,height))*255 alpha_seg = cv2.merge([alpha_seg,alpha_seg,alpha_seg]) alpha_seg[np.where(alpha_seg<127)] =0 alpha_seg[np.where(alpha_seg>=127)] =255 #plt.imshow(alpha) alpha_seg = alpha_seg[:,:,0] trimap = erode_dilate(alpha_seg) image_dict = generator_tensor_dict(image, trimap) pred, offset = single_inference(model_matting, image_dict) mat = pred bg = cv2.imread('plain-white-background.jpg') bg = cv2.resize(bg,(image.shape[1],image.shape[0])) pred = cv2.resize(pred,(width,height)).reshape(height, width, 1)/255 im_mat = pred*(image/255) + (1-pred)*(bg/255) visualize( image=image, predicted_mat = alpha_seg, mat = mat, trimap = trimap, image_mat = im_mat #using_mat=im_seg_mat, )

!pip install Pillow

def visualize(**images): """PLot images in one row.""" n = len(images) plt.figure(figsize=(16, 5)) for i, (name, image) in enumerate(images.items()): plt.subplot(1, n, i + 1) plt.xticks([]) plt.yticks([]) plt.title(' '.join(name.split('_')).title()) plt.imshow(image) import numpy as np import requests from PIL import Image def download_image(url): response = requests.get(url, stream=True) response.raw.decode_content = True return np.array(Image.open(response.raw))

def get_bg_removed_image(url,path,model,model_matting): if path: image= cv2.imread(path) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if url: image = download_image(url) height,width = image.shape[:2] if image.shape[2]>3: image = image[:,:] #actual_image = denormalize(image.squeeze()) #actual_image = cv2.cvtColor(actual_image,cv2.COLOR_BGR2RGB) image_resized = image.copy() image_resized = cv2.resize(image_resized,(224,224)) image_resized = np.expand_dims(image_resized, axis=0) print(image_resized.shape) pr_mask = model.predict(preprocess_input(image_resized).reshape(1,224,224,3)).round() alpha = pr_mask.reshape(224,224,1) alpha_seg = alpha.reshape(224,224) alpha_seg = cv2.resize(alpha_seg,(width,height))*255 alpha_seg = cv2.merge([alpha_seg,alpha_seg,alpha_seg]) alpha_seg[np.where(alpha_seg<127)] =0 alpha_seg[np.where(alpha_seg>=127)] =255 #plt.imshow(alpha) alpha_seg = alpha_seg[:,:,0] trimap = erode_dilate(alpha_seg) image_dict = generator_tensor_dict(image, trimap) pred, offset = single_inference(model_matting, image_dict) mat = pred bg = cv2.imread('plain-white-background.jpg') bg = cv2.resize(bg,(image.shape[1],image.shape[0])) pred = cv2.resize(pred,(width,height)).reshape(height, width, 1)/255 im_mat = pred*(image/255) + (1-pred)*(bg/255) visualize( image=image, predicted_mat = alpha_seg, mat = mat, trimap = trimap, image_mat = im_mat #using_mat=im_seg_mat, ) plt.show()