Ejemplo n.º 1
0
def get_train_ioput(train_data, bldshps):

    num_exp = train_data[0].exp.shape[0] - 1
    num_land = train_data[0].land.shape[0]
    num_angle = 3
    num_tslt = 3

    mean_ldmk, tri_idx, px_barycenter = load.load_tri_idx(
        '../const_file/tri_idx_px.txt', num_land)
    num_px = px_barycenter[0].shape[0]

    train_data_input = np.empty((len(train_data), num_px))
    train_data_output = np.empty(
        (len(train_data), num_angle + num_tslt + num_exp + num_land * 2))

    for counter, one_data in enumerate(train_data):
        #input:
        init_land = util.get_init_land(one_data, bldshps)

        # =============================================================================
        #         print(init_land[tri_idx[px_barycenter[0]][:,0]].shape)
        #         print(tri_idx[px_barycenter[0]][:,0].shape)
        #         print(tri_idx[px_barycenter[0]].shape)
        #         print(px_barycenter[1][:,0].shape)
        #         print(init_land[tri_idx[px_barycenter[0]][:,1]].shape)
        # =============================================================================

        #        t=init_land[tri_idx[px_barycenter[0]][:,0]]*px_barycenter[1][:,0].reshape(num_px,1)

        #        print(init_land[tri_idx[px_barycenter[0]][:,0]],px_barycenter[1][:,0].reshape(num_px,1),t)

        train_data_input[counter] = util.get_input_from_land_img(
            init_land, one_data.img, tri_idx, px_barycenter)

        #output:
        train_data_output[counter, 0:num_angle]\
            =(one_data.angle-one_data.init_angle).copy()
        train_data_output[counter,num_angle:num_angle+num_tslt]\
            =(one_data.tslt-one_data.init_tslt).copy()
        train_data_output[counter,num_angle+num_tslt:num_angle+num_tslt+num_exp]\
            =(one_data.exp-one_data.init_exp)[1:num_exp+1].copy()
        train_data_output[counter,num_angle+num_tslt+num_exp:num_angle+num_tslt+num_exp+num_land]\
            =(one_data.dis-one_data.init_dis)[:,0].copy()
        train_data_output[counter,num_angle+num_tslt+num_exp+num_land:num_angle+num_tslt+num_exp+num_land*2]\
            =(one_data.dis-one_data.init_dis)[:,1].copy()

    return train_data_input, train_data_output
