new create_output function that accepts a numpy array and does the wo…

…rk of creating fits, jpgs, and storing the output for you

datalab/datalab_session/data_operations/data_operation.py

@@ -85,7 +85,8 @@ def set_operation_progress(self, percent_completed: float):
     def get_operation_progress(self) -> float:
         return cache.get(f'operation_{self.cache_key}_progress', 0.0)
 
-    def set_output(self, output_data: dict):
+    def set_output(self, output):
+        output_data = {'output_files': output if isinstance(output, list) else [output]}
         self.set_status('COMPLETED')
         self.set_operation_progress(1.0)
         cache.set(f'operation_{self.cache_key}_output', output_data, CACHE_DURATION)

datalab/datalab_session/data_operations/long.py

@@ -44,7 +44,4 @@ def operate(self):
             sleep(per_image_timeout)
             self.set_operation_progress((i+1) / num_files)
         # Done "processing" the files so set the output which sets the final status
-        output = {
-            'output_files': self.input_data.get('input_files', [])
-        }
-        self.set_output(output)
+        self.set_output(self.input_data.get('input_files', []))

datalab/datalab_session/data_operations/median.py

@@ -4,8 +4,7 @@
 
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.file_utils import create_fits, crop_arrays, create_jpgs
-from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
+from datalab.datalab_session.file_utils import crop_arrays, create_output
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)
@@ -58,13 +57,7 @@ def operate(self):
         # using the numpy library's median method
         median = np.median(stacked_data, axis=2)
 
-        fits_file = create_fits(self.cache_key, median)
-
-        large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_file)
-
-        output_file = save_fits_and_thumbnails(self.cache_key, fits_file, large_jpg_path, small_jpg_path)
-
-        output =  {'output_files': [output_file]}
+        output = create_output(self.cache_key, median)
 
         self.set_output(output)
         log.info(f'Median output: {self.get_output()}')

datalab/datalab_session/data_operations/noop.py

@@ -42,7 +42,4 @@ def wizard_description():
 
     def operate(self):
         print("No-op triggered!")
-        output = {
-            'output_files': self.input_data.get('input_files', [])
-        }
-        self.set_output(output)
+        self.set_output(self.input_data.get('input_files', []))

datalab/datalab_session/data_operations/normalization.py

@@ -3,8 +3,7 @@
 import numpy as np
 
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
-from datalab.datalab_session.file_utils import create_fits, create_jpgs
-from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
+from datalab.datalab_session.file_utils import create_output
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)
@@ -53,14 +52,10 @@ def operate(self):
             median = np.median(image)
             normalized_image = image / median
 
-            fits_file = create_fits(self.cache_key, normalized_image)
-            large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_file)
-            output_file = save_fits_and_thumbnails(self.cache_key, fits_file, large_jpg_path, small_jpg_path, index=index)
-            output_files.append(output_file)
+            output = create_output(self.cache_key, normalized_image, index=index)
+            output_files.append(output)
 
             self.set_operation_progress(self.get_operation_progress() + .40 * (index + 1) / len(input))
-
-        output =  {'output_files': output_files}
 
-        self.set_output(output)
+        self.set_output(output_files)
         log.info(f'Normalization output: {self.get_output()}')

datalab/datalab_session/data_operations/rgb_stack.py

@@ -5,8 +5,7 @@
 
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.file_utils import get_fits, crop_arrays, create_fits, create_jpgs
-from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
+from datalab.datalab_session.file_utils import create_output, get_fits, crop_arrays, create_jpgs
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)
@@ -59,29 +58,26 @@ def wizard_description():
     def operate(self):
         rgb_input_list = self.input_data['red_input'] + self.input_data['green_input'] + self.input_data['blue_input']
 
-        if len(rgb_input_list) == 3:
-            log.info(f'Executing RGB Stack operation on files: {rgb_input_list}')
+        if len(rgb_input_list) != 3:
+            raise ClientAlertException('RGB stack requires exactly 3 files')
+
+        log.info(f'Executing RGB Stack operation on files: {rgb_input_list}')
 
-            fits_paths = []
-            for file in rgb_input_list:
-                fits_paths.append(get_fits(file.get('basename')))
-                self.set_operation_progress(self.get_operation_progress() + 0.2)
-            
-            large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_paths, color=True)
+        fits_paths = []
+        for file in rgb_input_list:
+            fits_paths.append(get_fits(file.get('basename')))
+            self.set_operation_progress(self.get_operation_progress() + 0.2)
+
+        large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_paths, color=True)
 
