From b59bdddd17455f77870bd6588f4e36426a156b5f Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Mon, 29 Sep 2025 20:40:35 +0000
Subject: [PATCH] feat: normalized clipped brightness

---
 src/aind_exaspim_image_compression/inference.py              | 5 +++++
 .../machine_learning/data_handling.py                        | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/src/aind_exaspim_image_compression/inference.py b/src/aind_exaspim_image_compression/inference.py
index 9c3203a..4a74dad 100644
--- a/src/aind_exaspim_image_compression/inference.py
+++ b/src/aind_exaspim_image_compression/inference.py
@@ -25,6 +25,7 @@ def predict(
     model,
     batch_size=32,
     normalization_percentiles=(0.5, 99.9),
+    normalized_brightness_clip=7,
     patch_size=64,
     overlap=12,
     trim=5,
@@ -45,6 +46,9 @@ def predict(
     normalization_percentiles : Tuple[int], optional
         Lower and upper percentiles used for normalization. Default is
         (0.5, 99.9).
+    normalized_brightness_clip : float, optional
+        Brightness value used as an upper limit that normalized intensities
+        are clipped to. Default is 10.
     patch_size : int, optional
         Size of the cubic patch extracted from the image. Default is 64.
     overlap : int, optional
@@ -262,6 +266,7 @@ def load_model(path, device="cuda"):
     """
     model = UNet()
     model.load_state_dict(torch.load(path, map_location=device))
+    model.to(device)
     model.eval()
     return model
 
diff --git a/src/aind_exaspim_image_compression/machine_learning/data_handling.py b/src/aind_exaspim_image_compression/machine_learning/data_handling.py
index fd121a3..59509b9 100644
--- a/src/aind_exaspim_image_compression/machine_learning/data_handling.py
+++ b/src/aind_exaspim_image_compression/machine_learning/data_handling.py
@@ -476,6 +476,9 @@ def __init__(
         normalization_percentiles : Tuple[float], optional
             Upper and lower percentiles used to normalize the input image.
             Default is (0.5, 99.5).
+        normalized_brightness_clip : float, optional
+            Brightness value used as an upper limit that normalized intensities
+            are clipped to. Default is 10.
         sigma_bm4d : float, optional
             Smoothing parameter used in the BM4D denoising algorithm. Default
             is 16.