Merge pull request #33 from LCOGT/feat/add-subtraction-operation

Add subtraction operation, update stack_arrays util to just crop_arrays, fixed basename output bug

datalab/datalab_session/data_operations/median.py

@@ -4,7 +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, stack_arrays, create_jpgs
+from datalab.datalab_session.file_utils import create_fits, crop_arrays, create_jpgs
 from datalab.datalab_session.s3_utils import save_fits_and_thumbnails
 
 log = logging.getLogger()
@@ -51,7 +51,8 @@ def operate(self):
 
         image_data_list = self.get_fits_npdata(input, percent=0.4, cur_percent=0.0)
 
-        stacked_data = stack_arrays(image_data_list)
+        cropped_data_list = crop_arrays(image_data_list)
+        stacked_data = np.stack(cropped_data_list, axis=2)
 
         # using the numpy library's median method
         median = np.median(stacked_data, axis=2)
@@ -64,6 +65,5 @@ def operate(self):
 
         output =  {'output_files': [output_file]}
 
-        self.set_percent_completion(1.0)
         self.set_output(output)
         log.info(f'Median output: {self.get_output()}')

datalab/datalab_session/data_operations/rgb_stack.py

@@ -1,10 +1,11 @@
 import logging
 
 from astropy.io import fits
+import numpy as np
 
 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, stack_arrays, create_fits, create_jpgs
+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
 
 log = logging.getLogger()
@@ -70,7 +71,9 @@ def operate(self):
 
             # 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]
-            stacked_data = stack_arrays(arrays)
+            cropped_data_list = crop_arrays(arrays)
+            stacked_data = np.stack(cropped_data_list, axis=2)
+
             fits_file = create_fits(self.cache_key, stacked_data)
 
             output_file = save_fits_and_thumbnails(self.cache_key, fits_file, large_jpg_path, small_jpg_path)

datalab/datalab_session/data_operations/subtraction.py

@@ -0,0 +1,90 @@
+import logging
+
+import numpy as np
+
+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
+
+log = logging.getLogger()
+log.setLevel(logging.INFO)
+
+
+class Subtraction(BaseDataOperation):
+
+    @staticmethod
+    def name():
+        return 'Subtraction'
+
+    @staticmethod
+    def description():
+        return """
+          The Subtraction operation takes in 1..n input images and calculated the subtraction value pixel-by-pixel.
+          The output is a subtraction image for the n input images. This operation is commonly used for background subtraction.
+        """
+
+    @staticmethod
+    def wizard_description():
+        return {
+            'name': Subtraction.name(),
+            'description': Subtraction.description(),
+            'category': 'image',
+            'inputs': {
+                'input_files': {
+                    'name': 'Input Files',
+                    'description': 'The input files to operate on',
+                    'type': 'file',
+                    'minimum': 1,
+                    'maximum': 999
+                },
+                'subtraction_file': {
+                    'name': 'Subtraction File',
+                    'description': 'This file will be subtracted from the input images.',
+                    'type': 'file',
+                    'minimum': 1,
+                    'maximum': 1
+                }
+            },
+        }
+
+    def operate(self):
+
+        input_files = self.input_data.get('input_files', [])
+        print(f'Input files: {input_files}')
+        subtraction_file_input = self.input_data.get('subtraction_file', [])
+        print(f'Subtraction file: {subtraction_file_input}')
+
+        if not subtraction_file_input:
+            raise ClientAlertException('Missing a subtraction file')
+
+        if len(input_files) < 1:
+            raise ClientAlertException('Need at least one input file')
+
+        log.info(f'Executing subtraction operation on {len(input_files)} files')
+
+        input_image_data_list = self.get_fits_npdata(input_files)
+        self.set_percent_completion(.30)
+
+        subtraction_image = self.get_fits_npdata(subtraction_file_input)[0]
+        self.set_percent_completion(.40)
+
+        outputs = []
+        for index, input_image in enumerate(input_image_data_list):
+            # crop the input_image and subtraction_image to the same size
+            input_image, subtraction_image = crop_arrays([input_image, subtraction_image])
+
+            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)
+
+            self.set_percent_completion(self.get_percent_completion() + .50 * (index + 1) / len(input_files))
+
+        output =  {'output_files': outputs}
+
+        self.set_output(output)
+        log.info(f'Subtraction output: {self.get_output()}')

datalab/datalab_session/file_utils.py

@@ -79,7 +79,7 @@ def create_jpgs(cache_key, fits_paths: str, color=False) -> list:
 
     return large_jpg_path, thumbnail_jpg_path
 
-def stack_arrays(array_list: list):
+def crop_arrays(array_list: list):
   """
   Takes a list of numpy arrays from fits images and stacks them to be a 3d numpy array
   cropped since fits images can be different sizes
@@ -88,8 +88,7 @@ def stack_arrays(array_list: list):
   min_y = min(arr.shape[1] for arr in array_list)
 
   cropped_data_list = [arr[:min_x, :min_y] for arr in array_list]
-
-  return np.stack(cropped_data_list, axis=2)
+  return cropped_data_list
 
 def scale_points(height_1: int, width_1: int, height_2: int, width_2: int, x_points=[], y_points=[], flip_y = False, flip_x = False):
   """

datalab/datalab_session/s3_utils.py

@@ -152,7 +152,7 @@ def save_fits_and_thumbnails(cache_key, fits_path, large_jpg_path, thumbnail_jpg
         'fits_url': fits_url,
         'large_url': large_jpg_url,
         'thumbnail_url': thumbnail_jpg_url,
-        'basename': f'{cache_key}',
+        'basename': f'{cache_key}-{index}' if index else cache_key,
         'source': 'datalab'}
     )
 

datalab/datalab_session/tests/test_utils.py

@@ -53,11 +53,10 @@ def test_stack_arrays(self):
     test_array_1 = np.zeros((10, 20))
     test_array_2 = np.ones((20, 10))
 
-    stacked_array = stack_arrays([test_array_1, test_array_2])
-    self.assertIsInstance(stacked_array, np.ndarray)
-    self.assertEqual(stacked_array.shape, (10, 10, 2))
-    self.assertEqual(stacked_array[:, :, 0].tolist(), np.zeros((10, 10)).tolist())
-    self.assertEqual(stacked_array[:, :, 1].tolist(), np.ones((10, 10)).tolist())
+    cropped_array = crop_arrays([test_array_1, test_array_2])
+    self.assertEqual(len(cropped_array), 2)
+    self.assertEqual(cropped_array[0].tolist(), np.zeros((10, 10)).tolist())
+    self.assertEqual(cropped_array[1].tolist(), np.ones((10, 10)).tolist())
 
   def test_scale_points(self):
     x_points = [1, 2, 3]