Skip to content

Commit

Permalink
feat: Add alpha loss
Browse files Browse the repository at this point in the history
  • Loading branch information
hugoycj committed Oct 16, 2024
1 parent dbcd181 commit 6b63529
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,13 @@ def training(dataset, opt, pipe, testing_iterations, saving_iterations, checkpoi
Ll1 = l1_loss_appearance(image, gt_image, appearances, viewpoint_idx) # use L1 loss for the transformed image if using decoupled appearance
loss = (1.0 - opt.lambda_dssim) * Ll1 + opt.lambda_dssim * (1.0 - ssim(image, gt_image))

# alpha loss
if opt.lambda_mask > 0:
opacity = 1 - render_pkg["rend_alpha"].clamp(1e-6, 1-1e-6)
bg = 1 - viewpoint_cam.gt_alpha_mask
mask_error = (- bg * torch.log(opacity)).mean()
loss += opt.lambda_mask * mask_error

# regularization
lambda_normal = opt.lambda_normal if iteration > 7000 else 0.0
lambda_depth = opt.propagation_begin if iteration > opt.propagation_begin else 0.0
Expand Down

0 comments on commit 6b63529

Please sign in to comment.