diff --git a/cutout/service/cutout_engine/astrocut_engine.py b/cutout/service/cutout_engine/astrocut_engine.py index ff185ae..f321e83 100644 --- a/cutout/service/cutout_engine/astrocut_engine.py +++ b/cutout/service/cutout_engine/astrocut_engine.py @@ -1,15 +1,93 @@ from __future__ import annotations -import shutil from pathlib import Path from typing import Any -from uuid import uuid4 +import numpy as np from astrocut import fits_cut +from astropy import units as u +from astropy.coordinates import SkyCoord +from astropy.io import fits +from astropy.wcs import WCS +from reproject import reproject_interp +from cutout import __version__ from cutout.service.stencils import Stencil from .base import CutoutEngine +from .color_composer import COLOR_PARAMS, _arcsinh_stretch, compose_rgb + + +def _mosaic_hdus( + data_hdus: list, + *, + center: SkyCoord, + cutout_size, + input_files: list[str], + ref_header: fits.Header, +) -> fits.PrimaryHDU: + """Reproject data HDUs onto a common grid. Returns a single PrimaryHDU.""" + + print(f"[astrocut] _mosaic_hdus: combining {len(data_hdus)} tiles") + + ref_wcs = WCS(ref_header) + pixel_scale = abs(ref_wcs.proj_plane_pixel_scales()[0].to_value(u.deg)) + + if hasattr(cutout_size, "unit"): + size_x = size_y = cutout_size.to(u.deg).value + else: + size_x = cutout_size[0].to(u.deg).value + size_y = cutout_size[1].to(u.deg).value + + nx = max(int(size_x / pixel_scale), 1) + ny = max(int(size_y / pixel_scale), 1) + + out_wcs = WCS(naxis=2) + out_wcs.wcs.ctype = [ref_wcs.wcs.ctype[0], ref_wcs.wcs.ctype[1]] + out_wcs.wcs.crval = [center.ra.deg, center.dec.deg] + out_wcs.wcs.crpix = [nx / 2.0, ny / 2.0] + out_wcs.wcs.cd = [[-pixel_scale, 0.0], [0.0, pixel_scale]] + out_wcs.wcs.radesys = "ICRS" + out_wcs.wcs.equinox = 2000.0 + + arrays = [] + for _idx, hdu in data_hdus: + arr, _ = reproject_interp(hdu, out_wcs, shape_out=(ny, nx), order="bilinear") + arrays.append(arr.astype(np.float32)) + hdu.data = None # free original cutout, no longer needed after reprojection + + stack = np.stack(arrays) + result = np.nanmean(stack, axis=0).astype(np.float32) + result[np.all(np.isnan(stack), axis=0)] = np.nan + + out_header = out_wcs.to_header() + out_header["HISTORY"] = f"Mosaic assembled from {len(data_hdus)} tiles using reproject_interp + nanmean" + out_header["NINPUTS"] = (len(data_hdus), "Number of input tiles combined") + out_header["METHOD"] = ("reproject_interp + nanmean", "Mosaicking method") + out_header["IMGTYPE"] = ("mosaic", "Image type") + for i, fpath in enumerate(input_files, 1): + out_header[f"INFILE{i:02d}"] = (str(Path(fpath).name), f"Input tile {i}") + out_header["ORIGIN"] = "data.linea.org.br" + out_header["SOFTNAME"] = "LIneA Cutout Service" + out_header["SOFTVERS"] = __version__ + + for kw in ("BUNIT", "MAGZERO", "FILTER", "BAND", "RADESYS", "EQUINOX", "TELESCOP", "INSTRUME"): + if kw in ref_header: + out_header[kw] = ref_header[kw] + + return fits.PrimaryHDU(data=result, header=out_header) + + +def _extract_data_hdus(hdul: fits.HDUList) -> list[tuple[int, fits.HDU]]: + """Return list of (index, hdu) for extensions that carry data.""" + return [(i, h) for i, h in enumerate(hdul) if getattr(h, "data", None) is not None] + + +def _add_provenance(header: fits.Header) -> None: + header["ORIGIN"] = "data.linea.org.br" + header["SOFTNAME"] = "LIneA Cutout Service" + header["SOFTVERS"] = __version__ + header["HISTORY"] = "Cutout produced by LIneA Cutout Service" class AstrocutEngine(CutoutEngine): @@ -26,7 +104,6 @@ def run_cutout( rgb_bands: str | None = None, persist: bool = False, ) -> Path: - # Accept png output (mono and rgb) in addition to fits. if output_format not in ("fits", "png"): raise ValueError("Astrocut engine currently supports only fits and png output") @@ -35,44 +112,52 @@ def run_cutout( output_path = Path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) - temp_tag = uuid4().hex[:8] stencil_obj = Stencil.from_dict(stencil) coordinate = stencil_obj.get_center() cutout_size = stencil_obj.get_cutout_size() stencil_type = stencil.get("type", "circle") - # Debug info: log input files and parameters - print(f"[astrocut] run_cutout: source_id={source_id} band={band} output_format={output_format} color={color} rgb_bands={rgb_bands} persist={persist}") + print( + f"[astrocut] run_cutout: source_id={source_id} band={band} output_format={output_format} color={color} rgb_bands={rgb_bands} persist={persist}" + ) print(f"[astrocut] run_cutout: stencil type={stencil_type} coordinate={coordinate} cutout_size={cutout_size}") print(f"[astrocut] run_cutout: input_files={input_files}") - # For FITS, reuse fits_cut + # --- FITS --- if output_format == "fits": - result = fits_cut( + results = fits_cut( input_files=input_files, coordinates=coordinate, cutout_size=cutout_size, single_outfile=True, - cutout_prefix=output_path.stem, - output_dir=output_path.parent, + memory_only=True, ) + hdul = results[0] + data_hdus = _extract_data_hdus(hdul) - print(f"[astrocut] fits_cut produced {result}") - - result_path = Path(result) - if result_path != output_path: - shutil.move(str(result_path), str(output_path)) + if len(data_hdus) > 1: + ref_header = data_hdus[0][1].header + primary = _mosaic_hdus( + data_hdus, + center=coordinate, + cutout_size=cutout_size, + input_files=list(input_files), + ref_header=ref_header, + ) + else: + hdu = data_hdus[0][1] + primary = fits.PrimaryHDU(data=hdu.data, header=hdu.header) + _add_provenance(primary.header) + primary.writeto(output_path, overwrite=True) + print(f"[astrocut] wrote FITS to {output_path}") return output_path - # For PNG, support mono and RGB composition - from astropy.io import fits - import numpy as np - from PIL import Image + # --- Color PNG --- + from PIL import Image, PngImagePlugin if output_format == "png" and color: - # Expect input_files as a dict: band -> list[str] if not isinstance(input_files, dict): raise ValueError("Color PNG requires input_files as a mapping band->files") @@ -84,106 +169,161 @@ def run_cutout( else: bands = list(raw) - temp_paths = [] arrays = [] + wcs_header = None for b in bands: files_b = input_files.get(b) if not files_b: raise ValueError(f"No input files provided for band {b}") - temp_fits = output_path.with_name(f"{output_path.stem}_{temp_tag}_{b}.fits") - print(f"[astrocut] creating temp fits for band {b} at {temp_fits} using files {files_b}") - res = fits_cut( + results = fits_cut( input_files=files_b, coordinates=coordinate, cutout_size=cutout_size, single_outfile=True, - cutout_prefix=temp_fits.stem, - output_dir=temp_fits.parent, + memory_only=True, ) - print(f"[astrocut] fits_cut for band {b} produced {res}") - temp_paths.append(Path(res)) - - # Read arrays: find first HDU with data in each FITS (astrocut writes CUTOUT extension) - for p in temp_paths: - with fits.open(p) as hdul: - data_hdu = None - for h in hdul: - if getattr(h, "data", None) is not None: - data_hdu = h.data - break - if data_hdu is None: - raise ValueError(f"No data HDU found in {p}") - arr = np.nan_to_num(data_hdu).astype(float) - print(f"[astrocut] read array from {p}: dtype={arr.dtype} shape={arr.shape} min={arr.min()} max={arr.max()}") - arrays.append(arr) - - # Ensure same shape by cropping to minimal shape + hdul = results[0] + data_hdus = _extract_data_hdus(hdul) + + if len(data_hdus) > 1: + ref_header = data_hdus[0][1].header + primary = _mosaic_hdus( + data_hdus, + center=coordinate, + cutout_size=cutout_size, + input_files=list(files_b), + ref_header=ref_header, + ) + arr = primary.data + if wcs_header is None: + wcs_header = primary.header + else: + hdu = data_hdus[0][1] + arr = hdu.data + if wcs_header is None: + wcs_header = hdu.header + + arr = np.nan_to_num(arr).astype(np.float32) + print(f"[astrocut] band {b}: dtype={arr.dtype} shape={arr.shape} min={arr.min()} max={arr.max()}") + arrays.append(arr) + min_rows = min(a.shape[0] for a in arrays) min_cols = min(a.shape[1] for a in arrays) - chans = [] - for a in arrays: - a_crop = a[:min_rows, :min_cols] - a_crop -= a_crop.min() - if a_crop.max() > 0: - a_crop = (a_crop / a_crop.max() * 255.0).astype('uint8') - else: - a_crop = a_crop.astype('uint8') - chans.append(a_crop) - - # Stack channels into RGB (order: channels[0]->R, [1]->G, [2]->B) - rgb = np.dstack(chans[:3]) if len(chans) >= 3 else np.dstack([chans[0]] * 3) - print(f"[astrocut] composing RGB image shape={rgb.shape} dtype={rgb.dtype}") - img = Image.fromarray(rgb, mode="RGB") - img.save(output_path) - print(f"[astrocut] saved PNG at {output_path} size={output_path.stat().st_size}") + arrays = [a[:min_rows, :min_cols] for a in arrays] + + rgb = compose_rgb(arrays, bands, source_id) - # Cleanup temps - for p in temp_paths: - try: - p.unlink() - except Exception: - pass + # --- Embed provenance + WCS into PNG --- + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("ORIGIN", "data.linea.org.br") + pnginfo.add_text("SOFTNAME", "LIneA Cutout Service") + pnginfo.add_text("SOFTVERS", __version__) + pnginfo.add_text("HISTORY", "RGB PNG composed from FITS cutouts using arcsinh stretch") + if wcs_header is not None: + for kw in ( + "CTYPE1", + "CTYPE2", + "CRPIX1", + "CRPIX2", + "CRVAL1", + "CRVAL2", + "CD1_1", + "CD1_2", + "CD2_1", + "CD2_2", + "CDELT1", + "CDELT2", + "PC1_1", + "PC1_2", + "PC2_1", + "PC2_2", + "NAXIS1", + "NAXIS2", + "RADESYS", + "EQUINOX", + ): + if kw in wcs_header: + pnginfo.add_text(kw, str(wcs_header[kw])) + + img = Image.fromarray(rgb) + img.save(output_path, pnginfo=pnginfo, compress_level=9, optimize=True) + print(f"[astrocut] saved PNG at {output_path} size={output_path.stat().st_size}") return output_path - # Fallback: single-band mono PNG conversion - temp_fits = output_path.with_name(f"{output_path.stem}_{temp_tag}_mono.fits") - result = fits_cut( - input_files=input_files if isinstance(input_files, list) else ([] if input_files is None else []), + # --- Mono PNG --- + from PIL import Image, PngImagePlugin + + results = fits_cut( + input_files=input_files if isinstance(input_files, list) else [], coordinates=coordinate, cutout_size=cutout_size, single_outfile=True, - cutout_prefix=temp_fits.stem, - output_dir=temp_fits.parent, + memory_only=True, ) + hdul = results[0] + data_hdus = _extract_data_hdus(hdul) - result_path = Path(result) - - print(f"[astrocut] mono fits_cut produced {result_path}") - with fits.open(result_path) as hdul: - data_hdu = None - for h in hdul: - if getattr(h, "data", None) is not None: - data_hdu = h.data - break - if data_hdu is None: - raise ValueError(f"No data HDU found in {result_path}") - data = data_hdu - - arr = np.nan_to_num(data).astype(float) - arr -= arr.min() - if arr.max() > 0: - arr = (arr / arr.max() * 255.0).astype('uint8') + if len(data_hdus) > 1: + ref_header = data_hdus[0][1].header + primary = _mosaic_hdus( + data_hdus, + center=coordinate, + cutout_size=cutout_size, + input_files=list(input_files) if isinstance(input_files, list) else [], + ref_header=ref_header, + ) + data = primary.data + wcs_header = primary.header else: - arr = arr.astype('uint8') + hdu = data_hdus[0][1] + data = hdu.data + wcs_header = hdu.header - img = Image.fromarray(arr) - img.save(output_path) + arr = np.nan_to_num(data).astype(np.float32) + cfg = COLOR_PARAMS.get(source_id, {}).get("arcsinh_clip", {}) + if band in cfg: + arr = _arcsinh_stretch(arr, *cfg[band]) + else: + arr -= arr.min() + if arr.max() > 0: + arr = (arr / arr.max() * 255.0).astype("uint8") + else: + arr = arr.astype("uint8") - try: - result_path.unlink() - except Exception: - pass + # --- Embed provenance + WCS into PNG --- + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("ORIGIN", "data.linea.org.br") + pnginfo.add_text("SOFTNAME", "LIneA Cutout Service") + pnginfo.add_text("SOFTVERS", __version__) + pnginfo.add_text("HISTORY", "Cutout produced by LIneA Cutout Service") + for kw in ( + "CTYPE1", + "CTYPE2", + "CRPIX1", + "CRPIX2", + "CRVAL1", + "CRVAL2", + "CD1_1", + "CD1_2", + "CD2_1", + "CD2_2", + "CDELT1", + "CDELT2", + "PC1_1", + "PC1_2", + "PC2_1", + "PC2_2", + "NAXIS1", + "NAXIS2", + "RADESYS", + "EQUINOX", + ): + if kw in wcs_header: + pnginfo.add_text(kw, str(wcs_header[kw])) + img = Image.fromarray(arr) + img.save(output_path, pnginfo=pnginfo, compress_level=9, optimize=True) + print(f"[astrocut] saved PNG at {output_path} size={output_path.stat().st_size}") return output_path diff --git a/cutout/service/cutout_engine/color_composer.py b/cutout/service/cutout_engine/color_composer.py new file mode 100644 index 0000000..c95f72c --- /dev/null +++ b/cutout/service/cutout_engine/color_composer.py @@ -0,0 +1,127 @@ +"""RGB color-composition for astronomical cutouts. + +Two methods, one fundamental difference: + +arcsinh_clip — per-band independent compression + Each band is normalized and arcsinh-stretched separately. + A pixel whose g-band value is faint will map to a dark blue, + regardless of what r and i are doing. + Result: high color contrast, saturated stars, dark background. + Use when: you want each filter to determine its own fate. + +lupton — shared-intensity compression (Lupton et al. 2004) + All three bands are normalized, then the average intensity + I = (R+G+B)/3 is computed. A single compression ratio f(I)/I + is applied to all three bands equally. + Result: preserves physical flux ratios between bands. + Use when: you want color fidelity over contrast. + +Why we do NOT use astropy's ``make_lupton_rgb``: + It places ``stretch`` in the *denominator* of the arcsinh + (``soften = Q / stretch``), which is the opposite of the + Lupton et al. 2004 paper (``asinh(alpha * beta * I) / alpha``). + After extensive testing, it cannot reproduce the behavior + of either arcsinh_clip or the paper's own algorithm. +""" + +from __future__ import annotations + +import numpy as np + +COLOR_PARAMS: dict[str, dict] = { + "des_dr2": { + "method": "lupton", + "arcsinh_clip": { + "g": (-1.2, 300.0, 80.0, 1.0), # (black, white, contrast, gain) + "r": (-1.2, 400.0, 80.0, 1.0), + "i": (-1.2, 400.0, 80.0, 1.0), + "z": (-1.2, 400.0, 80.0, 1.0), + "Y": (-1.0, 350.0, 80.0, 1.0), + }, + "lupton": { + "Q": 10.0, + "stretch": 8.0, + "gain": 1.9, + "g": (-1.2, 300), # (sky, clip_range = white - sky) + "r": (-1.2, 400), + "i": (-1.2, 400), + "z": (-1.2, 400), + "Y": (-1.0, 350), + }, + }, +} + + +def compose_rgb( + arrays: list[np.ndarray], + bands: list[str], + source_id: str, +) -> np.ndarray: + """Dispatch to the color method configured for *source_id*.""" + params = COLOR_PARAMS.get(source_id) + if params is None: + raise ValueError(f"No color params for catalog {source_id}") + + method = params.get("method", "arcsinh_clip") + if method == "arcsinh_clip": + return _arcsinh_clip(arrays, bands, params["arcsinh_clip"]) + elif method == "lupton": + return _lupton(arrays, bands, params["lupton"]) + raise ValueError(f"Unknown color method: {method}") + + +def _arcsinh_stretch( + array: np.ndarray, + black: float, + white: float, + contrast: float = 80.0, + gain: float = 1.0, +) -> np.ndarray: + """(pixel - black) / (white - black) → arcsinh(×contrast) / arcsinh(contrast) × gain.""" + scaled = np.clip((array - black) / (white - black), 0, 1) + stretched = np.arcsinh(scaled * contrast) / np.arcsinh(contrast) + return (stretched * gain * 255).clip(0, 255).astype(np.uint8) + + +def _arcsinh_clip( + arrays: list[np.ndarray], + bands: list[str], + cfg: dict, +) -> np.ndarray: + """Per-band arcsinh stretch, stacked as RGB.""" + chans = [] + for idx in [2, 1, 0]: + black, white, contrast, gain = cfg[bands[idx]] + chans.append(_arcsinh_stretch(arrays[idx], black, white, contrast, gain)) + return np.dstack(chans) + + +def _lupton( + arrays: list[np.ndarray], + bands: list[str], + cfg: dict, +) -> np.ndarray: + """Lupton et al. 2004: shared-intensity arcsinh with color-preserving ratio f(I)/I.""" + Q = cfg["Q"] + stretch = cfg["stretch"] + gain = cfg.get("gain", 1.0) + + # Per-band normalization to [0, 1] + norm = [] + for i, b in enumerate(bands): + sky, clip_range = cfg[b] + s = np.clip(arrays[i], sky, sky + clip_range) + s = (s - sky) / clip_range + norm.append(s) + + # Lupton 2004: I = (R+G+B)/3, ratio = f(I)/I + I = (norm[0] + norm[1] + norm[2]) / 3.0 + I = np.maximum(I, 1e-30) + fI = np.arcsinh(I * stretch * Q) / Q + ratio = np.clip(fI / I, 0, None) + + # Apply ratio + gain → uint8 + chans = [] + for idx in [2, 1, 0]: + chans.append((norm[idx] * ratio * gain * 255).clip(0, 255).astype(np.uint8)) + return np.dstack(chans) diff --git a/cutout/service/cutout_engine/tests/test_astrocut_engine.py b/cutout/service/cutout_engine/tests/test_astrocut_engine.py index b3fbb50..6651a38 100644 --- a/cutout/service/cutout_engine/tests/test_astrocut_engine.py +++ b/cutout/service/cutout_engine/tests/test_astrocut_engine.py @@ -6,22 +6,18 @@ def test_astrocut_engine_calls_fits_cut(monkeypatch): + import numpy as np + import cutout.service.cutout_engine.astrocut_engine as astro_module captured = {} def dummy_fits_cut(**kwargs): captured.update(kwargs) - return "/tmp/generated.fits" - - moved = {} - - def dummy_move(src, dst): - moved["src"] = src - moved["dst"] = dst + mock_hdu = astro_module.fits.PrimaryHDU(data=np.zeros((10, 10))) + return [astro_module.fits.HDUList([mock_hdu])] monkeypatch.setattr(astro_module, "fits_cut", dummy_fits_cut) - monkeypatch.setattr(astro_module.shutil, "move", dummy_move) engine = AstrocutEngine() result = engine.run_cutout( @@ -36,10 +32,10 @@ def dummy_move(src, dst): assert result == Path("/tmp/out.fits") assert captured["input_files"] == ["/data/tiles/a.fits.fz"] assert captured["single_outfile"] is True - assert moved == {"src": "/tmp/generated.fits", "dst": "/tmp/out.fits"} + assert captured["memory_only"] is True -def test_astrocut_engine_rejects_non_fits() -> None: +def test_astrocut_engine_rejects_unsupported_format() -> None: engine = AstrocutEngine() with pytest.raises(ValueError, match="supports only fits"): @@ -48,8 +44,8 @@ def test_astrocut_engine_rejects_non_fits() -> None: stencil={"type": "circle", "center": {"ra": 1.0, "dec": 2.0}, "radius": 0.1}, input_files=["/data/tiles/a.fits.fz"], band="g", - output_format="png", - output_path="/tmp/out.png", + output_format="jpg", + output_path="/tmp/out.jpg", ) diff --git a/cutout/service/tasks.py b/cutout/service/tasks.py index dec663b..14ec2ef 100644 --- a/cutout/service/tasks.py +++ b/cutout/service/tasks.py @@ -29,11 +29,17 @@ def _validate_input_files(files: list[str] | dict[str, list[str]] | None) -> Non raise FileNotFoundError(msg) -def _ensure_unpacked(files: list[str] | dict[str, list[str]] | None) -> list[str] | dict[str, list[str]] | None: - """If any input paths point to compressed `.fz` archives, unpack them to a tmp location and - return a structure of uncompressed paths suitable for engines like `astrocut`. +def _ensure_unpacked( + files: list[str] | dict[str, list[str]] | None, +) -> list[str] | dict[str, list[str]] | None: + """If any input paths point to compressed ``.fz`` archives, unpack them to a tmp location and + return a structure of uncompressed paths suitable for engines that require ``.fits`` files. - Returns the same shape as `files` but with `.fits` paths. + .. note:: + + ``fits_cut`` from astrocut handles ``.fz`` natively via ``.section``, so this helper + is currently **not used** by ``image_cutout``. It is kept as a legacy utility for + engines or ad-hoc scripts that need uncompressed files on disk. """ if not files: return files @@ -58,7 +64,6 @@ def _unpack_path(p: str) -> str: out: dict[str, list[str]] = {} for k, lst in files.items(): out[k] = [_unpack_path(p) for p in (lst or [])] - # log unpacked mapping print(f"[tasks] _ensure_unpacked: band={k} unpacked_paths={out[k]}") return out @@ -92,19 +97,15 @@ def image_cutout( f"color={color} rgb_bands={rgb_bands}" ) print(f"[tasks] image_cutout initial files={files}") - # If inputs reference compressed archives, unpack them for engines that require `.fits` files. cutout_engine = create_cutout_engine(engine) - unpacked = _ensure_unpacked(files) - print(f"[tasks] image_cutout unpacked files={unpacked}") - # validate the (possibly unpacked) input paths exist before running engine - _validate_input_files(unpacked) + _validate_input_files(files) try: print(f"[tasks] calling engine.run_cutout engine={engine} path={path}") result = cutout_engine.run_cutout( source_id=source_id, stencil=stencil, - input_files=unpacked, + input_files=files, band=band, output_format=format, output_path=path, diff --git a/requirements/base.txt b/requirements/base.txt index e0fae3b..1647f81 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -32,3 +32,4 @@ astropy==7.2.0 # https://www.astropy.org/ numpy==2.4.4 # https://numpy.org/ matplotlib==3.10.9 # https://matplotlib.org/ astrocut==1.2.0 # https://astrocut.readthedocs.io/ +reproject==0.14.1