import os
import pickle

import torch
from faceseg.u2net import U2NET
from utils import landmark_processor
import cv2
import numpy as np


class FaceSeg:
    def __init__(self, gpu_id = 0):
        model = U2NET(in_ch=4, out_ch=1)
        weights = torch.load('weights/20210927_01.pth', map_location='cpu')
        model_dict = model.state_dict()
        pretrained_dict = {}
        for ix, (k, v) in enumerate(model_dict.items()):
            if k in weights and weights[k].data.shape == v.data.shape:
                pretrained_dict[k] = weights[k]
            else:
                print('ignore {}'.format(k))
        model_dict.update(pretrained_dict)
        model.load_state_dict(model_dict)
        print('update success')
        model.cuda(gpu_id)
        model.eval()
        self.model = model
        self.last_mask = None
        self.output_img_size = 320
        self.gpu_id = gpu_id

    def inference(self, frame, pt1k, video_mode=False):
        image_to_face_mat = landmark_processor.get_transform_mat_full_face(pt1k, self.output_img_size)
        face_image = cv2.warpAffine(frame, image_to_face_mat, (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)
        if face_image.dtype == np.uint8: face_image = face_image.astype(np.float32) / 255
        if video_mode and self.last_mask is not None:
            last_small_mask = cv2.warpAffine(self.last_mask, image_to_face_mat,
                                             (self.output_img_size, self.output_img_size), flags=cv2.INTER_LANCZOS4)[:,:,np.newaxis]
            input_img = np.concatenate([face_image, last_small_mask], axis=2)
        else:
            zero_mask = np.zeros((face_image.shape[1], face_image.shape[0], 1), dtype=np.float32)
            input_img = np.concatenate([face_image, zero_mask], axis=2)

        face_image_tensor = input_img.transpose((2, 0, 1))[np.newaxis]
        face_image_tensor = torch.from_numpy(face_image_tensor).cuda(self.gpu_id)

        mask = self.model.test(face_image_tensor)
        mask = mask[0].detach().cpu().numpy().transpose((1, 2, 0))

        origin_mask = cv2.warpAffine(mask, image_to_face_mat, (frame.shape[1], frame.shape[0]),
                                     flags=cv2.WARP_INVERSE_MAP|cv2.INTER_LANCZOS4)[:, :, np.newaxis]
        if video_mode: self.last_mask = origin_mask.copy()

        return origin_mask


if __name__ == '__main__':
    face_segmentor = FaceSeg(gpu_id=0)
    testdata_dir = "/mnt/DataDisk/my_projects/faceswap_hq/train_data/example"
    for picname in os.listdir(testdata_dir):
        img_path = os.path.join(testdata_dir, picname)
        pkl_path = img_path[:-4]+".pkl"
        if not picname.endswith(".jpg"):
            continue
        if not os.path.exists(pkl_path):
            continue
        img = cv2.imread(img_path)
        with open(pkl_path, "rb") as fp:
            info = pickle.load(fp)
        pt1k = info["human_pt1k"]

        face_seg_mask = face_segmentor.inference(img, pt1k, video_mode=False)
        cv2.imshow("face_seg_mask", face_seg_mask)
        cv2.imshow("img", img)
        cv2.waitKey()