1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
|
#!/usr/bin/env python3
#
# Copyright (c) Weitian LI <weitian@aaronly.me>
# MIT license
#
"""
Create FITS image cube from a series of image slices.
"""
import os
import sys
import argparse
from datetime import datetime
import numpy as np
from astropy.io import fits
from astropy.wcs import WCS
class FITSCube:
"""
FITS image cube.
"""
def __init__(self, infile=None):
if infile is not None:
self.load(infile)
def load(self, infile):
with fits.open(infile) as f:
self.data = f[0].data
self.header = f[0].header
print("Loaded FITS cube from file: %s" % infile)
print("Cube dimensions: %dx%dx%d" %
(self.width, self.height, self.nslice))
def add_slices(self, slices, zbegin=0.0, zstep=1.0):
"""
Create a FITS cube from input image slices.
"""
nslice = len(slices)
header, image = self.open_image(slices[0])
shape = (nslice, ) + image.shape
data = np.zeros(shape, dtype=image.dtype)
for i, s in enumerate(slices):
print("[%d/%d] Adding image slice: %s ..." % (i, nslice, s))
hdr, img = self.open_image(s)
data[i, :, :] = img
self.data = data
wcs = self.make_wcs(header, zbegin=zbegin, zstep=zstep)
self.header = header.copy(strip=True)
self.header.extend(wcs.to_header(), update=True)
print("Created FITS cube of dimensions: %dx%dx%d" %
(self.width, self.height, self.nslice))
@staticmethod
def open_image(infile):
"""
Open the slice image and return its header and 2D image data.
NOTE
----
The input slice image may have following dimensions:
* NAXIS=2: [Y, X]
* NAXIS=3: [FREQ=1, Y, X]
* NAXIS=4: [FREQ=1, STOKES=1, Y, X]
NOTE
----
Only open slice image that has only ONE frequency and ONE Stokes
parameter.
Returns
-------
header : `~astropy.io.fits.Header`
image : 2D `~numpy.ndarray`
The 2D [Y, X] image part of the slice image.
"""
with fits.open(infile) as f:
header = f[0].header
data = f[0].data
if data.ndim == 2:
# NAXIS=2: [Y, X]
image = data
elif data.ndim == 3 and data.shape[0] == 1:
# NAXIS=3: [FREQ=1, Y, X]
image = data[0, :, :]
elif data.ndim == 4 and data.shape[0] == 1 and data.shape[1] == 1:
# NAXIS=4: [FREQ=1, STOKES=1, Y, X]
image = data[0, 0, :, :]
else:
raise ValueError("Slice '{0}' has invalid dimensions: {1}".format(
infile, data.shape))
return (header, image)
def make_wcs(self, header, zbegin, zstep):
w = WCS(naxis=3)
w.wcs.ctype = ["pixel", "pixel", "pixel"]
w.wcs.crpix = np.array([header.get("CRPIX1", 1.0),
header.get("CRPIX2", 1.0),
1.0])
w.wcs.crval = np.array([header.get("CRVAL1", 0.0),
header.get("CRVAL2", 0.0),
zbegin])
w.wcs.cdelt = np.array([header.get("CDELT1", 1.0),
header.get("CDELT2", 1.0),
zstep])
return w
def write(self, outfile, clobber):
self.header.add_history(datetime.now().isoformat())
self.header.add_history(" ".join(sys.argv))
hdu = fits.PrimaryHDU(data=self.data, header=self.header)
try:
hdu.writeto(outfile, overwrite=clobber)
except TypeError:
hdu.writeto(outfile, clobber=clobber)
@property
def width(self):
__, __, w = self.data.shape
return w
@property
def height(self):
__, h, __ = self.data.shape
return h
@property
def nslice(self):
ns, __, __ = self.data.shape
return ns
@property
def zbegin(self):
"""
The Z-axis position of the first slice.
"""
return self.header["CRVAL3"]
@property
def zstep(self):
"""
The Z-axis step/spacing between slices.
"""
return self.header["CDELT3"]
@property
def zvalues(self):
"""
Calculate the Z-axis positions for all slices
"""
nslice = self.nslice
wcs = WCS(self.header)
pix = np.zeros(shape=(nslice, 3), dtype=np.int)
pix[:, 2] = np.arange(nslice)
world = wcs.wcs_pix2world(pix, 0)
return world[:, 2]
def cmd_info(args):
"""
Sub-command: "info", show FITS cube information
"""
cube = FITSCube(args.infile)
print("Image/slice size: %dx%d" % (cube.width, cube.height))
print("Number of slices: %d" % cube.nslice)
print("Slice step/spacing: %.3f" % cube.zstep)
print("Slice positions: {0}".format(cube.zvalues))
def cmd_create(args):
"""
Sub-command: "create", create a FITS cube
"""
if not args.clobber and os.path.exists(args.outfile):
raise FileExistsError("output file already exists: %s" % args.outfile)
cube = FITSCube()
cube.add_slices(args.infiles, zbegin=args.zbegin, zstep=args.zstep)
cube.write(args.outfile, clobber=args.clobber)
print("Created FITS cube: %s" % args.outfile)
def main():
parser = argparse.ArgumentParser(
description="Create FITS cube from a series of image slices.")
subparsers = parser.add_subparsers(dest="subparser_name",
title="sub-commands",
help="additional help")
# sub-command: "info"
parser_info = subparsers.add_parser("info", help="show FITS cube info")
parser_info.add_argument("infile", help="FITS cube filename")
parser_info.set_defaults(func=cmd_info)
# sub-command: "create"
parser_create = subparsers.add_parser("create", help="create a FITS cube")
parser_create.add_argument("-C", "--clobber", dest="clobber",
action="store_true",
help="overwrite existing output file")
parser_create.add_argument("-z", "--z-begin", dest="zbegin",
type=float, default=0.0,
help="Z-axis position of the first slice")
parser_create.add_argument("-s", "--z-step", dest="zstep",
type=float, default=1.0,
help="Z-axis step/spacing between slices")
parser_create.add_argument("-o", "--outfile", dest="outfile",
required=True,
help="output FITS cube filename")
parser_create.add_argument("-i", "--infiles", dest="infiles",
nargs="+", required=True,
help="input image slices (in order)")
parser_create.set_defaults(func=cmd_create)
#
args = parser.parse_args()
args.func(args)
if __name__ == "__main__":
main()
|
