# Inference
tex, pred = model(iuv)
# Composite
pred = torch.where(mask > 0, pred, image)
# Loss
loss_rgb = masked_l1(tex[:, :3], image, mask.float())
loss_l1 = masked_l1(pred, image, mask.float())
loss_lpips = lpips(pred, image)
loss = loss_rgb + loss_l1 + loss_lpips
# Backpropagation
loss.backward()
optimizer.step()
optimizer.zero_grad()
# Logging
output = pred[0].detach().cpu().numpy().transpose(1, 2, 0)
cv2.imwrite(f"results/{e}_{i}.jpg", (255 * output).astype(np.uint8))
wandb.log({
"loss_rgb": loss_rgb.item(),
"loss_l1": loss_l1.item(),
"loss_lpips": loss_lpips.item(),
"loss": loss.item(),
})
torch.save(model.state_dict(), f"breatheless_v1.pth")
