From bbd3ddb978b4c78ca093ee5ad8dfa8709ab4fc8b Mon Sep 17 00:00:00 2001
From: Ziwen Liu <ziwen.liu@czbiohub.org>
Date: Fri, 5 Jan 2024 14:56:45 -0800
Subject: [PATCH] use new api in the model

---
 waveorder/models/inplane_oriented_thick_pol3d.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/waveorder/models/inplane_oriented_thick_pol3d.py b/waveorder/models/inplane_oriented_thick_pol3d.py
index 5d3ed46..aefbaaf 100644
--- a/waveorder/models/inplane_oriented_thick_pol3d.py
+++ b/waveorder/models/inplane_oriented_thick_pol3d.py
@@ -4,7 +4,7 @@
 import torch
 from torch import Tensor
 
-from waveorder import background_estimator, stokes, util
+from waveorder import correction, stokes, util
 
 
 def generate_test_phantom(yx_shape):
@@ -125,7 +125,6 @@ def apply_inverse_transfer_function(
 
     # Apply an "Estimated" background correction
     if remove_estimated_background:
-        estimator = background_estimator.BackgroundEstimator2D()
         for stokes_index in range(background_corrected_stokes.shape[0]):
             # Project to 2D
             z_projection = torch.mean(
@@ -134,10 +133,7 @@ def apply_inverse_transfer_function(
             # Estimate the background and subtract
             background_corrected_stokes[
                 stokes_index
-            ] -= estimator.get_background(
-                z_projection,
-                normalize=False,
-            )
+            ] -= correction.estimate_background(z_projection, normalize=False)
 
     # Project to 2D (typically for SNR reasons)
     if project_stokes_to_2d: