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 import os.path
33 import sys
34
35 from osgeo import gdal
36
37
39 print('Usage: rgb2pct.py [-n colors | -pct palette_file] [-of format] source_file dest_file')
40 sys.exit(1)
41
42
46
47
48 -def GetExtension(filename):
49 ext = os.path.splitext(filename)[1]
50 if ext.startswith('.'):
51 ext = ext[1:]
52 return ext
53
54
77
78
80 drv_list = GetOutputDriversFor(filename)
81 ext = GetExtension(filename)
82 if not drv_list:
83 if not ext:
84 return 'GTiff'
85 else:
86 raise Exception("Cannot guess driver for %s" % filename)
87 elif len(drv_list) > 1:
88 print("Several drivers matching %s extension. Using %s" % (ext if ext else '', drv_list[0]))
89 return drv_list[0]
90
91
93 color_count = 256
94 frmt = None
95 src_filename = None
96 dst_filename = None
97 pct_filename = None
98
99 gdal.AllRegister()
100 argv = gdal.GeneralCmdLineProcessor(argv)
101 if argv is None:
102 sys.exit(0)
103
104
105 i = 1
106 while i < len(argv):
107 arg = argv[i]
108
109 if arg == '-of' or arg == '-f':
110 i = i + 1
111 frmt = argv[i]
112
113 elif arg == '-n':
114 i = i + 1
115 color_count = int(argv[i])
116
117 elif arg == '-pct':
118 i = i + 1
119 pct_filename = argv[i]
120
121 elif src_filename is None:
122 src_filename = argv[i]
123
124 elif dst_filename is None:
125 dst_filename = argv[i]
126
127 else:
128 Usage()
129
130 i = i + 1
131
132 if dst_filename is None:
133 Usage()
134
135
136
137 src_ds = gdal.Open(src_filename)
138 if src_ds is None:
139 print('Unable to open %s' % src_filename)
140 sys.exit(1)
141
142 if src_ds.RasterCount < 3:
143 print('%s has %d band(s), need 3 for inputs red, green and blue.'
144 % (src_filename, src_ds.RasterCount))
145 sys.exit(1)
146
147
148
149 if frmt is None:
150 frmt = GetOutputDriverFor(dst_filename)
151
152 dst_driver = gdal.GetDriverByName(frmt)
153 if dst_driver is None:
154 print('"%s" driver not registered.' % frmt)
155 sys.exit(1)
156
157
158
159 ct = gdal.ColorTable()
160 if pct_filename is None:
161 err = gdal.ComputeMedianCutPCT(src_ds.GetRasterBand(1),
162 src_ds.GetRasterBand(2),
163 src_ds.GetRasterBand(3),
164 color_count, ct,
165 callback=gdal.TermProgress_nocb)
166 else:
167 pct_ds = gdal.Open(pct_filename)
168 ct = pct_ds.GetRasterBand(1).GetRasterColorTable().Clone()
169
170
171
172
173 if format == 'GTiff':
174 tif_filename = dst_filename
175 else:
176 import tempfile
177 tif_filedesc, tif_filename = tempfile.mkstemp(suffix='.tif')
178
179 gtiff_driver = gdal.GetDriverByName('GTiff')
180
181 tif_ds = gtiff_driver.Create(tif_filename,
182 src_ds.RasterXSize, src_ds.RasterYSize, 1)
183
184 tif_ds.GetRasterBand(1).SetRasterColorTable(ct)
185
186
187
188
189 tif_ds.SetProjection(src_ds.GetProjection())
190 tif_ds.SetGeoTransform(src_ds.GetGeoTransform())
191 if src_ds.GetGCPCount() > 0:
192 tif_ds.SetGCPs(src_ds.GetGCPs(), src_ds.GetGCPProjection())
193
194
195
196
197 err = gdal.DitherRGB2PCT(src_ds.GetRasterBand(1),
198 src_ds.GetRasterBand(2),
199 src_ds.GetRasterBand(3),
200 tif_ds.GetRasterBand(1),
201 ct,
202 callback=gdal.TermProgress_nocb)
203
204 tif_ds = None
205
206 if tif_filename != dst_filename:
207 tif_ds = gdal.Open(tif_filename)
208 dst_driver.CreateCopy(dst_filename, tif_ds)
209 tif_ds = None
210
211 os.close(tif_filedesc)
212 gtiff_driver.Delete(tif_filename)
213
214 return err
215
216
217 if __name__ == '__main__':
218 sys.exit(main(sys.argv))
219