Export JPG with AVM (WCS metadata)

CHANGES.rst

@@ -49,6 +49,8 @@ New Features
 - Provide better error reporting when attempting to load data via `load`
   and loaders infrastructure. [#4058]
 
+- Export image viewer screenshots to JPG with Astronomy Visualization Metadata (AVM). [#4065]
+
 Cubeviz
 ^^^^^^^
 

jdaviz/configs/default/aida.py

@@ -166,12 +166,13 @@ def _get_current_fov(self, sky_or_pixel=None):
 
         if sky_or_pixel == 'pixel':
             return pixel_fov
-        if not any(c.strip() for c in wcs.wcs.ctype):
-            raise ValueError("The image must have valid WCS to return `fov` in `sky`.")
 
         if isinstance(wcs, GWCS):
             wcs = WCS(wcs.to_fits_sip())
 
+        if not any(c.strip() for c in wcs.wcs.ctype):
+            raise ValueError("The image must have valid WCS to return `fov` in `sky`.")
+
         # compute the mean of the height and width of the
         # viewer's FOV on ``data`` in world units:
         x_corners = [

jdaviz/configs/default/plugins/export/avm.py

@@ -0,0 +1,76 @@
+
+import os
+import numpy as np
+from PIL import Image
+
+from astropy.wcs.utils import fit_wcs_from_points
+from pyavm import AVM
+
+
+def png_to_jpg_avm(viz, viewer, png_filename):
+    """Convert a PNG screenshot to JPG with Astronomy Visualization Metadata (AVM).
+
+    Parameters
+    ----------
+    viz : jdaviz.configs.imviz.helper.Imviz
+        Imviz helper.
+    viewer : jdaviz.configs.imviz.plugins.viewers.ImvizImageVew
+        Glue viewer exported in the PNG screenshot screenshot
+    png_filename : path-like, str
+        Path to PNG screenshot
+    """
+    # open temporary PNG
+    img = Image.open(png_filename)
+    png_shape = img.size
+
+    # get the viewer limits in refdata pixel coords
+    refdata_wcs = viewer.state.reference_data.coords
+    xy_ref = (
+        np.array([
+            viewer.state.x_min,
+            viewer.state.x_min,
+            viewer.state.x_max,
+            viewer.state.x_max,
+            0.5 * (viewer.state.x_min + viewer.state.x_max)]),
+        np.array([
+            viewer.state.y_min,
+            viewer.state.y_max,
+            viewer.state.y_min,
+            viewer.state.y_max,
+            0.5 * (viewer.state.y_min + viewer.state.y_max)]),
+    )
+    # pixel coordinates for the corners and center of the image
+    xy_png = (
+        np.array([0, 0, png_shape[0], png_shape[0], png_shape[0] / 2]),
+        np.array([0, png_shape[1], 0, png_shape[1], png_shape[1] / 2])
+    )
+
+    # convert observation pixel coords to world
+    world = refdata_wcs.pixel_to_world(*xy_ref)
+
+    # fit for WCS using these five points
+    png_wcs = fit_wcs_from_points(xy_png, world, proj_point=world[0])
+
+    # assemble AVM with this WCS:
+    png_avm = AVM().from_wcs(png_wcs, png_shape)
+
+    # add AVM tags required for the aladin-lite parser:
+    png_avm.Creator = 'jdaviz'
+    png_avm.Rights = ''
+    png_avm.Credit = ''
+    png_avm.ID = ''
+    png_avm.MetadataDate = ''
+
+    jpg_path = str(png_filename).replace('.png', '.jpg')
+
+    try:
+        jpg_path_tmp = str(png_filename).replace('.png', '_tmp.jpg')
+
+        # write out temporary JPG (drop alpha channel)
+        img.convert('RGB').save(jpg_path_tmp)
+
+        # embed AVM into final JPG
+        png_avm.embed(jpg_path_tmp, jpg_path, verify=True)
+    finally:
+        # ensure tmp file gets removed
+        os.remove(jpg_path_tmp)

jdaviz/configs/default/plugins/export/export.py

@@ -11,6 +11,7 @@
 from specutils import Spectrum
 from traitlets import Bool, List, Unicode, observe
 
+from jdaviz.configs.default.plugins.export.avm import png_to_jpg_avm
 from jdaviz.core.custom_traitlets import FloatHandleEmpty, IntHandleEmpty
 from jdaviz.core.marks import ShadowMixin
 from jdaviz.core.registries import tray_registry
@@ -137,7 +138,7 @@ def __init__(self, *args, **kwargs):
 
         self.viewer.add_filter('is_not_empty')
 
-        viewer_format_options = ['png', 'svg']
+        viewer_format_options = ['png', 'svg', 'jpg']
         if self.config == 'cubeviz':
             if not self.app.state.settings.get('server_is_remote'):
                 viewer_format_options += ['mp4']
@@ -443,7 +444,7 @@ def _normalize_filename(self, filename=None, filetype=None, overwrite=False, def
 
     @with_spinner()
     def export(self, filename=None, show_dialog=None, overwrite=False,
-               raise_error_for_overwrite=True):
+               raise_error_for_overwrite=True, embed_avm=False):
         """
         Export selected item(s)
 
@@ -462,6 +463,10 @@ def export(self, filename=None, show_dialog=None, overwrite=False,
             If `True`, raise exception when ``overwrite=False`` but
             output file already exists. Otherwise, a message will be sent
             to application snackbar instead.
+
+        embed_avm : bool
+            If `True` and the file type is 'jpg', embed
+            Astronomy Visualization Metadata (including WCS) using ``pyAVM``.
         """
         if self.multiselect:
             raise NotImplementedError("batch export not yet supported")
@@ -505,6 +510,33 @@ def export(self, filename=None, show_dialog=None, overwrite=False,
                 self.save_movie(viewer, filename, filetype,
                                 width=f"{self.image_width}px" if self.image_custom_size else None,
                                 height=f"{self.image_height}px" if self.image_custom_size else None)
+
+            elif filetype == "jpg":
+                # export screenshot to JPG with AVM by:
+                #   (1) export a temporary PNG
+                #   (2) convert temporary PNG to a temporary JPG with PIL
+                #   (3) use pyAVM to embed AVM into a final copy of the temporary JPG
+
+                try:
+                    # export PNG
+                    tmp_filename = Path(str(Path(filename)).replace('.jpg', '.png'))
+                    self.save_figure(
+                        viewer, tmp_filename, 'png', show_dialog=show_dialog,
+                        width=f"{self.image_width}px" if self.image_custom_size else None,
+                        height=f"{self.image_height}px" if self.image_custom_size else None
+                    )
+
+                    # wait for PNG to be available before continuing
+                    while viewer.figure._upload_png_callback is not None:
+                        time.sleep(0.1)
+
+                    # now convert to JPG with AVM
+                    png_to_jpg_avm(self.app._jdaviz_helper, viewer, tmp_filename)
+
+                finally:
+                    # remove temporary png file, even if there are exceptions
+                    os.remove(tmp_filename)
+
             else:
                 self.save_figure(viewer, filename, filetype, show_dialog=show_dialog,
                                  width=f"{self.image_width}px" if self.image_custom_size else None,
@@ -630,10 +662,11 @@ def save_figure(self, viewer, filename=None, filetype="png", show_dialog=False,
         else:
             app = viewer.app
 
-        def on_img_received(data):
+        def on_img_received(data, filename=filename):
             try:
                 with filename.open(mode='bw') as f:
                     f.write(data)
+
             except Exception as e:
                 self.hub.broadcast(SnackbarMessage(
                     f"{self.viewer.selected} failed to export to {str(filename)}: {e}",
@@ -643,7 +676,7 @@ def on_img_received(data):
                     f"{self.viewer.selected} exported to {str(filename)}",
                     sender=self, color="success"))
 
-        def get_png(figure):
+        def get_png(figure, on_img_received=on_img_received):
             if figure._upload_png_callback is not None:
                 raise ValueError("previous png export is still in progress. Wait to complete before making another call to save_figure")  # noqa: E501 # pragma: no cover
 
@@ -701,6 +734,7 @@ def wait_in_other_thread():
         elif filetype == 'png':
             # NOTE: get_png already check if _upload_png_callback is not None
             get_png(viewer.figure)
+
         elif filetype == 'svg':
             if viewer.figure._upload_svg_callback is not None:
                 raise ValueError("previous svg export is still in progress. Wait to complete before making another call to save_figure") # noqa

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "specutils>=2.0.0",
     "specreduce>=1.6.0",
     "photutils>=2.2",
+    "pyavm>=0.9",
     "glue-astronomy>=0.12.0",
     "asteval>=0.9.23",
     "idna",