Merge pull request #49 from LCOGT/feature/fits-align

Using fits_align package with RGB stack operation

datalab/datalab_session/data_operations/input_data_handler.py

@@ -33,7 +33,7 @@ def __init__(self, basename: str, source: str = None) -> None:
 
   def __str__(self) -> str:
     with fits.open(self.fits_file) as hdul:
-      return f"{self.basename}@{self.fits_file}\nHDU List\n{self.hdul.info()}"
+      return f"{self.basename}@{self.fits_file}\nHDU List\n{hdul.info()}"
 
   def get_hdu(self, extension: str=None):
     """Return an HDU from the FITS file.

datalab/datalab_session/data_operations/rgb_stack.py

@@ -1,7 +1,10 @@
 import logging
 
 import numpy as np
+from fits_align.ident import make_transforms
+from fits_align.align import affineremap
 
+from django.conf import settings
 from datalab.datalab_session.data_operations.input_data_handler import InputDataHandler
 from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
 from datalab.datalab_session.data_operations.fits_output_handler import FITSOutputHandler
@@ -11,8 +14,13 @@
 log = logging.getLogger()
 log.setLevel(logging.INFO)
 
-
 class RGB_Stack(BaseDataOperation):
+    REQUIRED_INPUTS = 3
+    PROGRESS_STEPS = {
+        'INPUT_PROCESSING': 0.4,
+        'ALIGNMENT': 0.6,
+        'STACKING': 0.8
+    }
 
     @staticmethod
     def name():
@@ -62,32 +70,82 @@ 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: raise ClientAlertException('RGB stack requires exactly 3 files')
-        rgb_comment = f'Datalab RGB Stack on files {", ".join([image["basename"] for image in rgb_input_list])}'
-        log.info(rgb_comment)
+    def _validate_inputs(self):
+        rgb_input_list = []
+        for color in ['red_input', 'green_input', 'blue_input']:
+            input_data = self.input_data[color][0]
+            if not input_data:
+                raise ClientAlertException(f'Missing {color}')
+            rgb_input_list.append(input_data)
 
+        if len(self.input_data) != self.REQUIRED_INPUTS:
+            raise ClientAlertException(f'RGB stack requires exactly {self.REQUIRED_INPUTS} files')
+
+        return rgb_input_list
+
+    def _process_inputs(self, rgb_input_list) -> tuple[list[InputDataHandler], list[float], list[float]]:
         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)))
+            self.set_operation_progress(self.PROGRESS_STEPS['INPUT_PROCESSING'] * (index / len(rgb_input_list)))
+
+        return input_fits_list, zmin_list, zmax_list
+
+    def _align_images(self, fits_files: list[str]) -> list[str]:
+        ref_image = fits_files[0]
+        images_to_align = fits_files[1:]
+        identifications = make_transforms(ref_image, images_to_align)
+
+        aligned_images = [ref_image]
+        for id in identifications:
+            if id.ok:
+                aligned_img = affineremap(id.ukn.filepath, id.trans, outdir=settings.TEMP_FITS_DIR)
+                aligned_images.append(aligned_img)
+
+        if len(aligned_images) != self.REQUIRED_INPUTS:
+            raise ClientAlertException('Failed to align all images')
+
+        self.set_operation_progress(self.PROGRESS_STEPS['ALIGNMENT'])
 
-        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, zmin=zmin_list, zmax=zmax_list)
-        self.set_operation_progress(0.6)
+        return aligned_images
+
+    # Currently storing the output fits SCI HDU as a 3D ndarray consisting of each input's SCI data
+    def _create_3d_array(self, input_handlers: list[InputDataHandler]) -> np.ndarray:
+        sci_data_list = [image.sci_data for image in input_handlers]
 
-        # color photos take three files, so we store it as one fits file with a 3d SCI ndarray
-        sci_data_list = [image.sci_data for image in input_fits_list]
         cropped_data_list = crop_arrays(sci_data_list)
+
         stacked_ndarray = np.stack(cropped_data_list, axis=2)
-        self.set_operation_progress(0.8)
+        self.set_operation_progress(self.PROGRESS_STEPS['STACKING'])
+
+        return stacked_ndarray
+
+    def operate(self):
+        rgb_inputs = self._validate_inputs()
+
+        input_handlers, zmin_list, zmax_list = self._process_inputs(rgb_inputs)
+
+        fits_files = [handler.fits_file for handler in input_handlers]
+
+        aligned_images = self._align_images(fits_files)
+
+        large_jpg_path, small_jpg_path = create_jpgs(self.cache_key, aligned_images, color=True, zmin=zmin_list, zmax=zmax_list)
+
+        stacked_ndarray = self._create_3d_array(input_handlers)
 
-        output = FITSOutputHandler(self.cache_key, stacked_ndarray, rgb_comment).create_and_save_data_products(large_jpg_path=large_jpg_path, small_jpg_path=small_jpg_path)
+        rgb_comment = f'Datalab RGB Stack on files {", ".join(input["basename"] for input in rgb_inputs)}'
+        output = FITSOutputHandler(
+            self.cache_key, 
+            stacked_ndarray, 
+            rgb_comment
+        ).create_and_save_data_products(
+            large_jpg_path=large_jpg_path, 
+            small_jpg_path=small_jpg_path
+        )
 
         log.info(f'RGB Stack output: {output}')
         self.set_output(output)