From 61df764988abcca4198ddb1736a1325f09379ddb Mon Sep 17 00:00:00 2001
From: Jon <jnation@lco.global>
Date: Sat, 16 Nov 2024 00:01:56 -0100
Subject: [PATCH] Add in histogram info to the raw_image endpoint to display
 the histogram slider

---
 datalab/datalab_session/analysis/raw_data.py | 36 ++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/datalab/datalab_session/analysis/raw_data.py b/datalab/datalab_session/analysis/raw_data.py
index a355252..38c43b9 100644
--- a/datalab/datalab_session/analysis/raw_data.py
+++ b/datalab/datalab_session/analysis/raw_data.py
@@ -1,4 +1,5 @@
 import numpy as np
+import math
 from PIL import Image
 from datalab.datalab_session.s3_utils import get_fits
 from datalab.datalab_session.file_utils import get_hdu
@@ -15,7 +16,7 @@ def raw_data(input: dict):
     # Compute the fits2image autoscale params to send with the image
     samples = extract_samples(image_data, sci_hdu.header, 2000)
     median = np.median(samples)
-    _, zmax, _ = calc_zscale_min_max(samples, contrast=0.1, iterations=1)
+    zmin, zmax, _ = calc_zscale_min_max(samples, contrast=0.1, iterations=1)
 
     # resize the image to max. 500 pixels on an axis by default for the UI
     max_size = input.get('max_size', 500)
@@ -25,17 +26,48 @@ def raw_data(input: dict):
     match bitpix:
         case 8:
             datatype = np.uint8
+            max_value = np.iinfo(datatype).max
         case 16:
             datatype = np.float16
+            max_value = np.finfo(datatype).max
         case 32:
             datatype = np.float32
+            max_value = np.finfo(datatype).max
     scaled_array = np.asarray(newImage).astype(datatype)
     scaled_array_flipped = np.flip(scaled_array, axis=0)
 
+    # Here we do a crazy histogram scaling to stretch the points in between zmin and zmax since that is where most detail is
+    # We have 10 bins before zmin, 100 between zmin and zmax and 10 after zmax.
+    lower_bound = int(zmin * 0.8)  # Increase resolution slightly below zmin
+    upper_bound = int(zmax*1.2)  # Increase resolution slightly beyond zmax
+    lower_step = int(lower_bound / 10)
+    upper_step = int((max_value - upper_bound) / 10)
+    step = int((upper_bound - lower_bound) / 100)
+    bins = np.arange(0, lower_bound, lower_step).tolist()
+    bins += np.arange(lower_bound, upper_bound, step).tolist()
+    bins += np.arange(upper_bound, max_value, upper_step).tolist()
+    histogram, bin_edges = np.histogram(samples, bins=bins)
+    bin_middles = []
+    previous_edge = 0
+    for edge in bin_edges:
+        if edge != 0:
+            bin_middles.append(previous_edge + int((edge-previous_edge) / 2.0))
+        previous_edge = edge
+
+    # Using np.log10 on the histogram made some wild results, so just apply log10 to each value
+    hist = []
+    for h in histogram:
+        if h > 0:
+            hist.append(math.log10(h))
+        else:
+            hist.append(0)
+
     return {'data': scaled_array_flipped.flatten().tolist(),
             'height': scaled_array.shape[0],
             'width': scaled_array.shape[1],
+            'histogram': hist,
+            'bins': bin_middles,
             'zmin': int(median),
             'zmax': int(zmax),
-            'bitdepth': 16
+            'bitdepth': bitpix
         }