| 50 | return step('test', opt, actions, val_loader, model, steps=steps) |
| 51 | |
| 52 | def step(split, args, actions, dataLoader, model, optimizer=None, epoch=None, steps=None): |
| 53 | |
| 54 | loss_all = {'loss': AccumLoss()} |
| 55 | |
| 56 | model_3d = model['CFM'] |
| 57 | if split == 'train': |
| 58 | model_3d.train() |
| 59 | else: |
| 60 | model_3d.eval() |
| 61 | |
| 62 | # determine steps for single-step evaluation per call |
| 63 | steps_to_use = steps |
| 64 | p1_error_sum = 0. |
| 65 | p2_error_sum = 0. |
| 66 | |
| 67 | data_lent = 0 |
| 68 | |
| 69 | for i, data in enumerate(tqdm(dataLoader, 0)): |
| 70 | |
| 71 | # batch_cam, gt_3D, input_2D, action, subject, cam_ind, vis_3D, start_3d, end_3d = data |
| 72 | input_2D, gt_3D = data['keypoints_2d'], data['keypoints_3d'] |
| 73 | # print(input_2D.shape,input_2D) |
| 74 | # print(gt_3D) |
| 75 | # input_2D shape: torch.Size([B, J, 2]) (normalized x,y coordinates) |
| 76 | # gt_3D shape: torch.Size([B, J, 4]) (x,y,z + homogeneous coordinate) |
| 77 | gt_3D = gt_3D[:,:,:3] # only use x,y,z for 3D ground truth |
| 78 | |
| 79 | # [input_2D, gt_3D, batch_cam, vis_3D] = get_variable(split, [input_2D, gt_3D, batch_cam, vis_3D]) |
| 80 | |
| 81 | # unsqueeze frame dimension |
| 82 | input_2D = input_2D.unsqueeze(1) # (B,F,J,C) |
| 83 | gt_3D = gt_3D.unsqueeze(1) # (B,F,J,C) |
| 84 | |
| 85 | device = next(model_3d.parameters()).device |
| 86 | |
| 87 | model_dtype = next(model_3d.parameters()).dtype |
| 88 | input_2D = input_2D.to(device=device, dtype=model_dtype) |
| 89 | gt_3D = gt_3D.to(device=device, dtype=model_dtype) |
| 90 | |
| 91 | B = input_2D.shape[0] |
| 92 | data_lent += B |
| 93 | |
| 94 | if split =='train': |
| 95 | B, F, J, C = input_2D.shape |
| 96 | |
| 97 | # Note: gt_3D is already root-relative from the dataloader |
| 98 | # Root joint should already be [0,0,0] |
| 99 | gt_3D = gt_3D.clone() |
| 100 | gt_3D[:, :, args.root_joint] = 0 |
| 101 | |
| 102 | # Conditional Flow Matching training |
| 103 | # gt_3D, input_2D shape: (B,F,J,C) |
| 104 | # vis_3D shape: (B,F,J,1) - visibility mask |
| 105 | # x0_noise = torch.randn_like(gt_3D) |
| 106 | x0_noise = torch.randn(B, F, J, 3, device=gt_3D.device, dtype=model_dtype) |
| 107 | x0 = x0_noise |
| 108 | |
| 109 | B = gt_3D.size(0) |