1#!/usr/local/bin/python3.8
2
3############################################################################
4#
5# MODULE:	i.colors.enhance (former i.landsat.rgb)
6#
7# AUTHOR(S):	Markus Neteler, original author
8#		Hamish Bowman, scripting enhancements
9#               Converted to Python by Glynn Clements
10#
11# PURPOSE:      create pretty RGBs: the trick is to remove outliers
12#               using percentiles (area under the histogram curve)
13#
14# COPYRIGHT:	(C) 2006, 2008, 2012-2014 by the GRASS Development Team
15#
16#		This program is free software under the GNU General Public
17#		License (>=v2). Read the file COPYING that comes with GRASS
18#		for details.
19#
20# TODO: implement better brightness control
21#############################################################################
22
23#%module
24#% description: Performs auto-balancing of colors for RGB images.
25#% keyword: imagery
26#% keyword: RGB
27#% keyword: satellite
28#% keyword: colors
29#%end
30#%option G_OPT_R_INPUT
31#% key: red
32#% description: Name of red channel
33#%end
34#%option G_OPT_R_INPUT
35#% key: green
36#% description: Name of green channel
37#%end
38#%option G_OPT_R_INPUT
39#% key: blue
40#% description: Name of blue channel
41#%end
42#%option
43#% key: strength
44#% type: double
45#% description: Cropping intensity (upper brightness level)
46#% options: 0-100
47#% answer : 98
48#% required: no
49#%end
50#%flag
51#% key: f
52#% description: Extend colors to full range of data on each channel
53#% guisection: Colors
54#%end
55#%flag
56#% key: p
57#% description: Preserve relative colors, adjust brightness only
58#% guisection: Colors
59#%end
60#%flag
61#% key: r
62#% description: Reset to standard color range
63#% guisection: Colors
64#%end
65#%flag
66#% key: s
67#% description: Process bands serially (default: run in parallel)
68#%end
69
70import sys
71
72import grass.script as gscript
73
74try:
75    # new for python 2.6, in 2.5 it may be easy_install'd.
76    import multiprocessing as mp
77    do_mp = True
78except:
79    do_mp = False
80
81
82def get_percentile(map, percentiles):
83    # todo: generalize for any list length
84    val1 = percentiles[0]
85    val2 = percentiles[1]
86    values = '%s,%s' % (val1, val2)
87
88    s = gscript.read_command('r.quantile', input=map,
89                             percentiles=values, quiet=True)
90
91    val_str1 = s.splitlines()[0].split(':')[2]
92    val_str2 = s.splitlines()[1].split(':')[2]
93    return (float(val_str1), float(val_str2))
94
95# wrapper to handle multiprocesses communications back to the parent
96
97
98def get_percentile_mp(map, percentiles, conn):
99    # Process() doesn't like storing connection parts in
100    #  separate dictionaries, only wants to pass through tuples,
101    #  so instead of just sending the sending the pipe we have to
102    #  send both parts then keep the one we want.  ??
103    output_pipe, input_pipe = conn
104    input_pipe.close()
105    result = get_percentile(map, percentiles)
106    gscript.debug('child (%s) (%.1f, %.1f)' % (map, result[0], result[1]))
107    output_pipe.send(result)
108    output_pipe.close()
109
110
111def set_colors(map, v0, v1):
112    rules = ''.join(["0% black\n", "%f black\n" % v0,
113                     "%f white\n" % v1, "100% white\n"])
114    gscript.write_command('r.colors', map=map, rules='-', stdin=rules,
115                          quiet=True)
116
117
118def main():
119    red = options['red']
120    green = options['green']
121    blue = options['blue']
122    brightness = options['strength']
123    full = flags['f']
124    preserve = flags['p']
125    reset = flags['r']
126
127    global do_mp
128
129    if flags['s']:
130        do_mp = False
131
132    check = True
133    for m in [red, green, blue]:
134        ex = gscript.find_file(m)
135        if ex['name'] == '':
136            check = False
137            gscript.warning("Raster map <{}> not found ".format(m))
138    if not check:
139        gscript.fatal("At least one of the input raster map was not found")
140    # 90 or 98? MAX value controls brightness
141    # think of percent (0-100), must be positive or 0
142    # must be more than "2" ?
143
144    if full:
145        for i in [red, green, blue]:
146            gscript.run_command('r.colors', map=i, color='grey', quiet=True)
147        sys.exit(0)
148
149    if reset:
150        for i in [red, green, blue]:
151            gscript.run_command('r.colors', map=i, color='grey255', quiet=True)
152        sys.exit(0)
153
154    if not preserve:
155        if do_mp:
156            gscript.message(_("Processing..."))
157            # set up jobs and launch them
158            proc = {}
159            conn = {}
160            for i in [red, green, blue]:
161                conn[i] = mp.Pipe()
162                proc[i] = mp.Process(target=get_percentile_mp,
163                                     args=(i, ['2', brightness],
164                                           conn[i],))
165                proc[i].start()
166            gscript.percent(1, 2, 1)
167
168            # collect results and wait for jobs to finish
169            for i in [red, green, blue]:
170                output_pipe, input_pipe = conn[i]
171                (v0, v1) = input_pipe.recv()
172                gscript.debug('parent (%s) (%.1f, %.1f)' % (i, v0, v1))
173                input_pipe.close()
174                proc[i].join()
175                set_colors(i, v0, v1)
176            gscript.percent(1, 1, 1)
177        else:
178            for i in [red, green, blue]:
179                gscript.message(_("Processing..."))
180                (v0, v1) = get_percentile(i, ['2', brightness])
181                gscript.debug("<%s>:  min=%f   max=%f" % (i, v0, v1))
182                set_colors(i, v0, v1)
183
184    else:
185        all_max = 0
186        all_min = 999999
187
188        if do_mp:
189            gscript.message(_("Processing..."))
190            # set up jobs and launch jobs
191            proc = {}
192            conn = {}
193            for i in [red, green, blue]:
194                conn[i] = mp.Pipe()
195                proc[i] = mp.Process(target=get_percentile_mp,
196                                     args=(i, ['2', brightness],
197                                           conn[i],))
198                proc[i].start()
199            gscript.percent(1, 2, 1)
200
201            # collect results and wait for jobs to finish
202            for i in [red, green, blue]:
203                output_pipe, input_pipe = conn[i]
204                (v0, v1) = input_pipe.recv()
205                gscript.debug('parent (%s) (%.1f, %.1f)' % (i, v0, v1))
206                input_pipe.close()
207                proc[i].join()
208                all_min = min(all_min, v0)
209                all_max = max(all_max, v1)
210            gscript.percent(1, 1, 1)
211        else:
212            for i in [red, green, blue]:
213                gscript.message(_("Processing..."))
214                (v0, v1) = get_percentile(i, ['2', brightness])
215                gscript.debug("<%s>:  min=%f   max=%f" % (i, v0, v1))
216                all_min = min(all_min, v0)
217                all_max = max(all_max, v1)
218
219        gscript.debug("all_min=%f   all_max=%f" % (all_min, all_max))
220        for i in [red, green, blue]:
221            set_colors(i, all_min, all_max)
222
223    # write cmd history:
224    mapset = gscript.gisenv()['MAPSET']
225    for i in [red, green, blue]:
226        if gscript.find_file(i)['mapset'] == mapset:
227            gscript.raster_history(i)
228
229if __name__ == "__main__":
230    options, flags = gscript.parser()
231    main()
232