-            # color photos take three files, so we store it as one fits file with a 3d SCI ndarray
-            arrays = [fits.open(file)['SCI'].data for file in fits_paths]
-            cropped_data_list = crop_arrays(arrays)
-            stacked_data = np.stack(cropped_data_list, axis=2)
-
-            fits_file = create_fits(self.cache_key, stacked_data)
+        # color photos take three files, so we store it as one fits file with a 3d SCI ndarray
+        arrays = [fits.open(file)['SCI'].data for file in fits_paths]
+        cropped_data_list = crop_arrays(arrays)
+        stacked_data = np.stack(cropped_data_list, axis=2)
 
-            output_file = save_fits_and_thumbnails(self.cache_key, fits_file, large_jpg_path, small_jpg_path)
-
-            output =  {'output_files': [output_file]}
-        else:
-            output = {'output_files': []}
-            raise ClientAlertException('RGB Stack operation requires exactly 3 input files')
+        self.set_operation_progress(0.8)
+
+        output =  create_output(self.cache_key, stacked_data, large_jpg=large_jpg_path, small_jpg=small_jpg_path)
 
         self.set_operation_progress(1.0)
         self.set_output(output)

datalab/datalab_session/data_operations/stacking.py

@@ -4,8 +4,7 @@
 
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.file_utils import create_fits, crop_arrays, create_jpgs
-from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
+from datalab.datalab_session.file_utils import create_output, crop_arrays
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)
@@ -64,13 +63,7 @@ def operate(self):
 
         self.set_operation_progress(0.8)
 
-        fits_file = create_fits(self.cache_key, stacked_sum)
-
-        large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_file)
-
-        output_file = save_fits_and_thumbnails(self.cache_key, fits_file, large_jpg_path, small_jpg_path)
-
-        output =  {'output_files': [output_file]}
+        output = create_output(self.cache_key, stacked_sum)
 
         self.set_output(output)
         log.info(f'Stacked output: {self.get_output()}')

datalab/datalab_session/data_operations/subtraction.py

@@ -4,8 +4,7 @@
 
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.file_utils import create_fits, create_jpgs, crop_arrays
-from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
+from datalab.datalab_session.file_utils import crop_arrays, create_output
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)
@@ -76,15 +75,9 @@ def operate(self):
 
             difference_array = np.subtract(input_image, subtraction_image)
 
-            fits_file = create_fits(self.cache_key, difference_array)
-            large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, fits_file)
-
-            output_file = save_fits_and_thumbnails(self.cache_key, fits_file, large_jpg_path, small_jpg_path, index)
-            outputs.append(output_file)
+            outputs.append(create_output(self.cache_key, difference_array, index=index))
 
             self.set_operation_progress(self.get_operation_progress() + .50 * (index + 1) / len(input_files))
 
-        output =  {'output_files': outputs}
-
-        self.set_output(output)
+        self.set_output(outputs)
         log.info(f'Subtraction output: {self.get_output()}')

datalab/datalab_session/file_utils.py

@@ -6,7 +6,7 @@
 from fits2image.conversions import fits_to_jpg, fits_to_tif
 
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.s3_utils import get_fits
+from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
 
 log = logging.getLogger()
 log.setLevel(logging.INFO)
@@ -108,3 +108,18 @@ def scale_points(height_1: int, width_1: int, height_2: int, width_2: int, x_poi
     x_points = width_2 - x_points
 
   return x_points, y_points
+
+def create_output(cache_key, np_array=None, fits_file=None, large_jpg=None, small_jpg=None, index=None):
+  """
+  A more automated way of creating output for a dev
+  Dev can specify just a cache_key and np array and the function will create the fits and jpgs
+  or the dev can pass the fits_file or jpgs and the function will save them
+  """
+
+  if np_array is not None and fits_file is None:
+    fits_file = create_fits(cache_key, np_array)
+
+  if not large_jpg or not small_jpg:
+    large_jpg, small_jpg = create_jpgs(cache_key, fits_file)
+
+  return save_fits_and_thumbnails(cache_key, fits_file, large_jpg, small_jpg, index)