Ejemplo n.º 2
0
def train_net(train_data_input, train_data_output):
    os.environ["CUDA_VISIBLE_DEVICES"] = '3'
    LR = 0.5
    BATCH_SIZE = 10000
    max_epoch = 20000
    test_epoch_num = 400
    hid_mul = 2
    server = '104'
    file_num = 6
    STEP_SIZE = 2000
    GAMMA = 0.5
    suffix='_7_24_cf3_f'+str(file_num)+'_lr'+str(LR*100)+'_bch'+\
            str(BATCH_SIZE/10000)+'w_hid'+str(hid_mul)+'_it'+str(max_epoch/10000)+'w_'+server

    train_data_input = train_data_input.astype(np.float32)
    train_data_output = train_data_output.astype(np.float32)

    num_input = train_data_input.shape[1]
    num_output = train_data_output.shape[1]

    print(torch.cuda.device_count())
    #    torch.cuda.set_device(6,7)

    net = dde_net_fc(num_input, num_input * hid_mul, num_output).cuda()
    # print(net)

    optimizer = torch.optim.SGD(net.parameters(), lr=LR)
    #    optimizer= torch.optim.Adam(net.parameters(), lr=LR, betas=(0.9, 0.99))
    loss_func = torch.nn.MSELoss()
    #
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=STEP_SIZE,
                                                gamma=GAMMA)

    # =============================================================================
    #     x=torch.tensor(train_data_input)
    #     y=torch.tensor(train_data_output)
    # =============================================================================

    train_data_num = train_data_input.shape[0]  #10000
    x = torch.Tensor(train_data_input[:train_data_num]).cuda()
    y = torch.Tensor(train_data_output[:train_data_num]).cuda()

    torch_dataset = Data.TensorDataset(x, y)

    train_loader = Data.DataLoader(dataset=torch_dataset,
                                   batch_size=BATCH_SIZE,
                                   shuffle=True,
                                   num_workers=6)

    data, bldshps = load.load_dataHEbldshps()
    mean_ldmk, tri_idx, px_barycenter = load.load_tri_idx(
        '../../const_file/tri_idx_px.txt', data[0].data[0].land.shape[0])

    loss_plt = []

    test_cnt = np.empty((max_epoch // test_epoch_num, max_cnt))

    for epoch in range(max_epoch):
        for step, (batch_x, batch_y) in enumerate(train_loader):

            predict = net.forward(batch_x)
            loss = loss_func(predict, batch_y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            loss_plt.append(loss.data.cpu().numpy())
            print('epoch: %d ' % epoch, 'step: %d ' % step,
                  'Loss=%.4f' % loss_plt[-1])
# =============================================================================
#         if epoch % 5 ==0:
#             print('epoch: %d '% epoch ,'Loss=%.4f'%loss.data.numpy())
# =============================================================================
        scheduler.step()
        if epoch % test_epoch_num == 0:
            test_cnt[epoch // test_epoch_num] = net_test_img_stc(
                data, net, bldshps, mean_ldmk, tri_idx, px_barycenter)
            rate = test_cnt[epoch // test_epoch_num, :10].sum() / test_cnt[
                epoch // test_epoch_num].sum()
            print('epoch: %d ' % epoch, 'test <1 rate=%.4f' % rate)

    torch.save(net, '../mid_data/net' + suffix + '.pkl')  # save entire net
    #    torch.save(net.state_dict(), '../mid_data/net_7_17_cf3_params_norm_gpu.pkl')   # save only the parameters

    np.save('../mid_data/test_cnt' + suffix, test_cnt)

    print(
        'final test <1 rate=%.4f' %
        (test_cnt[-1, :10].sum() / test_cnt[-1].sum()),
        '  <0.5=%.4f' % (test_cnt[-1, :5].sum() / test_cnt[-1].sum()))

    y = test_cnt[:, :10].sum(1)
    plt.ion()
    plt.cla()
    plt.figure(file_num, figsize=(16, 12))

    plt.subplot(221)
    plt.scatter(np.linspace(0, len(loss_plt), num=len(loss_plt)), loss_plt)
    plt.plot(np.linspace(0, len(loss_plt), num=len(loss_plt)),
             loss_plt,
             'r-',
             lw=1,
             label='loss')
    plt.legend(loc='best')

    plt.subplot(222)
    plt.scatter(np.linspace(0, len(y), num=len(y)), y)
    plt.plot(np.linspace(0, len(y), num=len(y)), y, 'r-', lw=1, label='<1')
    plt.legend(loc='best')

    plt.subplot(223)
    plt.scatter(np.linspace(0, len(y), num=len(y)), test_cnt[:, :5].sum(1))
    plt.plot(np.linspace(0, len(y), num=len(y)),
             test_cnt[:, :5].sum(1),
             'r-',
             lw=1,
             label='<0.5')
    plt.legend(loc='best')

    plt.subplot(224)
    plt.scatter(np.linspace(0, len(y), num=len(y)), test_cnt[:, :10].sum(1))
    plt.plot(np.linspace(0, len(y), num=len(y)),
             test_cnt[:, :10].sum(1),
             'r-',
             lw=1)
    plt.scatter(np.linspace(0, len(y), num=len(y)), test_cnt[:, :8].sum(1))
    plt.plot(np.linspace(0, len(y), num=len(y)),
             test_cnt[:, :8].sum(1),
             'b-',
             lw=1)
    plt.scatter(np.linspace(0, len(y), num=len(y)), test_cnt[:, :5].sum(1))
    plt.plot(np.linspace(0, len(y), num=len(y)),
             test_cnt[:, :5].sum(1),
             'y-',
             lw=1)
    plt.scatter(np.linspace(0, len(y), num=len(y)), test_cnt[:, :3].sum(1))
    plt.plot(np.linspace(0, len(y), num=len(y)),
             test_cnt[:, :3].sum(1),
             'g-',
             lw=1)

    plt.ioff()
    plt.show()
Ejemplo n.º 3
0
def test_video(train_data, net, path):
    def update_train_data(train_data, spf_bldshps, f_vd, center_vd, tslt2):
        test_all_data = []
        for one_ide in train_data:
            for one_img in one_ide.data:
                test_all_data.append(base.TestVideoOnePoint(one_img))

                test_all_data[-1].fcs = f_vd
                test_all_data[-1].center = center_vd
                test_all_data[-1].tslt[2] = tslt2

                test_all_data[-1].land_inv = util.get_land_spfbldshps_inv(
                    test_all_data[-1], spf_bldshps)
                #                test_all_data[-1].land=test_all_data[-1].land_inv.copy()
                #                test_all_data[-1].land[:,1]=one_img.img.shape[0]-test_all_data[-1].land_inv[:,1]

                test_all_data[-1].centroid_inv = test_all_data[
                    -1].land_inv.mean(0)

        return test_all_data

    def get_init_shape(init_num, all_test_data, test_data, spf_bldshps):
        ans = [0] * init_num
        ans_dis = [1e9] * init_num
        for i in range(init_num):
            ans[i] = all_test_data[i]

        for one_shape in all_test_data:
            dis = np.linalg.norm((one_shape.land_inv - test_data.land_inv) -
                                 (one_shape.centroid_inv -
                                  test_data.centroid_inv))
            for i in range(init_num):
                if (ans_dis[i] > dis):
                    ans[i + 1:] = ans[i:-1]
                    ans_dis[i + 1:] = ans_dis[i:-1]
                    ans[i] = one_shape
                    ans_dis[i] = dis
                    break
            #tslt?

#        return ans
#        print(ans[:].tslt)
#        print(test_data.tslt)
        print(ans_dis)
        ansx = copy.deepcopy(ans)
        ans_dis_aft = ans_dis.copy()
        for x, y in zip(ansx, ans_dis_aft):
            x.tslt[:2] += (-x.centroid_inv +
                           test_data.centroid_inv) / test_data.fcs * x.tslt[2]
            util.get_slt_land_cor(x, spf_bldshps, x.exp)
            land_inv = util.get_land_spfbldshps_inv(x, spf_bldshps)
            y = np.linalg.norm((land_inv - test_data.land_inv))

        print(ansx[0].tslt)
        print(ans[0].tslt)
        print(test_data.tslt)

        return ansx, ans_dis, ans_dis_aft

    def get_init_shape_inner(init_num, all_test_data, test_data, spf_bldshps):
        ans = [0] * init_num
        ans_dis = [1e9] * init_num
        for i in range(init_num):
            ans[i] = all_test_data[i]

        for one_shape in all_test_data:
            dis = np.linalg.norm(
                (one_shape.land_inv[15:] - test_data.land_inv[15:]) -
                (one_shape.centroid_inv - test_data.centroid_inv))
            for i in range(init_num):
                if (ans_dis[i] > dis):
                    ans[i + 1:] = ans[i:-1]
                    ans_dis[i + 1:] = ans_dis[i:-1]
                    ans[i] = one_shape
                    ans_dis[i] = dis
                    break
            #tslt?

#        return ans
#        print(ans[:].tslt)
#        print(test_data.tslt)
        print(ans_dis)
        ansx = copy.deepcopy(ans)
        ans_dis_aft = ans_dis.copy()
        for x, y in zip(ansx, ans_dis_aft):
            x.tslt[:2] += (-x.centroid_inv +
                           test_data.centroid_inv) / test_data.fcs * x.tslt[2]
            util.get_slt_land_cor(x, spf_bldshps, x.exp)
            land_inv = util.get_land_spfbldshps_inv(x, spf_bldshps)
            y = np.linalg.norm((land_inv[15:] - test_data.land_inv[15:]))

        print(ansx[0].tslt)
        print(ans[0].tslt)
        print(test_data.tslt)

        return ansx, ans_dis, ans_dis_aft

    bldshps_path = '/home/weiliu/fitting_dde/const_file/deal_data/blendshape_ide_svd_77.lv'
    bldshps = load.load_bldshps(bldshps_path,
                                num_ide=77,
                                num_exp=47,
                                num_vtx=11510)
    util.bld_reduce_neu(bldshps)

    num_exp = bldshps.shape[1] - 1
    num_land = train_data[0].data[0].land.shape[0]
    num_angle = 3
    num_tslt = 3
    num_ide = bldshps.shape[0]

    test_data = base.TestVideoOnePoint(train_data[0].data[0])
    user = np.empty((num_ide), dtype=np.float64)
    print(path + '_pre.psp_f')
    load.load_psp_f(path + '_pre.psp_f', test_data, user, num_ide, 47,
                    num_land)

    #    load.load_img(path+'_first_frame.jpg',test_data)
    #    load.load_land73(path+'_first_frame.land73',test_data,num_land)

    spf_bldshps = np.tensordot(bldshps, user, axes=(0, 0))
    test_image = cv2.imread(path + '_first_frame.jpg')

    all_test_data = update_train_data(train_data, spf_bldshps, test_data.fcs,
                                      test_data.center, test_data.tslt[2])

    #    for x in all_test_data:
    #        for pt in x.land_inv:
    #            cv2.circle(test_image,tuple(np.around(pt).astype(np.int)), 1, (255,255,255), -1)
    #    cv2.imshow('test frame',test_image)
    #    cv2.waitKey(0)

    mean_ldmk, tri_idx, px_barycenter = load.load_tri_idx(
        '../const_file/tri_idx_px.txt', num_land)

    ip_vd = cv2.VideoCapture(path + '.mp4')

    fps = ip_vd.get(cv2.CAP_PROP_FPS)
    size = (int(ip_vd.get(cv2.CAP_PROP_FRAME_WIDTH)),
            int(ip_vd.get(cv2.CAP_PROP_FRAME_HEIGHT)))

    init_num = 1
    ld_wt_vd = cv2.VideoWriter(
        path + model_suffix + '_initshape' + str(init_num) + '_inner.mp4',
        cv2.VideoWriter_fourcc(*'mp4v'), fps, size)

    test_data.land = test_data.land_inv.copy()
    test_data.land[:, 1] = size[1] - test_data.land_inv[:, 1]
    test_data.centroid_inv = test_data.land_inv.mean(0)

    ini_tt_dt = copy.deepcopy(test_data)
    init_land_data = copy.deepcopy(test_data)

    dbg_init_dis_file = open('dbg_init_dis_file.txt', 'w')
    dbg_dis_norm_file = open('dis_norm_file.txt', 'w')
    #    dbg_init_aft_dis_file=open('../mid_data/dbg_init_aft_dis_file.txt','w')

    for idx in range(10000):
        ret, rgb_frame = ip_vd.read()
        if not ret:
            break

        gray_frame = cv2.cvtColor(rgb_frame, cv2.COLOR_BGR2GRAY)

        for pt in test_data.land:
            cv2.circle(rgb_frame, tuple(np.around(pt).astype(np.int)), 2,
                       (255, 0, 0), -1)

        init_shape, dis, ans_dis_aft = get_init_shape_inner(
            init_num, all_test_data, test_data, spf_bldshps)

        dbg_init_dis_file.write(str([idx, dis, ans_dis_aft]) + '\n')

        #        result=np.zeros((num_angle+num_tslt+num_exp+2*num_land,),dtype=np.float64)
        angle, tslt = np.zeros(num_angle,
                               dtype=np.float64), np.zeros(num_tslt,
                                                           dtype=np.float64)
        exp, dis = np.zeros(num_exp, dtype=np.float64), np.zeros(
            (num_land, 2), dtype=np.float64)

        for x in init_shape:
            land = util.get_land_spfbldshps_inv(x, spf_bldshps)
            for pt in land:
                cv2.circle(rgb_frame, tuple(np.around(pt).astype(np.int)), 1,
                           (122, 122, 122), -1)
            land[:, 1] = gray_frame.shape[0] - land[:, 1]
            for pt in land:
                cv2.circle(rgb_frame, tuple(np.around(pt).astype(np.int)), 2,
                           (255, 255, 255), -1)
            input_now = util.get_input_from_land_img(land, gray_frame, tri_idx,
                                                     px_barycenter)
            print(input_now.shape)
            result = net(torch.tensor(input_now.astype(
                np.float32))).data.numpy()

            print(result)
            print('f: ', test_data.fcs, x.fcs)
            print(test_data.angle, test_data.tslt)
            print(x.angle, x.tslt)
            print(result[0:num_angle], result[num_angle:num_angle + num_tslt])
            print(x.angle + result[0:num_angle],
                  x.tslt + result[num_angle:num_angle + num_tslt])
            print(
                'dis x:',
                x.dis[:, 0] + result[num_angle + num_tslt + num_exp:num_angle +
                                     num_tslt + num_exp + num_land])
            print(
                'dis y:', x.dis[:, 0] +
                result[num_angle + num_tslt + num_exp + num_land:num_angle +
                       num_tslt + num_exp + num_land * 2])
            angle += x.angle + result[0:num_angle]
            tslt += x.tslt + result[num_angle:num_angle + num_tslt]
            exp += x.exp[1:] + result[num_angle + num_tslt:num_angle +
                                      num_tslt + num_exp]
            dis[:, 0] += x.dis[:, 0] + result[num_angle + num_tslt +
                                              num_exp:num_angle + num_tslt +
                                              num_exp + num_land]
            dis[:, 1] += x.dis[:, 1] + result[num_angle + num_tslt + num_exp +
                                              num_land:num_angle + num_tslt +
                                              num_exp + num_land * 2]

            init_land_data.angle = x.angle + result[0:num_angle]
            init_land_data.tslt = x.tslt + result[num_angle:num_angle +
                                                  num_tslt]
            init_land_data.exp[1:] = x.exp[1:] + result[num_angle +
                                                        num_tslt:num_angle +
                                                        num_tslt + num_exp]
            init_land_data.dis[:,
                               0] = x.dis[:, 0] + result[num_angle + num_tslt +
                                                         num_exp:num_angle +
                                                         num_tslt + num_exp +
                                                         num_land]
            init_land_data.dis[:,
                               0] = x.dis[:, 1] + result[num_angle + num_tslt +
                                                         num_exp +
                                                         num_land:num_angle +
                                                         num_tslt + num_exp +
                                                         num_land * 2]

            init_shape_land = util.get_land_spfbldshps_inv(
                init_land_data, spf_bldshps)
            init_shape_land[:, 1] = gray_frame.shape[0] - init_shape_land[:, 1]
# =============================================================================
#             print(result)
#             for pt in init_shape_land:
#                 cv2.circle(rgb_frame,tuple(np.around(pt).astype(np.int)), 1, (0,0,255), -1)
#             cv2.imshow('test frame',rgb_frame)
#             cv2.waitKey(0)
#             cv2.destroyAllWindows()
# =============================================================================

        test_data.angle = angle / init_num
        test_data.tslt = tslt / init_num
        test_data.exp[1:num_exp + 1] = exp / init_num
        test_data.dis = dis / init_num
        test_data.land_cor = init_shape[0].land_cor

        test_data.land_inv = util.get_land_spfbldshps_inv(
            test_data, spf_bldshps)
        test_data.land = test_data.land_inv.copy()
        test_data.land[:, 1] = gray_frame.shape[0] - test_data.land_inv[:, 1]

        test_data.centroid_inv = test_data.land_inv.mean(0)

        print('dif vs ini:', ini_tt_dt.angle - test_data.angle)
        print('dif vs ini:', ini_tt_dt.tslt - test_data.tslt)

        dbg_dis_norm_file.write(str(np.linalg.norm(test_data.dis)) + '\n')

        for pt in test_data.land:
            cv2.circle(rgb_frame, tuple(np.around(pt).astype(np.int)), 2,
                       (0, 255, 0), -1)

        ld_wt_vd.write(rgb_frame)

        load.save_psp_f(path + '_psp_f/lv_out_psp_f_tt_' + str(idx) + '.psp_f',
                        test_data, user)

# =============================================================================
#         if idx>90:
#             cv2.imshow('test frame',rgb_frame)
#             cv2.waitKey(0)
#             cv2.destroyAllWindows()
# =============================================================================

    ip_vd.release()
    ld_wt_vd.release()
    dbg_init_dis_file.close()
    dbg_dis_norm_file.close()
Ejemplo n.º 4
0
def test_set_img_stc(data,net):
    
    def gen_rand_data(data,all_data):
        for i in range(3):        
            data.angle[i]+=2*train.rand_angle_range[i]*np.random.random((1,))-train.rand_angle_range[i]
            
        for i in range(2):
            data.tslt[i]\
                +=2*train.rand_tslt_range[i]*np.random.random((1,))-train.rand_tslt_range[i]         
        data.tslt[2]\
            +=2*train.rand_tslt_range[2]*data.tslt[2]*np.random.random((1,))-train.rand_tslt_range[2]*data.tslt[2]
        
        for i in range(data.exp.shape[0]):
            if (i==0):
                continue
            data.exp[i]\
                +=2*train.rand_exp_range*np.random.random((1,))-train.rand_exp_range
# =============================================================================
#         first=random.randrange(len(all_data))
#         second=random.randrange(len(all_data[first].data))
#         data.exp[i]=all_data[first].data[second].exp.copy()
# =============================================================================
        
        first=random.randrange(len(all_data))
        second=random.randrange(len(all_data[first].data))
        data.dis=all_data[first].data[second].dis.copy()
            
    
    bldshps_path='/home/weiliu/fitting_dde/const_file/deal_data/blendshape_ide_svd_77.lv'
    bldshps=load.load_bldshps(bldshps_path,num_ide=77,num_exp=47,num_vtx=11510)
    util.bld_reduce_neu(bldshps)
    
    num_exp=bldshps.shape[1]-1
    num_land=data[0].data[0].land.shape[0]
    num_angle=3
    num_tslt=3
    #util.get_land(data[0].data[0],bldshps,data[0].user)
    
    ITE_num=1
    
    mean_ldmk,tri_idx,px_barycenter=load.load_tri_idx('../const_file/tri_idx_px.txt',num_land)
    
    max_cnt=100
    cnt=np.zeros((max_cnt),int)
    
    idx_n=0
    tot_num=0
    for one_ide in data:
        spf_bldshps=np.tensordot(bldshps,one_ide.user,axes=(0,0))
        
        for one_img in one_ide.data:
            p=np.random.random(1)[0]
            idx_n+=1
            print('image num now: %d tot image: %d' % (idx_n,tot_num))
            if (p>0.01):
                continue
            tot_num+=1             
            for ite in range(ITE_num):
                test_data=base.TestOnePoint(one_img,one_ide.user)
                std_land=test_data.land
                gen_rand_data(test_data,data)
#                bf_land=util.get_land(test_data,bldshps,one_ide.user)
                bf_land=util.get_land_spfbldshps(test_data,spf_bldshps)
                data_input=util.get_input_from_land_img(bf_land,test_data.img,tri_idx,px_barycenter)
                                
                out=net(torch.tensor(data_input.astype(np.float32)))
                out_data=out.data.numpy()
                angle,tslt,exp=out_data[0:num_angle],out_data[num_angle:num_angle+num_tslt],out_data[num_angle+num_tslt:num_angle+num_tslt+num_exp]
                dis=np.empty((num_land,2))
                dis[:,0]=out_data[num_angle+num_tslt+num_exp:num_angle+num_tslt+num_exp+num_land]
                dis[:,1]=out_data[num_angle+num_tslt+num_exp+num_land:num_angle+num_tslt+num_exp+num_land*2]
    
                test_data.angle+=angle
                test_data.tslt+=tslt
                test_data.exp[1:num_exp+1]+=exp
                test_data.dis+=dis
                
#                aft_land=util.get_land(test_data,bldshps,one_ide.user)
                aft_land=util.get_land_spfbldshps(test_data,spf_bldshps)
                
                rate=\
                    math.sqrt((np.linalg.norm(aft_land-std_land)**2)/num_land)\
                    /math.sqrt((np.linalg.norm(bf_land-std_land)**2)/num_land)
                
                p=int(rate*10)
                if (p>=max_cnt):
                    cnt[-1]+=1
                else:
                    cnt[p]+=1
                    
    print(cnt)    
    pro=cnt.copy()    
    s=pro.sum()
    
    pro=pro/s
    for i in range(1,pro.shape[0]):
        pro[i]+=pro[i-1]
    print(pro)