Inference the Pose Regression Network Inputs: args: parsed argument data: Input image in shape (batchsize, channels, H, W) device: gpu device model: PoseNet model Outputs: pose: Predicted Pose in shape (batchsize, 3, 4)
(args, data, device, model)
| 48 | return model |
| 49 | |
| 50 | def inference_pose_regression(args, data, device, model): |
| 51 | """ |
| 52 | Inference the Pose Regression Network |
| 53 | Inputs: |
| 54 | args: parsed argument |
| 55 | data: Input image in shape (batchsize, channels, H, W) |
| 56 | device: gpu device |
| 57 | model: PoseNet model |
| 58 | Outputs: |
| 59 | pose: Predicted Pose in shape (batchsize, 3, 4) |
| 60 | """ |
| 61 | inputs = data.to(device) |
| 62 | if args.preprocess_ImgNet: |
| 63 | inputs = preprocess_data(inputs, device) |
| 64 | predict_pose = model(inputs) |
| 65 | pose = predict_pose.reshape(args.batch_size, 3, 4) |
| 66 | |
| 67 | if args.svd_reg: |
| 68 | # pdb.set_trace() |
| 69 | R_torch = pose[:,:3,:3].clone() # debug |
| 70 | u,s,v=torch.svd(R_torch) |
| 71 | Rs = torch.matmul(u, v.transpose(-2,-1)) |
| 72 | pose[:,:3,:3] = Rs |
| 73 | return pose |
| 74 | |
| 75 | def rgb_loss(rgb, target, extras): |
| 76 | ''' Compute RGB MSE Loss, original from NeRF Paper ''' |
no test coverage detected