-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate_rgb.py
98 lines (83 loc) · 3.46 KB
/
generate_rgb.py
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
import os
import sys
import numpy as np
import astropy.io.fits as pyfits
from compose_class import channel_RGB
# import aplpy
# def generate_rgb(src_path, dst_path, number=20000):
# start = 100000
# bands = ['Band1', 'Band2', 'Band3', 'Band4']
# band_letter = ['R', 'I', 'G', 'U']
# for i in range(number):
# image_src = []
# for j in range(3):
# image = f'{bands[j]}/imageSDSS_{band_letter[j]}-{start + i}.fits'
# image = os.path.join(src_path, image)
# image_src.append(image)
# temp_image = f'/tmp/temp_cube_{start + i}.fits'
# temp_2d_image = f'/tmp/temp_cube_{start + i}_2d.fits'
# dst_image = f'ground_based_{start + i}.png'
# dst_image = os.path.join(dst_path, dst_image)
# aplpy.make_rgb_cube(image_src, temp_image)
# aplpy.make_rgb_image(temp_image, dst_image)
# os.remove(temp_image)
# os.remove(temp_2d_image)
def img_preproc(ifile, rfile, gfile):
img_g = pyfits.getdata(gfile)
img_g[np.where(img_g == 100.0)] = 0.0
img_r = pyfits.getdata(rfile)
img_r[np.where(img_r == 100.0)] = 0.0
img_i = pyfits.getdata(ifile)
img_i[np.where(img_i == 100.0)] = 0.0
img_rscl = 1e-11
img_thds = 1e-11
img_g_rscl = img_g/img_rscl
img_r_rscl = img_r/img_rscl
img_i_rscl = img_i/img_rscl
img_g_rscl[np.where(img_g_rscl <= img_thds)] = img_thds
img_r_rscl[np.where(img_r_rscl <= img_thds)] = img_thds
img_i_rscl[np.where(img_i_rscl <= img_thds)] = img_thds
return img_g_rscl, img_r_rscl, img_i_rscl
def generate_rgb_single(img_g_rscl, img_r_rscl, img_i_rscl):
scales, offset, Q, alpha = (0.75, 1.05, 1.5), 0.01, 2.0, 0.7
masklevel = -1.0
saturation, itype = 'white', 'sum'
object_gri = channel_RGB(RED=img_i_rscl, GREEN=img_r_rscl, BLUE=img_g_rscl)
object_gri.apply_scale(scales=scales)
object_gri.lupton_stretch(Q=Q, alpha=alpha, itype=itype)
object_gri.pjm_mask(masklevel=masklevel)
object_gri.pjm_offset(offset=offset)
object_gri.lupton_saturate(saturation=saturation)
object_gri.pack_up()
object_gri.imgRGB.save(dst_image)
return
def generate_rgb(src_path, dst_path, number=20000):
# scales, offset, Q, alpha, masklevel, saturation,
# itype = (0.75,1.05,1.5),
# 0.01, 2.0,0.7, -1.0, 'color', 'rms' # gri
scales, offset, Q, alpha = (0.75, 1.05, 1.5), 0.01, 2.0, 0.7
masklevel = -1.0
saturation, itype = 'white', 'sum'
start = 100000
bands = ['Band2', 'Band1', 'Band3', 'Band4']
band_letter = ['I', 'R', 'G', 'U']
for i in range(number):
image_src = []
for j in range(3):
image = f'{bands[j]}/imageSDSS_{band_letter[j]}-{start + i}.fits'
image = os.path.join(src_path, image)
image_src.append(image)
img_g_rscl, img_r_rscl, img_i_rscl = img_preproc(*image_src)
dst_image = f'ground_based_{start + i}.png'
dst_image = os.path.join(dst_path, dst_image)
object_gri = channel_RGB(RED=img_i_rscl, GREEN=img_r_rscl, BLUE=img_g_rscl)
object_gri.apply_scale(scales=scales)
object_gri.lupton_stretch(Q=Q, alpha=alpha, itype=itype)
object_gri.pjm_mask(masklevel=masklevel)
object_gri.pjm_offset(offset=offset)
object_gri.lupton_saturate(saturation=saturation)
object_gri.pack_up()
object_gri.imgRGB.save(dst_image)
print(f"{dst_image} saved.")
if __name__ == "__main__":
generate_rgb(sys.argv[1], sys.argv[2])