def estimate_foreground_multilevel(image, alpha):
import pymatting
return pymatting.estimate_foreground_ml(
image,
alpha,
gradient_weight=0.1,
regularization=5e-3,
n_small_iterations=10,
# n_big_iterations=4 decreases error by a few percent, but will increase runtime
n_big_iterations=2,
)
def extend(self, fg_name):
fg_name = fg_name.strip()
alpha_path = join_first_contain(self.alpha_dirs, fg_name,
self.data_root)
fg_path = join_first_contain(self.fg_dirs, fg_name, self.data_root)
alpha_path = osp.join(self.data_root, alpha_path)
fg_path = osp.join(self.data_root, fg_path)
extended_path = re.sub('/fg/', '/fg_extended/', fg_path)
extended_path = extended_path.replace('jpg', 'png')
if not osp.exists(alpha_path):
raise FileNotFoundError(f'{alpha_path} does not exist!')
if not osp.exists(fg_path):
raise FileNotFoundError(f'{fg_path} does not exist!')
image = load_image(fg_path, 'RGB')
alpha = load_image(alpha_path, 'GRAY')
F = estimate_foreground_ml(image, alpha, return_background=False)
fg = Image.fromarray(np.uint8(F * 255))
fg.save(extended_path)
fix_png_file(osp.basename(extended_path), osp.dirname(extended_path))
data_info = dict()
data_info['alpha_path'] = alpha_path
data_info['fg_path'] = extended_path
return data_info
def main():
# --------- 1. get image path and name ---------
model_name = 'u2net'
model_dir = os.path.join(os.path.expanduser('~/data'), model_name,
model_name + '.pth')
# --------- 3. model define ---------
if (model_name == 'u2net'):
print("...load U2NET---173.6 MB")
net = U2NET(3, 1)
elif (model_name == 'u2netp'):
print("...load U2NEP---4.7 MB")
net = U2NETP(3, 1)
net.load_state_dict(torch.load(model_dir))
net = net.to(device)
net.eval()
# --------- 4. inference for each image ---------
input_paths = ["Girl_in_front_of_a_green_background.jpg"]
for path in input_paths:
image = Image.open(path).convert("RGB")
width, height = image.size
# resize to prevent out of memory error
scale = 500.0 / max(width, height)
width = int(scale * width)
height = int(scale * height)
image = image.resize((width, height), Image.BOX)
image0 = np.array(image) / 255.0
image = torch.from_numpy(image0.transpose(2, 0, 1)[np.newaxis,
...]).float()
image = image.to(device)
print("running u2net")
d1, d2, d3, d4, d5, d6, d7 = net(Variable(image))
pred = d1[:, 0, :, :]
del d1, d2, d3, d4, d5, d6, d7
print("converting")
# normalization
pred = normPRED(pred)
pred = pred.detach().cpu().numpy()
pred = pred[0, :, :]
is_foreground = pred > 0.95
is_background = pred < 0.05
from scipy.ndimage.morphology import binary_erosion
size = 11
structure = np.ones((size, size), dtype=np.int)
is_foreground = binary_erosion(is_foreground, structure=structure)
is_background = binary_erosion(is_background,
structure=structure,
border_value=1)
trimap = 0.5 * (np.ones_like(pred) + is_foreground - is_background)
from pymatting import estimate_foreground_ml, stack_images, save_image, estimate_alpha_lkm
print("alpha matting")
alpha = estimate_alpha_lkm(image0,
trimap,
laplacian_kwargs=dict(radius=10))
print("foreground estimation")
foreground = estimate_foreground_ml(image0, alpha)
print("saving")
cutout = stack_images(foreground, alpha)
save_image("cutout.png", cutout)
save_image("foreground.png", foreground)
save_image("timap.png", trimap)
save_image("pred.png", pred)
save_image("alpha.png", alpha)