コード例 #1
0
def crop_progress(image):

    rects = face_detector(image, 1)

    if len(rects) == 0:
        return

    for rect in rects:
        offset = 0
        top = rect.top()
        bottom = rect.bottom() - 0
        left = rect.left() + offset
        right = rect.right() - offset

        faceBoxRectangleS = dlib.rectangle(left=left,
                                           top=top,
                                           right=right,
                                           bottom=bottom)

        pts = face_regressor(image, faceBoxRectangleS).parts()
        pts = np.array([[pt.x, pt.y] for pt in pts]).T
        roi_box = parse_roi_box_from_landmark(pts)

        cropped_image = crop_img(image, roi_box)

        cropped_image = cv2.resize(cropped_image,
                                   dsize=(STD_SIZE, STD_SIZE),
                                   interpolation=cv2.INTER_LINEAR)

        return cropped_image
コード例 #2
0
def crop_process(image, filename, folder_path, save_path, size=224):

    if image is None:
        print("Could not read input image")
        return False


    start = time.time()

    # apply face detection (cnn)
    faces_cnn = cnn_face_detector(image, 1)

    end = time.time()
    print("CNN : ", format(end - start, '.2f'))
    print(len(faces_cnn))
    if len(faces_cnn) == 0:
        print("no face found")
        return False    

    print(len(faces_cnn))
    # loop over detected faces
    face =  faces_cnn[0]
    left = face.rect.left()
    top = face.rect.top()
    right = face.rect.right()
    bottom = face.rect.bottom()

    faceBoxRectangleS =  dlib.rectangle(left=left,top=top,right=right, bottom=bottom)

    # - use landmark for cropping
    pts = face_regressor(image, faceBoxRectangleS).parts()
    pts = np.array([[pt.x, pt.y] for pt in pts]).T
    roi_box = parse_roi_box_from_landmark(pts)

    height, width, _ =  image.shape

    ########## left
    if roi_box[0] < 0:
        roi_box[0] = 0
    ########## right        
    if roi_box[1] < 0:
        roi_box[1] = 0
    ########## width
    if roi_box[2] > width:
        roi_box[2] = width
    ########## height
    if roi_box[3] > height:
        roi_box[3] = height


    cropped_image = crop_img(image, roi_box)

    # forward: one step
    cropped_image = cv2.resize(cropped_image, dsize=(size, size), interpolation=cv2.INTER_LINEAR)
    save_img(cropped_image, filename, save_path)
    print('saved')        
    return cropped_image
コード例 #3
0
def crop_process(image, filename, folder_path, save_path, size=224):

    if image is None:
        print("Could not read input image")
        return False

    rects = face_detector(image, 1)

    if len(rects) == 0:
        suffix = get_suffix(filename)  #suffix = '.jpg'
        image_name = filename.replace(folder_path + '\\', '')
        with open('./notfound.txt', 'a+') as f:
            f.write(image_name + '\n')
        print("face not found")
        return False

    rect = rects[0]
    offset = 0
    top = rect.top()
    bottom = rect.bottom() - 0
    left = rect.left() + offset
    right = rect.right() - offset

    faceBoxRectangleS = dlib.rectangle(left=left,
                                       top=top,
                                       right=right,
                                       bottom=bottom)

    # - use landmark for cropping
    pts = face_regressor(image, faceBoxRectangleS).parts()
    pts = np.array([[pt.x, pt.y] for pt in pts]).T
    roi_box = parse_roi_box_from_landmark(pts)

    height, width, _ = image.shape

    ########## left
    if roi_box[0] < 0:
        roi_box[0] = 0
    ########## right
    if roi_box[1] < 0:
        roi_box[1] = 0
    ########## width
    if roi_box[2] > width:
        roi_box[2] = width
    ########## height
    if roi_box[3] > height:
        roi_box[3] = height

    cropped_image = crop_img(image, roi_box)

    # forward: one step
    cropped_image = cv2.resize(cropped_image,
                               dsize=(size, size),
                               interpolation=cv2.INTER_LINEAR)
    save_img(cropped_image, filename, save_path, folder_path)
    # print('saved')
    return cropped_image
コード例 #4
0
def get_landmark(path):
    device = torch.device("cuda")
    cap = cv2.VideoCapture(path)

    cudnn.benchmark = True
    smooth_flag = False

    ret, frame = cap.read()
    frame = resize_frame(frame, cfg.resize_size)
    height, width, _ = frame.shape

    face_detect = detector.SFDDetector(device, cfg.rect_model_path)
    face_landmark = FAN_landmark.FANLandmarks(device, cfg.landmark_model_path,
                                              cfg.detect_type)
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    model = mobilenet_v1.mobilenet(cfg.checkpoint_fp, num_classes=240)

    landmark_list = []
    frame_list = []

    while ret:
        frame = resize_frame(frame, cfg.resize_size)
        bbox = face_detect.extract(frame, cfg.thresh)
        if len(bbox) > 0:
            frame, landmarks = face_landmark.extract([frame, bbox])
            roi = roi_box_from_landmark(landmarks[0])
            img = crop_img(frame, roi)
            img = cv2.resize(img,
                             dsize=(cfg.STD_SIZE, cfg.STD_SIZE),
                             interpolation=cv2.INTER_LINEAR)
            input_data = transform(img).unsqueeze(0)
            with torch.no_grad():
                param = model(input_data.cuda())
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)
            pts68 = predict_68pts(param, roi, img.shape[0])
            if smooth_flag:
                show_frame, landmark_3d = get_smooth_data(
                    frame, pts68.T, frame_list, landmark_list, cfg.list_size)
            else:
                show_frame = frame
                landmark_3d = pts68.T
            for x, y in landmark_3d[:, :2]:
                cv2.circle(show_frame, (int(x), int(y)), 1, (255, 255, 0))
            cv2.imshow('3ddfa', show_frame)
            cv2.waitKey(1)
        else:
            print("cannot find face")
            continue

        ret, frame = cap.read()
    cap.release()
    cv2.destroyAllWindows()
コード例 #5
0
    def __getitem__(self, idx):

        filename = self.img_list[idx]
        ori_img = cv2.imread(filename)

        landmark = self.landmarks[idx]

        # preprocess img
        roi_box = parse_roi_box_from_landmark(landmark.T)
        img = crop_img(ori_img, roi_box)
        img = cv2.resize(img, dsize=(self.std_size, self.std_size), interpolation=cv2.INTER_LINEAR)
        if self.transform is not None:
            img = self.transform(img)

        return img, ori_img, filename, np.array(roi_box)
コード例 #6
0
def get_inputs(files: List[str]) -> Tuple[torch.Tensor, List[np.ndarray]]:
    """Get Inputs

    Parameters
    ----------
    files: List[ str ]
           List of frames to process.

    Returns
    -------
    inputs  : torch.Tensor
              Frame cropped to the face.
    imgs_roi: List[ np.ndarray ]
              Frame face rectangles ( region of interest ).
    """
    face_reg, face_det = get_face_model()
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    imgs_ori = [cv2.imread(file) for file in files]
    imgs_rect = [face_det(img)[0] for img in imgs_ori]
    imgs_pts = [
        face_reg(imgs_ori[i], imgs_rect[i]).parts()
        for i in range(len(imgs_ori))
    ]
    imgs_pts = [
        np.array([[pt.x, pt.y] for pt in imgs_pts[i]]).T
        for i in range(len(imgs_ori))
    ]
    imgs_roi = [parse_roi_box_from_landmark(pts) for pts in imgs_pts]
    imgs_cropped = [
        crop_img(imgs_ori[i], imgs_roi[i]) for i in range(len(imgs_ori))
    ]
    imgs_new = [
        cv2.resize(img, dsize=(120, 120), interpolation=cv2.INTER_LINEAR)
        for img in imgs_cropped
    ]

    inputs = [transform(img).unsqueeze(0) for img in imgs_new]
    inputs = torch.cat(inputs, axis=0)

    return inputs, imgs_roi
