u2net_portrait-keras.py

from keras.models import load_model
import cv2
import numpy as np
import sys


input = sys.argv[1]
output = sys.argv[2]

model = './u2net_portrait_keras.h5'


# load model
u2netp_keras = load_model(model, compile=False)

# load image
image = cv2.imread(input)

# normalize input image
input = cv2.resize(image, (512, 512), interpolation=cv2.INTER_CUBIC)

tmpImg = np.zeros((512, 512, 3))
input = input/np.max(input)

tmpImg[:,:,0] = (input[:,:,2]-0.406)/0.225
tmpImg[:,:,1] = (input[:,:,1]-0.456)/0.224
tmpImg[:,:,2] = (input[:,:,0]-0.485)/0.229

# convert BGR to RGB
tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = tmpImg[np.newaxis,:,:,:]

# predict
d1,d2,d3,d4,d5,d6,d7 = u2netp_keras.predict(tmpImg)
pred = np.array(1.0 - d1[:,0,:,:])[0]


# normalize
ma = np.max(pred)
mi = np.min(pred)
pred = (pred-mi)/(ma-mi)

pred = pred.squeeze()

pred = (pred*255).astype(np.uint8)

out = cv2.resize(pred, image.shape[1::-1], interpolation=cv2.INTER_CUBIC)

# save image
cv2.imwrite(output, out)