皮肤病分割学习笔记

AI算法网奇

已于 2024-01-07 01:55:09 修改

阅读量460

点赞数 7

分类专栏：深度学习基础文章标签：学习笔记

于 2023-12-26 12:37:40 首次发布

本文链接：https://blog.csdn.net/jacke121/article/details/135218969

版权

深度学习基础专栏收录该内容

166 篇文章 17 订阅

订阅专栏

皮肤分割：

MALUNet 皮肤病分割

BA-Transformer

皮肤分割：

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