目录
皮肤分割:
easyphoto/face_process_utils.py
class Face_Skin(object):
'''
Inputs:
image input image.
Outputs:
mask output mask.
'''
def __init__(self, model_path) -> None:
n_classes = 19
self.model = BiSeNet(n_classes=n_classes)
self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
self.model.eval()
self.cuda = torch.cuda.is_available()
if self.cuda:
self.model.cuda(1)
# transform for input image
self.trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# index => label
# 1:'skin', 2:'nose', 3:'eye_g', 4:'left_eye', 5:'right_eye', 6:'left_brow', 7:'right_brow', 8:'left_ear',
# 9:'right_ear', 10:'mouth', 11:'upper_lip', 12:'low_lip', 13:'hair'
def __call__(self, image, retinaface_detection, needs_index=[[12, 13]]):
# needs_index 12, 13 means seg the lip
with torch.no_grad():
total_mask = np.zeros_like(np.uint8(image))
# detect image
retinaface_boxes, _, _ = call_face_crop(retinaface_detection, image, 1.5, prefix="tmp")
retinaface_box = retinaface_boxes[0]
# sub_face for seg skin
sub_image = image.crop(retinaface_box)
image_h, image_w, c = np.shape(np.uint8(sub_image))
PIL_img = Image.fromarray(np.uint8(sub_image))
PIL_img = PIL_img.resize((512, 512), Image.BILINEAR)
torch_img = self.trans(PIL_img)
torch_img = torch.unsqueeze(torch_img, 0)
if self.cuda:
torch_img = torch_img.to(next(self.model.parameters()).device)
out = self.model(torch_img)[0]
model_mask = out.squeeze(0).cpu().numpy().argmax(0)
masks = []
for _needs_index in needs_index:
total_mask = np.zeros_like(np.uint8(image))
sub_mask = np.zeros_like(model_mask)
for index in _needs_index:
sub_mask += np.uint8(model_mask == index)
sub_mask = np.clip(sub_mask, 0, 1) * 255
sub_mask = np.tile(np.expand_dims(cv2.resize(np.uint8(sub_mask), (image_w, image_h)), -1), [1, 1, 3])
total_mask[retinaface_box[1]:retinaface_box[3], retinaface_box[0]:retinaface_box[2], :] = sub_mask
masks.append(Image.fromarray(np.uint8(total_mask)))
return masks
MALUNet 皮肤病分割
没有预训练,有数据
https://github.com/JCruan519/MALUNet
BAT:用于皮肤病变分割的边界感知Transformer - 知乎
BA-Transformer
论文标题:Boundary-aware Transformers for Skin Lesion Segmentation