コード例 #7
0
def detect_pose(img_ori, target, model):
    STD_SIZE = 120
    pose_list = []
    for rect in target:
        roi_box = rect[0:4]
        # print(roi_box)
        try:
            img = crop_img(img_ori, roi_box)
        except Exception as e:
            print(e)
            continue

        shape_img = img.shape
        if shape_img[0] < 100 or shape_img[1] < 100:
            continue
        if len(img) == 0:
            return pose_list
        try:
            img = cv2.resize(img,
                             dsize=(STD_SIZE, STD_SIZE),
                             interpolation=cv2.INTER_LINEAR)
        except Exception as e:
            print(e)
            continue
        transform = transforms.Compose(
            [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
        input = transform(img).unsqueeze(0)
        with torch.no_grad():
            param = model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        pts68 = predict_68pts(param, roi_box)
        P, pose, pose_angle = parse_pose(param)

        # if shape_img[0]>100 and shape_img[1]>100 and np.abs(pose_angle[0]) <= 60 and np.abs(pose_angle[1]) <= 60 and np.abs(pose_angle[2])<= 60:
        #   print(shape_img)
        #   plt.imshow(img)
        #   plt.show()
        pose_list.append(pose_angle)
    return pose_list
コード例 #8
0
ファイル: draw_face.py プロジェクト: wuxiaolianggit/R3FA
def getFaceTextureCoords(textImag):
    checkpoint_fp = 'models/MobDenseNet.pth.tar'
    arch = 'densemobilenetv4_19'
    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(MobDenseNet, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)
    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    cudnn.benchmark = True
    model = model.cuda()
    model.eval()
    face_detector = dlib.get_frontal_face_detector()
    # 3. forward
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    img_ori = textImag
    rects = face_detector(img_ori, 1)
    for rect in rects:
        # - use detected face bbox
        bbox = [rect.left(), rect.top(), rect.right(), rect.bottom()]
        roi_box = parse_roi_box_from_bbox(bbox)
        img = crop_img(img_ori, roi_box)

        # forward: one step
        img = cv2.resize(img,
                         dsize=(STD_SIZE, STD_SIZE),
                         interpolation=cv2.INTER_LINEAR)
        input = transform(img).unsqueeze(0)
        with torch.no_grad():
            input = input.cuda()
            param = model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        # 68 pts
        pts68 = predict_68pts(param, roi_box)
        return pts68[[0, 1], :]
コード例 #9
0
def crop_progress(image):

    rects = face_detector(image, 1)

    if len(rects) == 0: ## frontal face detector가 얼굴을 찾았는가?
        print('using cnn')
        rects = cnn_face_detector(image, 1)
    
        if len(rects) == 0: ### mmod가 얼굴을 찾았는가?
            print('could not find face.')
            return

        rect = rects[0].rect ###rect를 mmod가 찾은 것으로 진행한다
        
    else:
        rect = rects[0] #### frontal face detector가 찾은 얼굴로 진행한다

    top = rect.top()
    bottom = rect.bottom()
    left = rect.left()
    right = rect.right()


    faceBoxRectangleS =  dlib.rectangle(left=left,top=top,right=right, bottom=bottom)

    # - use landmark for cropping
    pts = face_regressor(image, faceBoxRectangleS).parts()
    pts = np.array([[pt.x, pt.y] for pt in pts]).T
    roi_box = parse_roi_box_from_landmark(pts)


    cropped_image = crop_img(image, roi_box)

    # forward: one step
    cropped_image = cv2.resize(cropped_image, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
    
    return cropped_image
コード例 #10
0
ファイル: R3FR_predictor.py プロジェクト: wuxiaolianggit/R3FA
def r3fa_landmark(image,rect,model,predictor,imgScale):

    transform=transforms.Compose([ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    Ps=[]  # Camera matrix collection
    pts_res=[]
    faceRectangle=rectangle(int(rect.left() / imgScale), int(rect.top() / imgScale), int(rect.right() / imgScale),
                            int(rect.bottom() / imgScale))

    # - use landmark for cropping
    pts=predictor(image, faceRectangle).parts()
    pts=np.array([[pt.x, pt.y] for pt in pts]).T
    roi_box=parse_roi_box_from_landmark(pts)
    #     # - use detected face bbox
    # bbox=[rects.left(), rects.top(), rects.right(), rects.bottom()]
    # roi_box=parse_roi_box_from_bbox(bbox)
    img=crop_img(image, roi_box)

    # forward: one step
    img=cv2.resize(img, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
    input=transform(img).unsqueeze(0)
    with torch.no_grad():
        input=input.cuda()
        param=model(input)
        param=param.squeeze().cpu().numpy().flatten().astype(np.float32)
    # 68 pts
    pts68=predict_68pts(param, roi_box)
    pts_res.append(pts68)
    P, pose=parse_pose(param)
    Ps.append(P)
    # # 绘制landmark
    # for indx in range(68):
    #     pos=(pts68[0, indx], pts68[1, indx])
    #     # pts.append(pos)
    #     cv2.circle(img, pos, 3, color=(255, 255, 255), thickness=-1)
    # img_ori=plot_pose_box(img, Ps, pts_res)
    return Ps,pts68[[0,1],:]
コード例 #11
0
def crop_process(image, filename, folder_path, save_path, size=224):

    if image is None:
        print("Could not read input image")
        return False

    h, w, _ = image.shape
    blob = cv2.dnn.blobFromImage(image, 1.0, (300, 300), [104, 117, 123],
                                 False, False)
    net.setInput(blob)

    # inference, find faces
    detections = net.forward()
    # postprocessing

    faceBoxRectangleS = None
    accuracy = 0.
    for i in range(detections.shape[2]):
        confidence = detections[0, 0, i, 2]
        if confidence > conf_threshold:
            x1 = int(detections[0, 0, i, 3] * w)
            y1 = int(detections[0, 0, i, 4] * h)
            x2 = int(detections[0, 0, i, 5] * w)
            y2 = int(detections[0, 0, i, 6] * h)
            #, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)

            accuracy = confidence * 100
            faceBoxRectangleS = dlib.rectangle(left=x1,
                                               top=y1,
                                               right=x2,
                                               bottom=y2)

    if faceBoxRectangleS == None:
        suffix = get_suffix(filename)  #suffix = '.jpg'
        image_name = filename.replace(folder_path + '\\', '')
        with open('./notfound2.txt', 'a+') as f:
            f.write(image_name + '\n')
        return False

    # - use landmark for cropping
    pts = face_regressor(image, faceBoxRectangleS).parts()
    pts = np.array([[pt.x, pt.y] for pt in pts]).T
    roi_box = parse_roi_box_from_landmark(pts)

    height, width, _ = image.shape

    ########## left
    if roi_box[0] < 0:
        roi_box[0] = 0
    ########## right
    if roi_box[1] < 0:
        roi_box[1] = 0
    ########## width
    if roi_box[2] > width:
        roi_box[2] = width
    ########## height
    if roi_box[3] > height:
        roi_box[3] = height

    cropped_image = crop_img(image, roi_box)

    # forward: one step
    cropped_image = cv2.resize(cropped_image,
                               dsize=(size, size),
                               interpolation=cv2.INTER_LINEAR)
    save_img(cropped_image, filename, save_path, folder_path)
    # print('saved')
    return cropped_image
コード例 #12
0
ファイル: PoseQueue.py プロジェクト: accubits/AI-Research
def run(image):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    model.eval()

    # 2. load dlib model for face detection and landmark used for face cropping
    dlib_landmark_model = 'models/shape_predictor_68_face_landmarks.dat'
    face_regressor = dlib.shape_predictor(dlib_landmark_model)
    face_detector = dlib.get_frontal_face_detector()

    # 3. forward
    tri = sio.loadmat('tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    img_ori = cv2.imread(image)
    # img_ori = img_ori[:, :, [2, 1, 0]]
    rects = face_detector(img_ori, 1)

    # if len(rects) == 0:
    #     rects = dlib.rectangles()
    #     rect_fp = img_fp + '.bbox'
    #     lines = open(rect_fp).read().strip().split('\n')[1:]
    #     for l in lines:
    #         l, r, t, b = [int(_) for _ in l.split(' ')[1:]]
    #         rect = dlib.rectangle(l, r, t, b)
    #         rects.append(rect)

    pts_res = []
    Ps = []  # Camera matrix collection
    poses = []  # pose collection, [todo: validate it]
    vertices_lst = []  # store multiple face vertices
    ind = 0
    suffix = get_suffix(image)
    for rect in rects:
        # whether use dlib landmark to crop image, if not, use only face bbox to calc roi bbox for cropping
        # - use landmark for cropping
        pts = face_regressor(img_ori, rect).parts()
        pts = np.array([[pt.x, pt.y] for pt in pts]).T
        roi_box = parse_roi_box_from_landmark(pts)

        img = crop_img(img_ori, roi_box)

        # forward: one step
        img = cv2.resize(img,
                         dsize=(STD_SIZE, STD_SIZE),
                         interpolation=cv2.INTER_LINEAR)
        input = transform(img).unsqueeze(0)
        with torch.no_grad():
            param = model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        # 68 pts
        pts68 = predict_68pts(param, roi_box)

        # two-step for more accurate bbox to crop face

        pts_res.append(pts68)
        P, pose = parse_pose(param)
        Ps.append(P)
        poses.append(pose)

        vertices = predict_dense(param, roi_box)
        vertices_lst.append(vertices)

        # dense face 3d vertices
        wfp = './output/obj_{}.obj'.format(image.split('/')[-1])
        colors = get_colors(img_ori, vertices)
        write_obj_with_colors(wfp, vertices, tri, colors)
        print('Dump obj with sampled texture to {}'.format(wfp))
        ind += 1
コード例 #13
0
def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose([ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    # 2. parse images list 
    with open(args.img_list) as f:
        img_list = [x.strip() for x in f.readlines()]
    landmark_list = []

    alignment_model = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)
    if not os.path.exists(args.save_lmk_dir):
        os.mkdir(args.save_lmk_dir)

    for img_idx, img_fp in enumerate(tqdm(img_list)):
        img_ori = cv2.imread(os.path.join(args.img_prefix, img_fp))

        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = get_suffix(img_fp)

        # face alignment model use RGB as input, result is a tuple with landmarks and boxes, cv2读取图片的书序是 BGR
        preds = alignment_model.get_landmarks(img_ori[:, :, ::-1])
        pts_2d_68 = preds[0]
        pts_2d_5 = get_5lmk_from_68lmk(pts_2d_68)
        landmark_list.append(pts_2d_5)
        roi_box = parse_roi_box_from_landmark(pts_2d_68.T)      # 根据68个关键点,确定roi区域

        img = crop_img(img_ori, roi_box)
        # import pdb; pdb.set_trace()

        # forward: one step
        img = cv2.resize(img, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
        input = transform(img).unsqueeze(0)
        with torch.no_grad():
            if args.mode == 'gpu':
                input = input.cuda()
            param = model(input)                        # 人脸的RT矩阵,形状和表情的系数
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        # 68 pts
        pts68 = predict_68pts(param, roi_box)                   # 此时pts68是一个68*3的顶点坐标列表

        # two-step for more accurate bbox to crop face
        if args.bbox_init == 'two':
            roi_box = parse_roi_box_from_landmark(pts68)
            img_step2 = crop_img(img_ori, roi_box)
            img_step2 = cv2.resize(img_step2, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
            input = transform(img_step2).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

            pts68 = predict_68pts(param, roi_box)

        pts_res.append(pts68)
        P, pose = parse_pose(param)
        Ps.append(P)
        poses.append(pose)

        # dense face 3d vertices
        vertices = predict_dense(param, roi_box)

        if args.dump_2d_img:                # 人脸区域的2d图像
            wfp_2d_img = os.path.join(args.save_dir, os.path.basename(img_fp))
            colors = get_colors(img_ori, vertices)
            # aligned_param = get_aligned_param(param)
            # vertices_aligned = predict_dense(aligned_param, roi_box)
            # h, w, c = 120, 120, 3
            h, w, c = img_ori.shape
            img_2d = crender_colors(vertices.T, (tri - 1).T, colors[:, ::-1], h, w)
            cv2.imwrite(wfp_2d_img, img_2d[:, :, ::-1])
        if args.dump_param:
            split = img_fp.split('/')
            save_name = os.path.join(args.save_dir, '{}.txt'.format(os.path.splitext(split[-1])[0]))
            this_param = param * param_std + param_mean
            this_param = np.concatenate((this_param, roi_box))
            this_param.tofile(save_name, sep=' ')
    if args.dump_lmk:                   # 存储
        save_path = os.path.join(args.save_lmk_dir, 'realign_lmk')
        with open(save_path, 'w') as f:
            for idx, (fname, land) in enumerate(zip(img_list, landmark_list)):
                # f.write('{} {} {} {}')
                land = land.astype(np.int)
                land_str = ' '.join([str(x) for x in land])
                msg = f'{fname} {idx} {land_str}\n'
                f.write(msg)
コード例 #14
0
ファイル: inference.py プロジェクト: hchouchen/faceswap
def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    # 2. parse images list
    img_list = listdir(args.img_prefix)

    alignment_model = face_alignment.FaceAlignment(
        face_alignment.LandmarksType._2D, flip_input=False, device=args.mode)

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    for img_idx, img_fp in enumerate(tqdm(img_list)):
        img_ori = cv2.imread(os.path.join(args.img_prefix, img_fp))

        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = get_suffix(img_fp)

        # face alignment model use RGB as input, result is a tuple with landmarks and boxes
        preds = alignment_model.get_landmarks(img_ori[:, :, ::-1])
        try:
            pts_2d_68 = preds[0]
        except:
            continue
        roi_box = parse_roi_box_from_landmark(pts_2d_68.T)

        img = crop_img(img_ori, roi_box)
        # import pdb; pdb.set_trace()

        # forward: one step
        img = cv2.resize(img,
                         dsize=(STD_SIZE, STD_SIZE),
                         interpolation=cv2.INTER_LINEAR)
        input = transform(img).unsqueeze(0)
        with torch.no_grad():
            if args.mode == 'gpu':
                input = input.cuda()
            param = model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        # 68 pts
        pts68 = predict_68pts(param, roi_box)

        # two-step for more accurate bbox to crop face
        if args.bbox_init == 'two':
            roi_box = parse_roi_box_from_landmark(pts68)
            img_step2 = crop_img(img_ori, roi_box)
            img_step2 = cv2.resize(img_step2,
                                   dsize=(STD_SIZE, STD_SIZE),
                                   interpolation=cv2.INTER_LINEAR)
            input = transform(img_step2).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)

            pts68 = predict_68pts(param, roi_box)

        pts_res.append(pts68)
        P, pose = parse_pose(param)
        Ps.append(P)
        poses.append(pose)

        # dense face 3d vertices
        vertices = predict_dense(param, roi_box)

        if args.dump_2d_img:
            wfp_2d_img = os.path.join(args.save_dir, os.path.basename(img_fp))
            colors = get_colors(img_ori, vertices)
            # aligned_param = get_aligned_param(param)
            # vertices_aligned = predict_dense(aligned_param, roi_box)
            # h, w, c = 120, 120, 3
            h, w, c = img_ori.shape
            img_2d = crender_colors(vertices.T, (tri - 1).T, colors[:, ::-1],
                                    h, w)
            cv2.imwrite(wfp_2d_img, img_2d[:, :, ::-1])
        del img_ori
        del img
        del img_2d
コード例 #15
0
def test_video(args):
    start_time = time.time()
    x = 1  # displays the frame rate every 1 second
    counter = 0
    # 1. load pre-tained model
    # checkpoint_fp='models/phase1_wpdc_vdc_v2.pth.tar'
    # arch='mobilenet_1'
    checkpoint_fp = 'models/MobDenseNet.pth.tar'
    arch = 'mobdensenet_v1'
    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(MobDenseNet, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)
    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()
    # 2. load dlib model for face detection and landmark used for face cropping
    if args.dlib_landmark:
        dlib_landmark_model = 'models/shape_predictor_68_face_landmarks.dat'
        face_regressor = dlib.shape_predictor(dlib_landmark_model)
    if args.dlib_bbox:
        face_detector = dlib.get_frontal_face_detector()
    # 3. forward
    tri = sio.loadmat('visualize/tri.mat')['tri'] - 1
    tri_pts68 = sio.loadmat('visualize/pats68_tri.mat')['tri']
    textureImg = cv2.imread(image_name)
    cameraImg = cap.read()[1]
    # textureCoords=df.getFaceTextureCoords(textureImg)
    # drawface=Drawing3DFace.Draw3DFace(cameraImg,textureImg,textureCoords,tri_pts68.T)
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    while True:
        # get a frame
        img_ori = cap.read()[1]
        imgScale = 1
        scaledImg = img_ori
        if max(img_ori.shape) > maxImgSizeForDetection:
            imgScale = maxImgSizeForDetection / float(max(img_ori.shape))
            scaledImg = cv2.resize(img_ori, (int(img_ori.shape[1] * imgScale),
                                             int(img_ori.shape[0] * imgScale)))
        rects = face_detector(scaledImg, 1)

        Ps = []  # Camera matrix collection
        poses = []  # pose collection
        pts_res = []
        # suffix=get_suffix(img_ori)
        for rect in rects:
            if args.dlib_landmark:
                faceRectangle = rectangle(int(rect.left() / imgScale),
                                          int(rect.top() / imgScale),
                                          int(rect.right() / imgScale),
                                          int(rect.bottom() / imgScale))

                # - use landmark for cropping
                pts = face_regressor(img_ori, faceRectangle).parts()
                pts = np.array([[pt.x, pt.y] for pt in pts]).T
                roi_box = parse_roi_box_from_landmark(pts)
            else:
                bbox = [
                    int(rect.left() / imgScale),
                    int(rect.top() / imgScale),
                    int(rect.right() / imgScale),
                    int(rect.bottom() / imgScale)
                ]
                roi_box = parse_roi_box_from_bbox(bbox)
        img = crop_img(img_ori, roi_box)
        # forward: one step
        img = cv2.resize(img,
                         dsize=(STD_SIZE, STD_SIZE),
                         interpolation=cv2.INTER_LINEAR)
        input = transform(img).unsqueeze(0)
        with torch.no_grad():
            if args.mode == 'gpu':
                input = input.cuda()
            param = model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)
        # 68 pts
        pts68 = predict_68pts(param, roi_box)
        # df.triDelaunay(pts68)
        densePts = predict_dense(param, roi_box)
        P, pose = parse_pose(param)
        Ps.append(P)
        poses.append(pose)
        # two-step for more accurate bbox to crop face
        if args.bbox_init == 'two':
            roi_box = parse_roi_box_from_landmark(pts68)
            img_step2 = crop_img(img_ori, roi_box)
            img_step2 = cv2.resize(img_step2,
                                   dsize=(STD_SIZE, STD_SIZE),
                                   interpolation=cv2.INTER_LINEAR)
            input = transform(img_step2).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)
            pts68 = predict_68pts(param, roi_box)
        pts_res.append(pts68)
        pts = []
        #draw landmark
        for indx in range(68):
            pos = (pts68[0, indx], pts68[1, indx])
            pts.append(pos)
            cv2.circle(img_ori, pos, 3, color=(255, 255, 255), thickness=-1)
        ##draw pose box
        if args.dump_pose:
            img_ori = plot_pose_box(img_ori, Ps, pts_res)
        #draw face mesh
        if args.dump_2D_face_mesh:
            img_ori = df.drawMesh(img_ori, densePts.T, tri.T)
        if args.dump_3D_face_mesh:
            pass
            # img=drawface.render(pts68)
        cv2.imshow("faceDetector", img_ori)
        counter += 1
        if (time.time() - start_time) > x:
            print("FPS: ", counter / (time.time() - start_time))
            counter = 0
            start_time = time.time()
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
コード例 #16
0
    def process_one_frame(frame):
        if frame is None: print("No frame, check your camera")
        img_ori = frame
        img_fp = "camera.png"  # 随便给一个filename 反正之后用不到
        rects = face_detector(img_ori, 1)

        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = get_suffix(img_fp)

        for i in range(len(rects)):
            if i != 0: break  # 这里只检测第一个脸
            rect = rects[i]
            # whether use dlib landmark to crop image, if not, use only face bbox to calc roi bbox for cropping
            # - use landmark for cropping
            pts = face_regressor(img_ori, rect).parts()
            pts = np.array([[pt.x, pt.y] for pt in pts]).T
            roi_box = parse_roi_box_from_landmark(pts)
            img = crop_img(img_ori, roi_box)
            # forward: one step
            img = cv2.resize(img,
                             dsize=(STD_SIZE, STD_SIZE),
                             interpolation=cv2.INTER_LINEAR)
            input = transform(img).unsqueeze(0)
            with torch.no_grad():
                if use_gpu:
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)
            # 68 pts
            pts68 = predict_68pts(param, roi_box)
            # two-step for more accurate bbox to crop face
            if use_two_step_bbox_init:
                roi_box = parse_roi_box_from_landmark(pts68)
                img_step2 = crop_img(img_ori, roi_box)
                img_step2 = cv2.resize(img_step2,
                                       dsize=(STD_SIZE, STD_SIZE),
                                       interpolation=cv2.INTER_LINEAR)
                input = transform(img_step2).unsqueeze(0)
                with torch.no_grad():
                    if use_gpu:
                        input = input.cuda()
                    param = model(input)
                    param = param.squeeze().cpu().numpy().flatten().astype(
                        np.float32)
                pts68 = predict_68pts(param, roi_box)
            pts_res.append(pts68)
            P, pose = parse_pose(param)
            Ps.append(P)
            poses.append(pose)
            # dense face 3d vertices
            vertices = predict_dense(param, roi_box)
            vertices_lst.append(vertices)
            if is_dump_to_ply:
                dump_to_ply(
                    vertices, tri,
                    '{}_{}.ply'.format(img_fp.replace(suffix, ''), ind))
            if is_dump_vertex:
                dump_vertex(
                    vertices, '{}_{}.mat'.format(img_fp.replace(suffix, ''),
                                                 ind))
            if is_dump_pts:
                wfp = '{}_{}.txt'.format(img_fp.replace(suffix, ''), ind)
                np.savetxt(wfp, pts68, fmt='%.3f')
                print('Save 68 3d landmarks to {}'.format(wfp))
            if is_dump_roi_box:
                wfp = '{}_{}.roibox'.format(img_fp.replace(suffix, ''), ind)
                np.savetxt(wfp, roi_box, fmt='%.3f')
                print('Save roi box to {}'.format(wfp))
            if is_dump_paf:
                wfp_paf = '{}_{}_paf.jpg'.format(img_fp.replace(suffix, ''),
                                                 ind)
                wfp_crop = '{}_{}_crop.jpg'.format(img_fp.replace(suffix, ''),
                                                   ind)
                paf_feature = gen_img_paf(img_crop=img,
                                          param=param,
                                          kernel_size=paf_size)
                cv2.imwrite(wfp_paf, paf_feature)
                cv2.imwrite(wfp_crop, img)
                print('Dump to {} and {}'.format(wfp_crop, wfp_paf))
            if is_dump_obj:
                wfp = '{}_{}.obj'.format(img_fp.replace(suffix, ''), ind)
                colors = get_colors(img_ori, vertices)
                write_obj_with_colors(wfp, vertices, tri, colors)
                print('Dump obj with sampled texture to {}'.format(wfp))
            ind += 1
        if is_dump_pose:
            # P, pose = parse_pose(param)  # Camera matrix (without scale), and pose (yaw, pitch, roll, to verify)
            img_pose = plot_pose_box(img_ori, Ps, pts_res)
            wfp = img_fp.replace(suffix, '_pose.jpg')
            cv2.imwrite(wfp, img_pose)
            print('Dump to {}'.format(wfp))
        if is_dump_depth:
            wfp = img_fp.replace(suffix, '_depth.png')
            # depths_img = get_depths_image(img_ori, vertices_lst, tri-1)  # python version
            depths_img = cget_depths_image(img_ori, vertices_lst,
                                           tri - 1)  # cython version
            cv2.imwrite(wfp, depths_img)
            print('Dump to {}'.format(wfp))
        if is_dump_pncc:
            wfp = img_fp.replace(suffix, '_pncc.png')
            pncc_feature = cpncc(img_ori, vertices_lst,
                                 tri - 1)  # cython version
            cv2.imwrite(
                wfp,
                pncc_feature[:, :, ::-1])  # cv2.imwrite will swap RGB -> BGR
            print('Dump to {}'.format(wfp))
        if is_dump_res:
            draw_landmarks(img_ori,
                           pts_res,
                           wfp=img_fp.replace(suffix, '_3DDFA.jpg'),
                           show_flg=show_landmarks_fig)

        # 下面这一句将原始图像关掉
        img_ori = np.ones_like(img_ori) * 255

        img_pose = plot_pose_box(img_ori, Ps, pts_res)
        draw_landmarks_opencv(img_pose, pts_res)
        cv2.imshow("pose", img_pose)
コード例 #17
0
ファイル: uv_map.py プロジェクト: huyentls1114/3ddfa-uv-gan
def main(args):
    # 1. load pre-trained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    # 2. load pre-trained model uv-gan
    if args.uvgan:
        if args.checkpoint_uv_gan == "":
            print("Specify the path to checkpoint uv_gan")
            exit()
        uvgan = infer_uv_gan.UV_GAN(args.checkpoint_uv_gan)

    # 3. load dlib model for face detection and landmark used for face cropping
    if args.dlib_landmark:
        dlib_landmark_model = 'models/shape_predictor_68_face_landmarks.dat'
        face_regressor = dlib.shape_predictor(dlib_landmark_model)
    if args.dlib_bbox:
        face_detector = dlib.get_frontal_face_detector()

    # 4. forward
    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    for img_fp in args.files:
        img_ori = cv2.imread(img_fp)
        if args.dlib_bbox:
            rects = face_detector(img_ori, 1)
        else:
            rects = []

        if len(rects) == 0:
            rects = dlib.rectangles()
            rect_fp = img_fp + '.bbox'
            lines = open(rect_fp).read().strip().split('\n')[1:]
            for l in lines:
                l, r, t, b = [int(_) for _ in l.split(' ')[1:]]
                rect = dlib.rectangle(l, r, t, b)
                rects.append(rect)

        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = get_suffix(img_fp)
        for rect in rects:
            # whether use dlib landmark to crop image, if not, use only face bbox to calc roi bbox for cropping
            if args.dlib_landmark:
                # - use landmark for cropping
                pts = face_regressor(img_ori, rect).parts()
                pts = np.array([[pt.x, pt.y] for pt in pts]).T
                roi_box = parse_roi_box_from_landmark(pts)
            else:
                # - use detected face bbox
                bbox = [rect.left(), rect.top(), rect.right(), rect.bottom()]
                roi_box = parse_roi_box_from_bbox(bbox)

            img = crop_img(img_ori, roi_box)

            # forward: one step
            img = cv2.resize(img,
                             dsize=(STD_SIZE, STD_SIZE),
                             interpolation=cv2.INTER_LINEAR)
            input = transform(img).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)

            # 68 pts
            pts68 = predict_68pts(param, roi_box)

            # two-step for more accurate bbox to crop face
            if args.bbox_init == 'two':
                roi_box = parse_roi_box_from_landmark(pts68)
                img_step2 = crop_img(img_ori, roi_box)
                img_step2 = cv2.resize(img_step2,
                                       dsize=(STD_SIZE, STD_SIZE),
                                       interpolation=cv2.INTER_LINEAR)
                input = transform(img_step2).unsqueeze(0)
                with torch.no_grad():
                    if args.mode == 'gpu':
                        input = input.cuda()
                    param = model(input)
                    param = param.squeeze().cpu().numpy().flatten().astype(
                        np.float32)

                pts68 = predict_68pts(param, roi_box)

            pts_res.append(pts68)
            P, pose = parse_pose(param)

            Ps.append(P)
            poses.append(pose)

            if args.dump_obj:
                vertices = predict_dense(param, roi_box)
                vertices_lst.append(vertices)
                wfp = '{}_{}.obj'.format(img_fp.replace(suffix, ''), ind)
                colors = get_colors(img_ori, vertices)

                p, offset, alpha_shp, alpha_exp = _parse_param(param)

                vertices = (u + w_shp @ alpha_shp + w_exp @ alpha_exp).reshape(
                    3, -1, order='F') + offset
                vertices = vertices.T
                tri = tri.T - 1
                print('Dump obj with sampled texture to {}'.format(wfp))
                unwraps = create_unwraps(vertices)
                h, w = args.height, args.width
                tcoords = process_uv(unwraps[:, :2], h, w)
                texture = render_colors(tcoords, tri, colors, h, w,
                                        c=3).astype('uint8')
                scaled_tcoords = scale_tcoords(tcoords)
                if args.uvgan:
                    texture = uvgan.infer(texture)
                else:
                    texture = cv2.cvtColor(texture, cv2.COLOR_BGR2RGB)
                vertices, colors, uv_coords = vertices.astype(
                    np.float32).copy(), colors.astype(
                        np.float32).copy(), scaled_tcoords.astype(
                            np.float32).copy()
                write_obj_with_colors_texture(wfp, vertices, colors, tri,
                                              texture * 255.0, uv_coords)

            ind += 1
コード例 #18
0
ファイル: fit_3DDFA.py プロジェクト: loomai/RI_render_DPR
    def forward(self, img_ori, saveFolder, imgName):
        rects = self.face_detector(img_ori, 1)
        if len(rects) == 0:
            return
        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = saveFolder

        rect = rects[0]
        # whether use dlib landmark to crop image, if not, use only face bbox to calc roi bbox for cropping
        # - use landmark for cropping
        pts = self.face_regressor(img_ori, rect).parts()
        pts = np.array([[pt.x, pt.y] for pt in pts
                        ]).T  # detected landmarks, should record this
        roi_box = parse_roi_box_from_landmark(pts)
        img = crop_img(img_ori, roi_box)

        # forward: one step
        img = cv2.resize(img,
                         dsize=(STD_SIZE, STD_SIZE),
                         interpolation=cv2.INTER_LINEAR)
        input = self.transform(img).unsqueeze(0)
        with torch.no_grad():
            if self.mode == 'gpu':
                input = input.cuda()
            param = self.model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        # 68 pts
        pts68 = predict_68pts(param, roi_box)

        #pts_res.append(pts68)
        pts_res.append(pts)
        P, pose = parse_pose(param)
        Ps.append(P)
        poses.append(pose)

        # dense face 3d vertices
        vertices = predict_dense(param, roi_box)
        vertices_lst.append(vertices)

        # save projected 3D points
        wfp = '{}_projected.txt'.format(os.path.join(saveFolder, imgName))
        np.savetxt(wfp, pts68, fmt='%.3f')
        #print('Save 68 3d landmarks to {}'.format(wfp))

        wfp = '{}_detected.txt'.format(os.path.join(saveFolder, imgName))
        np.savetxt(wfp, pts, fmt='%.3f')
        #print('Save 68 3d landmarks to {}'.format(wfp))

        # save obj file
        wfp = '{}.obj'.format(os.path.join(saveFolder, imgName))
        colors = get_colors(img_ori, vertices)
        write_obj_with_colors(wfp, vertices, self.tri, colors)
        #print('Dump obj with sampled texture to {}'.format(wfp))

        wfp = os.path.join(saveFolder, imgName) + '_depth.png'
        depths_img = cget_depths_image(img_ori, vertices_lst,
                                       self.tri - 1)  # cython version
        cv2.imwrite(wfp, depths_img)
        #print('Dump to {}'.format(wfp))

        draw_landmarks(img_ori,
                       pts_res,
                       wfp=os.path.join(saveFolder, imgName) + '_3DDFA.png',
                       show_flg=False)
コード例 #19
0
    for item in tqdm(folder):
        try:
            img_ori = cv2.imread(str(folder[item]))
            rects = face_detector(img_ori, 1)

            if len(rects) != 0:
                for rect in rects:
                    bbox = [
                        rect.left(),
                        rect.top(),
                        rect.right(),
                        rect.bottom()
                    ]
                    roi_box = parse_roi_box_from_bbox(bbox)
                    img = crop_img(img_ori, roi_box)
                    img = cv2.resize(img,
                                     dsize=(STD_SIZE, STD_SIZE),
                                     interpolation=cv2.INTER_LINEAR)
                    input = transform(img).unsqueeze(0)
                    with torch.no_grad():
                        if args.mode == 'gpu':
                            input = input.cuda()
                        param = model(input)
                        param = param.squeeze().cpu().numpy().flatten().astype(
                            np.float32)

                vertices_lst = []
                vertices = predict_dense(param, roi_box)
                vertices_lst.append(vertices)
コード例 #20
0
def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    # 2. load dlib model for face detection and landmark used for face cropping
    if args.dlib_landmark:
        dlib_landmark_model = 'models/shape_predictor_68_face_landmarks.dat'
        face_regressor = dlib.shape_predictor(dlib_landmark_model)
    if args.dlib_bbox:
        face_detector = dlib.get_frontal_face_detector()

    # 3. forward
    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    files = sorted(
        glob.glob(os.path.join(args.folder, '*.jpg')) +
        glob.glob(os.path.join(args.folder, '*.png')))
    p_bar = tqdm(total=len(files))
    for img_fp in files:

        img_ori = cv2.imread(img_fp)
        if args.dlib_bbox:
            rects = face_detector(img_ori, 1)
        else:
            rects = []

        if len(rects) == 0:
            rects = dlib.rectangles()
            rect_fp = img_fp + '.bbox'
            lines = open(rect_fp).read().strip().split('\n')[1:]
            for l in lines:
                l, r, t, b = [int(_) for _ in l.split(' ')[1:]]
                rect = dlib.rectangle(l, r, t, b)
                rects.append(rect)

        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = get_suffix(img_fp)
        for rect in rects:
            # whether use dlib landmark to crop image, if not, use only face bbox to calc roi bbox for cropping
            if args.dlib_landmark:
                # - use landmark for cropping
                pts = face_regressor(img_ori, rect).parts()
                pts = np.array([[pt.x, pt.y] for pt in pts]).T
                roi_box = parse_roi_box_from_landmark(pts)
            else:
                # - use detected face bbox
                bbox = [rect.left(), rect.top(), rect.right(), rect.bottom()]
                roi_box = parse_roi_box_from_bbox(bbox)

            img = crop_img(img_ori, roi_box)

            # forward: one step
            img = cv2.resize(img,
                             dsize=(STD_SIZE, STD_SIZE),
                             interpolation=cv2.INTER_LINEAR)
            input_ = transform(img).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input_ = input_.cuda()
                param = model(input_)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)

            # 68 pts
            pts68 = predict_68pts(param, roi_box)

            # two-step for more accurate bbox to crop face
            if args.bbox_init == 'two':
                roi_box = parse_roi_box_from_landmark(pts68)
                img_step2 = crop_img(img_ori, roi_box)
                img_step2 = cv2.resize(img_step2,
                                       dsize=(STD_SIZE, STD_SIZE),
                                       interpolation=cv2.INTER_LINEAR)
                input_ = transform(img_step2).unsqueeze(0)
                with torch.no_grad():
                    if args.mode == 'gpu':
                        input_ = input_.cuda()
                    param = model(input_)
                    param = param.squeeze().cpu().numpy().flatten().astype(
                        np.float32)

                pts68 = predict_68pts(param, roi_box)

            pts_res.append(pts68)
            P, pose = parse_pose(param)
            Ps.append(P)
            poses.append(pose)

            points = np.array(pts_res)[0].T

            rotated = eulerAnglesToRotationMatrix(np.array([0., np.pi, 0.]))
            points = points.dot(rotated)

            scaler = MinMaxScaler(feature_range=(-1., 1))
            scaled_points = scaler.fit_transform(points)
            points = scaled_points

            f_name = img_fp.replace(args.folder + '/',
                                    '').replace('.png',
                                                '').replace('.jpg', '')
            np.save('./results/{}.npy'.format(f_name), points.reshape(-1))

            if args.plot:
                plot_face(points, img_fp)
                if args.show_flg:
                    plt.show()
                else:
                    plt.savefig('./results/{}.png'.format(f_name))

            p_bar.update(1)
コード例 #21
0
def getPoses(img_ori):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    cudnn.benchmark = True
    model = model.cuda()
    model.eval()

    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose([ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    alignment_model = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False,device='cuda')

    # face alignment model use RGB as input, result is a tuple with landmarks and boxes
    preds = alignment_model.get_landmarks(img_ori[:, :, ::-1])
    pts_2d_68 = preds[0]
    roi_box = parse_roi_box_from_landmark(pts_2d_68.T)

    img = crop_img(img_ori, roi_box)
    # import pdb; pdb.set_trace()

    # forward: one step
    img = cv2.resize(img, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
    input = transform(img).unsqueeze(0)
    with torch.no_grad():
        input = input.cuda()
        param = model(input)
        param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

    # 68 pts
    pts68 = predict_68pts(param, roi_box)

    roi_box = parse_roi_box_from_landmark(pts68)
    img_step2 = crop_img(img_ori, roi_box)
    img_step2 = cv2.resize(img_step2, dsize=(STD_SIZE, STD_SIZE), interpolation=cv2.INTER_LINEAR)
    input = transform(img_step2).unsqueeze(0)
    with torch.no_grad():
        input = input.cuda()
        param = model(input)
        param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

    P, pose = parse_pose(param)        
    # dense face 3d vertices
    vertices = predict_dense(param, roi_box)
    colors = get_colors(img_ori, vertices)
    # aligned_param = get_aligned_param(param)
    # vertices_aligned = predict_dense(aligned_param, roi_box)
    # h, w, c = 120, 120, 3
    h, w, c = img_ori.shape
    img_2d = crender_colors(vertices.T, (tri - 1).T, colors[:, ::-1], h, w)
    img_2d = img_2d[:,:,::-1]
        
    return img_2d, pose
コード例 #22
0
ファイル: runway_model.py プロジェクト: agermanidis/3DDFA
def classify(model, inputs):
    in_img = inputs['photo']
    img_ori = np.array(in_img)
    img_fp = 'samples/test1.jpg'

    face_detector = dlib.get_frontal_face_detector()

    # 3. forward
    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    #print(transform)
    rects = face_detector(img_ori, 1)

    pts_res = []
    Ps = []  # Camera matrix collection
    poses = []  # pose collection, [todo: validate it]
    vertices_lst = []  # store multiple face vertices
    ind = 0
    suffix = get_suffix(img_fp)
    for rect in rects:
        # - use detected face bbox
        bbox = [rect.left(), rect.top(), rect.right(), rect.bottom()]
        roi_box = parse_roi_box_from_bbox(bbox)

        img = crop_img(img_ori, roi_box)

        # forward: one step
        img = cv2.resize(img,
                         dsize=(STD_SIZE, STD_SIZE),
                         interpolation=cv2.INTER_LINEAR)
        input = transform(img).unsqueeze(0)
        print(input)
        with torch.no_grad():

            if mode == 'gpu':
                input = input.cuda()

            param = model(input)
            param = param.squeeze().cpu().numpy().flatten().astype(np.float32)

        # 68 pts
        pts68 = predict_68pts(param, roi_box)

        # two-step for more accurate bbox to crop face
        if bbox_init == 'two':
            roi_box = parse_roi_box_from_landmark(pts68)
            img_step2 = crop_img(img_ori, roi_box)
            img_step2 = cv2.resize(img_step2,
                                   dsize=(STD_SIZE, STD_SIZE),
                                   interpolation=cv2.INTER_LINEAR)
            input = transform(img_step2).unsqueeze(0)
            with torch.no_grad():
                if mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)

            pts68 = predict_68pts(param, roi_box)

        pts_res.append(pts68)
        P, pose = parse_pose(param)
        Ps.append(P)
        poses.append(pose)

        vertices = predict_dense(param, roi_box)
        vertices_lst.append(vertices)
        ind += 1

    pncc_feature = cpncc(img_ori, vertices_lst, tri - 1)
    output = pncc_feature[:, :, ::-1]
    print(type(output))
    pilImg = transforms.ToPILImage()(np.uint8(output))

    return {"image": pilImg}
コード例 #23
0
ファイル: main.py プロジェクト: sin00528/3DEmo
def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    # 2. load dlib model for face detection and landmark used for face cropping
    if args.dlib_landmark:
        dlib_landmark_model = 'models/shape_predictor_68_face_landmarks.dat'
        face_regressor = dlib.shape_predictor(dlib_landmark_model)
    if args.dlib_bbox:
        face_detector = dlib.get_frontal_face_detector()

    # 3. forward
    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    for img_fp in args.files:
        img_ori = cv2.imread(img_fp)
        if args.dlib_bbox:
            rects = face_detector(img_ori, 1)
        else:
            rects = []

        if len(rects) == 0:
            rects = dlib.rectangles()
            rect_fp = img_fp + '.bbox'
            try:
                lines = open(rect_fp).read().strip().split('\n')[1:]
            except FileNotFoundError:
                print('Cannot load bbox file')
                continue
            for l in lines:
                l, r, t, b = [int(_) for _ in l.split(' ')[1:]]
                rect = dlib.rectangle(l, r, t, b)
                rects.append(rect)

        pts_res = []
        Ps = []  # Camera matrix collection
        poses = []  # pose collection, [todo: validate it]
        vertices_lst = []  # store multiple face vertices
        ind = 0
        suffix = get_suffix(img_fp)
        for rect in rects:
            # whether use dlib landmark to crop image, if not, use only face bbox to calc roi bbox for cropping
            if args.dlib_landmark:
                # - use landmark for cropping
                pts = face_regressor(img_ori, rect).parts()
                pts = np.array([[pt.x, pt.y] for pt in pts]).T
                roi_box = parse_roi_box_from_landmark(pts)
            else:
                # - use detected face bbox
                bbox = [rect.left(), rect.top(), rect.right(), rect.bottom()]
                roi_box = parse_roi_box_from_bbox(bbox)

            img = crop_img(img_ori, roi_box)

            # forward: one step
            img = cv2.resize(img,
                             dsize=(STD_SIZE, STD_SIZE),
                             interpolation=cv2.INTER_LINEAR)
            input = transform(img).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)

            # 68 pts
            pts68 = predict_68pts(param, roi_box)

            # two-step for more accurate bbox to crop face
            if args.bbox_init == 'two':
                roi_box = parse_roi_box_from_landmark(pts68)
                img_step2 = crop_img(img_ori, roi_box)
                img_step2 = cv2.resize(img_step2,
                                       dsize=(STD_SIZE, STD_SIZE),
                                       interpolation=cv2.INTER_LINEAR)
                input = transform(img_step2).unsqueeze(0)
                with torch.no_grad():
                    if args.mode == 'gpu':
                        input = input.cuda()
                    param = model(input)
                    param = param.squeeze().cpu().numpy().flatten().astype(
                        np.float32)

                pts68 = predict_68pts(param, roi_box)

            pts_res.append(pts68)
            P, pose = parse_pose(param)
            Ps.append(P)
            poses.append(pose)

            # dense face 3d vertices
            if args.dump_ply or args.dump_vertex or args.dump_depth or args.dump_pncc or args.dump_obj:
                vertices = predict_dense(param, roi_box)
                vertices_lst.append(vertices)
            if args.dump_ply:
                dump_to_ply(
                    vertices, tri,
                    '{}_{}.ply'.format(img_fp.replace(suffix, ''), ind))
            if args.dump_vertex:
                dump_vertex(
                    vertices, '{}_{}.mat'.format(img_fp.replace(suffix, ''),
                                                 ind))
            if args.dump_pts:
                wfp = '{}_{}.txt'.format(img_fp.replace(suffix, ''), ind)
                np.savetxt(wfp, pts68, fmt='%.3f')
                print('Save 68 3d landmarks to {}'.format(wfp))
            if args.dump_roi_box:
                wfp = '{}_{}.roibox'.format(img_fp.replace(suffix, ''), ind)
                np.savetxt(wfp, roi_box, fmt='%.3f')
                print('Save roi box to {}'.format(wfp))
            if args.dump_paf:
                wfp_paf = '{}_{}_paf.jpg'.format(img_fp.replace(suffix, ''),
                                                 ind)
                wfp_crop = '{}_{}_crop.jpg'.format(img_fp.replace(suffix, ''),
                                                   ind)
                paf_feature = gen_img_paf(img_crop=img,
                                          param=param,
                                          kernel_size=args.paf_size)

                cv2.imwrite(wfp_paf, paf_feature)
                cv2.imwrite(wfp_crop, img)
                print('Dump to {} and {}'.format(wfp_crop, wfp_paf))
            if args.dump_obj:
                wfp = '{}_{}.obj'.format(img_fp.replace(suffix, ''), ind)
                colors = get_colors(img_ori, vertices)
                write_obj_with_colors(wfp, vertices, tri, colors)
                print('Dump obj with sampled texture to {}'.format(wfp))
            ind += 1

        if args.dump_pose:
            # P, pose = parse_pose(param)  # Camera matrix (without scale), and pose (yaw, pitch, roll, to verify)
            img_pose = plot_pose_box(img_ori, Ps, pts_res)
            wfp = img_fp.replace(suffix, '_pose.jpg')
            cv2.imwrite(wfp, img_pose)
            print('Dump to {}'.format(wfp))
        if args.dump_depth:
            wfp = img_fp.replace(suffix, '_depth.png')
            # depths_img = get_depths_image(img_ori, vertices_lst, tri-1)  # python version
            depths_img = cget_depths_image(img_ori, vertices_lst,
                                           tri - 1)  # cython version
            cv2.imwrite(wfp, depths_img)
            print('Dump to {}'.format(wfp))
        if args.dump_pncc:
            wfp = img_fp.replace(suffix, '_pncc.png')
            pncc_feature = cpncc(img_ori, vertices_lst,
                                 tri - 1)  # cython version
            cv2.imwrite(
                wfp,
                pncc_feature[:, :, ::-1])  # cv2.imwrite will swap RGB -> BGR
            print('Dump to {}'.format(wfp))
        if args.dump_res:
            draw_landmarks(img_ori,
                           pts_res,
                           wfp=img_fp.replace(suffix, '_3DDFA.jpg'),
                           show_flg=args.show_flg)
コード例 #24
0
ファイル: cap_demo.py プロジェクト: zjjMaiMai/3DDFA
def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    app = RenderPipeline(**cfg)

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    face_detector = dlib.get_frontal_face_detector()

    # 3. forward
    tri = sio.loadmat('tri_refine.mat')['tri']
    tri = _to_ctype(tri).astype(np.int32)  # for type compatible
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    last_frame_lmks = []

    vc = cv2.VideoCapture("C:\\Users\\zh13\\Videos\\TrackingTest.mov")
    success, frame = vc.read()
    while success:
        roi_box = []

        if len(last_frame_lmks) == 0:
            rects = face_detector(frame, 1)
            for rect in rects:
                bbox = [rect.left(), rect.top(), rect.right(), rect.bottom()]
                roi_box.append(parse_roi_box_from_bbox(bbox))
        else:
            for lmk in last_frame_lmks:
                roi_box.append(parse_roi_box_from_landmark(lmk))

        this_frame_lmk = []
        params = []
        for box in roi_box:
            img_to_net = crop_img(frame, box)
            img_to_net = cv2.resize(img_to_net,
                                    dsize=(STD_SIZE, STD_SIZE),
                                    interpolation=cv2.INTER_LINEAR)
            input = transform(img_to_net).unsqueeze(0)
            with torch.no_grad():
                if args.mode == 'gpu':
                    input = input.cuda()
                param = model(input)
                param = param.squeeze().cpu().numpy().flatten().astype(
                    np.float32)
                params.append(param)
            this_frame_lmk.append(predict_68pts(param, box))
        last_frame_lmks = this_frame_lmk

        if args.render_mesh:
            for box, param in zip(roi_box, params):
                vertices = predict_dense(param, box)
                frame = app(_to_ctype(vertices.T), tri,
                            _to_ctype(frame.astype(np.float32) / 255.))
        else:
            for lmk in last_frame_lmks:
                for p in lmk.T:
                    cv2.circle(frame, (int(round(p[0] * draw_multiplier)),
                                       int(round(p[1] * draw_multiplier))),
                               draw_multiplier, (255, 0, 0), 1, cv2.LINE_AA,
                               draw_shiftbits)
        cv2.imshow("3ddfa video demo", frame)
        cv2.waitKey(1)
        success, frame = vc.read()
コード例 #25
0
def get_landmark_2d(root, image_path):
    # 0.read image
    img_ori = cv2.imread(os.path.join(root, image_path))
    # 1. load pre-tained model
    checkpoint_fp = 'models/MobDenseNet.pth.tar'
    arch = 'mobdensenet_v1'
    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(MobDenseNet, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)
    model_dict = model.state_dict()
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()
    # 2. load dlib model for face detection and landmark used for face cropping
    if args.dlib_landmark:
        dlib_landmark_model = 'models/shape_predictor_68_face_landmarks.dat'
        face_regressor = dlib.shape_predictor(dlib_landmark_model)
    if args.dlib_bbox:
        face_detector = dlib.get_frontal_face_detector()
    # 3. forward
    tri = sio.loadmat('visualize/tri.mat')['tri'] - 1
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])
    imgScale = 1
    scaledImg = img_ori
    if max(img_ori.shape) > maxImgSizeForDetection:
        imgScale = maxImgSizeForDetection / float(max(img_ori.shape))
        scaledImg = cv2.resize(img_ori, (int(
            img_ori.shape[1] * imgScale), int(img_ori.shape[0] * imgScale)))
    rects = face_detector(scaledImg, 1)
    for rect in rects:
        if args.dlib_landmark:
            faceRectangle = rectangle(int(rect.left() / imgScale),
                                      int(rect.top() / imgScale),
                                      int(rect.right() / imgScale),
                                      int(rect.bottom() / imgScale))
            # - use landmark for cropping
            pts = face_regressor(img_ori, faceRectangle).parts()
            pts = np.array([[pt.x, pt.y] for pt in pts]).T
            roi_box = parse_roi_box_from_landmark(pts)
        else:
            bbox = [
                int(rect.left() / imgScale),
                int(rect.top() / imgScale),
                int(rect.right() / imgScale),
                int(rect.bottom() / imgScale)
            ]
            roi_box = parse_roi_box_from_bbox(bbox)
    img = crop_img(img_ori, roi_box)
    # forward: one step
    img = cv2.resize(img,
                     dsize=(STD_SIZE, STD_SIZE),
                     interpolation=cv2.INTER_LINEAR)
    input = transform(img).unsqueeze(0)
    with torch.no_grad():
        if args.mode == 'gpu':
            input = input.cuda()
        param = model(input)
        param = param.squeeze().cpu().numpy().flatten().astype(np.float32)
    # 68 pts
    pts68 = predict_68pts(param, roi_box)
    return pts68
コード例 #26
0
def main(args):
    # 1. load pre-tained model
    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
    arch = 'mobilenet_1'

    checkpoint = torch.load(
        checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
    model = getattr(mobilenet_v1, arch)(
        num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)

    model_dict = model.state_dict()
    # because the model is trained by multiple gpus, prefix module should be removed
    for k in checkpoint.keys():
        model_dict[k.replace('module.', '')] = checkpoint[k]
    model.load_state_dict(model_dict)
    if args.mode == 'gpu':
        cudnn.benchmark = True
        model = model.cuda()
    model.eval()

    tri = sio.loadmat('visualize/tri.mat')['tri']
    transform = transforms.Compose(
        [ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)])

    if not os.path.exists(args.save_dir):
        os.mkdir(args.save_dir)

    # 2. parse images list and landmark
    lmk_file = args.lmk_file
    ts = time.time()
    rank_land, rank_img_list, start, end = parse_quality_list_part(
        lmk_file, args.world_size, args.rank, args.resume_idx)
    print('parse land file in {:.3f} seconds'.format(time.time() - ts))

    # for batch processing
    print('World size {}, rank {}, start from {}, end with {}'.format(
        args.world_size, args.rank, start, end))
    dataset = McDataset(rank_img_list,
                        rank_land,
                        transform=transform,
                        std_size=STD_SIZE)
    dataloader = DataLoader(dataset,
                            batch_size=args.batch_size,
                            shuffle=False,
                            num_workers=2,
                            pin_memory=True)

    for img_idx, (inputs, ori_imgs, img_fps,
                  roi_boxes) in enumerate(tqdm(dataloader)):

        # forward: one step
        with torch.no_grad():
            if args.mode == 'gpu':
                inputs = inputs.cuda()
            params = model(inputs)
            params = params.cpu().numpy()

        roi_boxes = roi_boxes.numpy()
        outputs_roi_boxes = roi_boxes
        if args.bbox_init == 'two':
            step_two_ori_imgs = []
            step_two_roi_boxes = []
            ori_imgs = ori_imgs.numpy()
            for ii in range(params.shape[0]):
                # 68 pts
                pts68 = predict_68pts(params[ii], roi_boxes[ii])

                # two-step for more accurate bbox to crop face
                roi_box = parse_roi_box_from_landmark(pts68)
                img_step2 = crop_img(ori_imgs[ii], roi_box)
                img_step2 = cv2.resize(img_step2,
                                       dsize=(STD_SIZE, STD_SIZE),
                                       interpolation=cv2.INTER_LINEAR)
                # input = transform(img_step2).unsqueeze(0)
                step_two_ori_imgs.append(transform(img_step2))
                step_two_roi_boxes.append(roi_box)
            with torch.no_grad():
                step_two_ori_imgs = torch.stack(step_two_ori_imgs, dim=0)
                inputs = step_two_ori_imgs
                if args.mode == 'gpu':
                    inputs = inputs.cuda()
                params = model(inputs)
                params = params.cpu().numpy()
            outputs_roi_boxes = step_two_roi_boxes

        # dump results
        if args.dump_param:
            for img_fp, param, roi_box in zip(img_fps, params,
                                              outputs_roi_boxes):
                split = img_fp.split('/')
                save_name = os.path.join(
                    args.save_dir,
                    '{}.txt'.format(os.path.splitext(split[-1])[0]))
                this_param = param * param_std + param_mean
                this_param = np.concatenate((this_param, roi_box))
                this_param.tofile(save_name, sep=' ')