Update server to have raw-data analysis operation and update rgb oper…

datalab/datalab_session/analysis/raw_data.py

@@ -0,0 +1,33 @@
+import numpy as np
+from PIL import Image
+from datalab.datalab_session.s3_utils import get_fits
+from datalab.datalab_session.file_utils import get_hdu
+from fits2image.scaling import extract_samples, calc_zscale_min_max
+
+# TODO: This analysis endpoint assumes the image to be of 16 bitdepth. We should make this agnositc to bit depth in the future
+
+def raw_data(input: dict):
+    fits_path = get_fits(input['basename'], input.get('source', 'archive'))
+
+    sci_hdu = get_hdu(fits_path, 'SCI')
+    image_data = sci_hdu.data
+
+    # 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)
+
+    # resize the image to max. 500 pixels on an axis by default for the UI
+    max_size = input.get('max_size', 500)
+    image = Image.fromarray(image_data)
+    newImage = image.resize((max_size, max_size), Image.LANCZOS)
+    scaled_array = np.asarray(newImage).astype(np.float16)
+    scaled_array_flipped = np.flip(scaled_array, axis=0)
+
+    return {'data': scaled_array_flipped.flatten().tolist(),
+            'height': scaled_array.shape[0],
+            'width': scaled_array.shape[1],
+            'zmin': int(median),
+            'zmax': int(zmax),
+            'bitdepth': 16
+        }

datalab/datalab_session/data_operations/rgb_stack.py

@@ -36,6 +36,8 @@ def wizard_description():
                     'type': 'file',
                     'minimum': 1,
                     'maximum': 1,
+                    'include_custom_scale': True,
+                    'combine_custom_scale': 'rgb',
                     'filter': ['rp', 'r']
                 },
                 'green_input': {
@@ -44,6 +46,8 @@ def wizard_description():
                     'type': 'file',
                     'minimum': 1,
                     'maximum': 1,
+                    'include_custom_scale': True,
+                    'combine_custom_scale': 'rgb',
                     'filter': ['V', 'gp']
                 },
                 'blue_input': {
@@ -52,9 +56,11 @@ def wizard_description():
                     'type': 'file',
                     'minimum': 1,
                     'maximum': 1,
+                    'include_custom_scale': True,
+                    'combine_custom_scale': 'rgb',
                     'filter': ['B']
                 }
-            }
+            },
         }
 
     def operate(self):
@@ -64,12 +70,16 @@ def operate(self):
         log.info(rgb_comment)
 
         input_fits_list = []
+        zmin_list = []
+        zmax_list = []
         for index, input in enumerate(rgb_input_list, start=1):
             input_fits_list.append(InputDataHandler(input['basename'], input['source']))
+            zmin_list.append(input['zmin'])
+            zmax_list.append(input['zmax'])
             self.set_operation_progress(0.4 * (index / len(rgb_input_list)))
 
         fits_file_list = [image.fits_file for image in input_fits_list]
-        large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_file_list, color=True)
+        large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_file_list, color=True, zmin=zmin_list, zmax=zmax_list)
         self.set_operation_progress(0.6)
 
         # color photos take three files, so we store it as one fits file with a 3d SCI ndarray

datalab/datalab_session/file_utils.py

@@ -55,7 +55,7 @@ def create_tif(key: str, fits_path: np.ndarray) -> str:
 
   return tif_path
 
-def create_jpgs(cache_key, fits_paths: str, color=False) -> list:
+def create_jpgs(cache_key, fits_paths: str, color=False, zmin=None, zmax=None) -> list:
     """
     Create jpgs from fits files and save them to S3
     If using the color option fits_paths need to be in order R, G, B
@@ -71,8 +71,8 @@ def create_jpgs(cache_key, fits_paths: str, color=False) -> list:
 
     max_height, max_width = max(get_fits_dimensions(path) for path in fits_paths)
 
-    fits_to_jpg(fits_paths, large_jpg_path, width=max_width, height=max_height, color=color)
-    fits_to_jpg(fits_paths, thumbnail_jpg_path, color=color)
+    fits_to_jpg(fits_paths, large_jpg_path, width=max_width, height=max_height, color=color, zmin=zmin, zmax=zmax)
+    fits_to_jpg(fits_paths, thumbnail_jpg_path, color=color, zmin=zmin, zmax=zmax)
 
     return large_jpg_path, thumbnail_jpg_path
 

datalab/datalab_session/tests/test_operations.py

@@ -237,9 +237,9 @@ def test_operate(self, mock_get_fits, mock_named_tempfile, mock_create_jpgs, moc
         mock_save_fits_and_thumbnails.return_value = self.temp_rgb_path
 
         input_data = {
-            'red_input': [{'basename': 'red_fits', 'source': 'local'}],
-            'green_input': [{'basename': 'green_fits', 'source': 'local'}],
-            'blue_input': [{'basename': 'blue_fits', 'source': 'local'}]
+            'red_input': [{'basename': 'red_fits', 'source': 'local', 'zmin': 0, 'zmax': 255}],
+            'green_input': [{'basename': 'green_fits', 'source': 'local', 'zmin': 0, 'zmax': 255}],
+            'blue_input': [{'basename': 'blue_fits', 'source': 'local', 'zmin': 0, 'zmax': 255}]
         }
 
         rgb = RGB_Stack(input_data)

datalab/datalab_session/views.py

@@ -6,6 +6,8 @@
 from datalab.datalab_session.analysis.line_profile import line_profile
 from datalab.datalab_session.analysis.source_catalog import source_catalog
 from datalab.datalab_session.analysis.get_tif import get_tif
+from datalab.datalab_session.analysis.raw_data import raw_data
+
 
 class OperationOptionsApiView(RetrieveAPIView):
     """ View to retrieve the set of operations available, for the UI to use """
@@ -33,6 +35,8 @@ def post(self, request, action):
                 output = source_catalog(input)
             case 'get-tif':
                 output = get_tif(input)
+            case 'raw-data':
+                output = raw_data(input)
             case _:
                 raise Exception(f'Analysis action {action} not found')