Spaces:
Runtime error
Runtime error
File size: 2,526 Bytes
0b756df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import torch
import torchvision.transforms as transforms
import cv2
import numpy as np
from .model import BiSeNet
mask_regions = {
"Background":0,
"Skin":1,
"L-Eyebrow":2,
"R-Eyebrow":3,
"L-Eye":4,
"R-Eye":5,
"Eye-G":6,
"L-Ear":7,
"R-Ear":8,
"Ear-R":9,
"Nose":10,
"Mouth":11,
"U-Lip":12,
"L-Lip":13,
"Neck":14,
"Neck-L":15,
"Cloth":16,
"Hair":17,
"Hat":18
}
run_with_cuda = False
def init_parser(pth_path, use_cuda=False):
global run_with_cuda
run_with_cuda = use_cuda
n_classes = 19
net = BiSeNet(n_classes=n_classes)
if run_with_cuda:
net.cuda()
net.load_state_dict(torch.load(pth_path))
else:
net.load_state_dict(torch.load(pth_path, map_location=torch.device('cpu')))
net.eval()
return net
def image_to_parsing(img, net):
img = cv2.resize(img, (512, 512))
img = img[:,:,::-1]
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
img = transform(img.copy())
img = torch.unsqueeze(img, 0)
with torch.no_grad():
if run_with_cuda:
img = img.cuda()
out = net(img)[0]
parsing = out.squeeze(0).cpu().numpy().argmax(0)
return parsing
def get_mask(parsing, classes):
res = parsing == classes[0]
for val in classes[1:]:
res += parsing == val
return res
def swap_regions(source, target, net, includes=[1,2,3,4,5,10,11,12,13], excludes=[7,8], blur_size=25):
parsing = image_to_parsing(source, net)
if len(includes) == 0:
return source, np.zeros_like(source)
include_mask = get_mask(parsing, includes)
include_mask = np.repeat(np.expand_dims(include_mask.astype('float32'), axis=2), 3, 2)
if len(excludes) > 0:
exclude_mask = get_mask(parsing, excludes)
exclude_mask = np.repeat(np.expand_dims(exclude_mask.astype('float32'), axis=2), 3, 2)
include_mask -= exclude_mask
mask = 1 - cv2.GaussianBlur(include_mask.clip(0,1), (0, 0), blur_size)
result = (1 - mask) * cv2.resize(source, (512, 512)) + mask * cv2.resize(target, (512, 512))
result = cv2.resize(result.astype("float32"), (source.shape[1], source.shape[0]))
return result, mask.astype('float32')
def mask_regions_to_list(values):
out_ids = []
for value in values:
if value in mask_regions.keys():
out_ids.append(mask_regions.get(value))
return out_ids
|