代码改变世界

gdal带的影像合并Python代码

2007-05-10 13:37  flyingfish  阅读(2993)  评论(0编辑  收藏  举报
    ps:最近用vc作遥感影像的入库功能,借助GDAL来做,基本已经可以使用了。最后要写段代码对建立好的金字塔作合并输出,再次翻看了gdal_merge.py,其实这段代码很早就看过,尽管道理不复杂,但看到Frank把代码写的这么精致,不由肃然起敬,放到这里经常来看看勉励自己。
    写代码有好几年了,现在才深刻体会到软件设计模式、架构、文档、规范的重要性,反而感觉做工作时有些瞻前顾后,总想把代码的逻辑整的很清晰以提高可维护行和复用性,真是那个 阶段有哪个阶段的烦恼,呵呵!

  1
 #!/usr/bin/env python
  2 ###############################################################################
  3 # $Id: gdal_merge.py,v 1.25 2006/04/20 13:27:57 fwarmerdam Exp $
  4 #
  5 # Project:  InSAR Peppers
  6 # Purpose:  Module to extract data from many rasters into one output.
  7 # Author:   Frank Warmerdam, warmerdam@pobox.com
  8 #
  9 ###############################################################################
 10 # Copyright (c) 2000, Atlantis Scientific Inc. (www.atlsci.com)
 11 # 
 12 # This library is free software; you can redistribute it and/or
 13 # modify it under the terms of the GNU Library General Public
 14 # License as published by the Free Software Foundation; either
 15 # version 2 of the License, or (at your option) any later version.
 16 # 
 17 # This library is distributed in the hope that it will be useful,
 18 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 19 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 20 # Library General Public License for more details.
 21 # 
 22 # You should have received a copy of the GNU Library General Public
 23 # License along with this library; if not, write to the
 24 # Free Software Foundation, Inc., 59 Temple Place - Suite 330,
 25 # Boston, MA 02111-1307, USA.
 26 ###############################################################################
 27 # 
 28 #  $Log: gdal_merge.py,v $
 29 #  Revision 1.25  2006/04/20 13:27:57  fwarmerdam
 30 #  Added error checks on Driver's Create support and success of Create.
 31 #
 32 #  Revision 1.24  2006/01/26 23:40:39  fwarmerdam
 33 #  treat nodata as a float.
 34 #
 35 #  Revision 1.23  2005/11/25 02:13:36  fwarmerdam
 36 #  Added --help-general.
 37 #
 38 #  Revision 1.22  2005/11/16 18:30:49  fwarmerdam
 39 #  Fixed round off issue with output file size.
 40 #
 41 #  Revision 1.21  2005/08/18 15:45:15  fwarmerdam
 42 #  Added the -createonly switch.
 43 #
 44 #  Revision 1.20  2005/07/19 03:33:39  fwarmerdam
 45 #  removed left over global_list
 46 #
 47 #  Revision 1.19  2005/06/23 19:51:51  fwarmerdam
 48 #  Fixed support for non-square pixels c/o Matt Giger
 49 #  http://bugzilla.remotesensing.org/show_bug.cgi?id=874
 50 #
 51 #  Revision 1.18  2005/03/29 22:40:00  fwarmerdam
 52 #  Added -ot option.
 53 #
 54 #  Revision 1.17  2005/02/23 18:29:07  fwarmerdam
 55 #  Accept "either spelling" of separate.
 56 #
 57 #  Revision 1.16  2005/02/23 18:23:00  fwarmerdam
 58 #  Added -seperate to the usage message.
 59 #
 60 #  Revision 1.15  2004/09/02 22:06:24  warmerda
 61 #  Added a bit of commandline error reporting.
 62 #
 63 #  Revision 1.14  2004/08/23 15:05:27  warmerda
 64 #  Added projection setting for new files.
 65 #
 66 #  Revision 1.13  2004/04/02 22:31:26  warmerda
 67 #  Use -of for format.
 68 #
 69 #  Revision 1.12  2004/04/02 17:40:44  warmerda
 70 #  added GDALGeneralCmdLineProcessor() support
 71 #
 72 #  Revision 1.11  2004/03/26 17:11:42  warmerda
 73 #  added -init
 74 #
 75 #  Revision 1.10  2003/04/22 14:42:45  warmerda
 76 #  Don't import Numeric unless we need it.
 77 #
 78 #  Revision 1.9  2003/04/22 13:30:05  warmerda
 79 #  Added -co flag.
 80 #
 81 #  Revision 1.8  2003/03/07 16:26:39  warmerda
 82 #  fixed up for ungeoreferenced files, supress extra error
 83 #
 84 #  Revision 1.7  2003/01/28 15:00:13  warmerda
 85 #  applied patch for multi-band support from Ken Boss
 86 #
 87 #  Revision 1.6  2003/01/20 22:19:08  warmerda
 88 #  added nodata support
 89 #
 90 #  Revision 1.5  2002/12/12 14:54:42  warmerda
 91 #  added the -pct flag to copy over a pct
 92 #
 93 #  Revision 1.4  2002/12/12 14:48:12  warmerda
 94 #  removed broken options arg to gdal.Create()
 95 #
 96 #  Revision 1.3  2002/04/03 21:12:05  warmerda
 97 #  added -separate flag for Gerald Buckmaster
 98 #
 99 #  Revision 1.2  2000/11/29 20:36:18  warmerda
100 #  allow output file to be preexisting
101 #
102 #  Revision 1.1  2000/11/29 20:23:13  warmerda
103 #  New
104 #
105 #
106 
107 import gdal
108 import sys
109 
110 verbose = 0
111 
112 
113 # =============================================================================
114 def raster_copy( s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
115                  t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
116                  nodata=None ):
117 
118     if nodata is not None:
119         return raster_copy_with_nodata(
120             s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
121             t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
122             nodata )
123 
124     if verbose != 0:
125         print 'Copy %d,%d,%d,%d to %d,%d,%d,%d.' \
126               % (s_xoff, s_yoff, s_xsize, s_ysize,
127              t_xoff, t_yoff, t_xsize, t_ysize )
128 
129     s_band = s_fh.GetRasterBand( s_band_n )
130     t_band = t_fh.GetRasterBand( t_band_n )
131 
132     data = s_band.ReadRaster( s_xoff, s_yoff, s_xsize, s_ysize,
133                               t_xsize, t_ysize, t_band.DataType )
134     t_band.WriteRaster( t_xoff, t_yoff, t_xsize, t_ysize,
135                         data, t_xsize, t_ysize, t_band.DataType )
136         
137 
138     return 0
139     
140 # =============================================================================
141 def raster_copy_with_nodata( s_fh, s_xoff, s_yoff, s_xsize, s_ysize, s_band_n,
142                              t_fh, t_xoff, t_yoff, t_xsize, t_ysize, t_band_n,
143                              nodata ):
144 
145     import Numeric
146     
147     if verbose != 0:
148         print 'Copy %d,%d,%d,%d to %d,%d,%d,%d.' \
149               % (s_xoff, s_yoff, s_xsize, s_ysize,
150              t_xoff, t_yoff, t_xsize, t_ysize )
151 
152     s_band = s_fh.GetRasterBand( s_band_n )
153     t_band = t_fh.GetRasterBand( t_band_n )
154 
155     data_src = s_band.ReadAsArray( s_xoff, s_yoff, s_xsize, s_ysize,
156                                    t_xsize, t_ysize )
157     data_dst = t_band.ReadAsArray( t_xoff, t_yoff, t_xsize, t_ysize )
158 
159     nodata_test = Numeric.equal(data_src,nodata)
160     to_write = Numeric.choose( nodata_test, (data_src, data_dst) )
161                                
162     t_band.WriteArray( to_write, t_xoff, t_yoff )
163 
164     return 0
165     
166 # =============================================================================
167 def names_to_fileinfos( names ):
168     """
169     Translate a list of GDAL filenames, into file_info objects.
170 
171     names -- list of valid GDAL dataset names.
172 
173     Returns a list of file_info objects.  There may be less file_info objects
174     than names if some of the names could not be opened as GDAL files.
175     """
176     
177     file_infos = []
178     for name in names:
179         fi = file_info()
180         if fi.init_from_name( name ) == 1:
181             file_infos.append( fi )
182 
183     return file_infos
184 
185 # *****************************************************************************
186 class file_info:
187     """A class holding information about a GDAL file."""
188 
189     def init_from_name(self, filename):
190         """
191         Initialize file_info from filename
192 
193         filename -- Name of file to read.
194 
195         Returns 1 on success or 0 if the file can't be opened.
196         """
197         fh = gdal.Open( filename )
198         if fh is None:
199             return 0
200 
201         self.filename = filename
202         self.bands = fh.RasterCount
203         self.xsize = fh.RasterXSize
204         self.ysize = fh.RasterYSize
205         self.band_type = fh.GetRasterBand(1).DataType
206         self.projection = fh.GetProjection()
207         self.geotransform = fh.GetGeoTransform()
208         self.ulx = self.geotransform[0]
209         self.uly = self.geotransform[3]
210         self.lrx = self.ulx + self.geotransform[1* self.xsize
211         self.lry = self.uly + self.geotransform[5* self.ysize
212 
213         ct = fh.GetRasterBand(1).GetRasterColorTable()
214         if ct is not None:
215             self.ct = ct.Clone()
216         else:
217             self.ct = None
218 
219         return 1
220 
221     def report( self ):
222         print 'Filename: '+ self.filename
223         print 'File Size: %dx%dx%d' \
224               % (self.xsize, self.ysize, self.bands)
225         print 'Pixel Size: %f x %f' \
226               % (self.geotransform[1],self.geotransform[5])
227         print 'UL:(%f,%f)   LR:(%f,%f)' \
228               % (self.ulx,self.uly,self.lrx,self.lry)
229 
230     def copy_into( self, t_fh, s_band = 1, t_band = 1, nodata_arg=None ):
231         """
232         Copy this files image into target file.
233 
234         This method will compute the overlap area of the file_info objects
235         file, and the target gdal.Dataset object, and copy the image data
236         for the common window area.  It is assumed that the files are in
237         a compatible projection  no checking or warping is done.  However,
238         if the destination file is a different resolution, or different
239         image pixel type, the appropriate resampling and conversions will
240         be done (using normal GDAL promotion/demotion rules).
241 
242         t_fh -- gdal.Dataset object for the file into which some or all
243         of this file may be copied.
244 
245         Returns 1 on success (or if nothing needs to be copied), and zero one
246         failure.
247         """
248         t_geotransform = t_fh.GetGeoTransform()
249         t_ulx = t_geotransform[0]
250         t_uly = t_geotransform[3]
251         t_lrx = t_geotransform[0] + t_fh.RasterXSize * t_geotransform[1]
252         t_lry = t_geotransform[3+ t_fh.RasterYSize * t_geotransform[5]
253 
254         # figure out intersection region
255         tgw_ulx = max(t_ulx,self.ulx)
256         tgw_lrx = min(t_lrx,self.lrx)
257         if t_geotransform[5< 0:
258             tgw_uly = min(t_uly,self.uly)
259             tgw_lry = max(t_lry,self.lry)
260         else:
261             tgw_uly = max(t_uly,self.uly)
262             tgw_lry = min(t_lry,self.lry)
263         
264         # do they even intersect?
265         if tgw_ulx >= tgw_lrx:
266             return 1
267         if t_geotransform[5< 0 and tgw_uly <= tgw_lry:
268             return 1
269         if t_geotransform[5> 0 and tgw_uly >= tgw_lry:
270             return 1
271             
272         # compute target window in pixel coordinates.
273         tw_xoff = int((tgw_ulx - t_geotransform[0]) / t_geotransform[1+ 0.1)
274         tw_yoff = int((tgw_uly - t_geotransform[3]) / t_geotransform[5+ 0.1)
275         tw_xsize = int((tgw_lrx - t_geotransform[0])/t_geotransform[1+ 0.5) \
276                    - tw_xoff
277         tw_ysize = int((tgw_lry - t_geotransform[3])/t_geotransform[5+ 0.5) \
278                    - tw_yoff
279 
280         if tw_xsize < 1 or tw_ysize < 1:
281             return 1
282 
283         # Compute source window in pixel coordinates.
284         sw_xoff = int((tgw_ulx - self.geotransform[0]) / self.geotransform[1])
285         sw_yoff = int((tgw_uly - self.geotransform[3]) / self.geotransform[5])
286         sw_xsize = int((tgw_lrx - self.geotransform[0]) \
287                        / self.geotransform[1+ 0.5- sw_xoff
288         sw_ysize = int((tgw_lry - self.geotransform[3]) \
289                        / self.geotransform[5+ 0.5- sw_yoff
290 
291         if sw_xsize < 1 or sw_ysize < 1:
292             return 1
293 
294         # Open the source file, and copy the selected region.
295         s_fh = gdal.Open( self.filename )
296 
297         return \
298             raster_copy( s_fh, sw_xoff, sw_yoff, sw_xsize, sw_ysize, s_band,
299                          t_fh, tw_xoff, tw_yoff, tw_xsize, tw_ysize, t_band,
300                          nodata_arg )
301 
302 
303 # =============================================================================
304 def Usage():
305     print 'Usage: gdal_merge.py [-o out_filename] [-of out_format] [-co NAME=VALUE]*'
306     print '                     [-ps pixelsize_x pixelsize_y] [-separate] [-v] [-pct]'
307     print '                     [-ul_lr ulx uly lrx lry] [-n nodata_value] [-init value]'
308     print '                     [-ot datatype] [-createonly] input_files'
309     print '                     [--help-general]'
310     print
311 
312 # =============================================================================
313 #
314 # Program mainline.
315 #
316 
317 if __name__ == '__main__':
318 
319     names = []
320     format = 'GTiff'
321     out_file = 'out.tif'
322 
323     ulx = None
324     psize_x = None
325     separate = 0
326     copy_pct = 0
327     nodata = None
328     create_options = []
329     pre_init = None
330     band_type = None
331     createonly = 0
332 
333     gdal.AllRegister()
334     argv = gdal.GeneralCmdLineProcessor( sys.argv )
335     if argv is None:
336         sys.exit( 0 )
337 
338     # Parse command line arguments.
339     i = 1
340     while i < len(argv):
341         arg = argv[i]
342 
343         if arg == '-o':
344             i = i + 1
345             out_file = argv[i]
346 
347         elif arg == '-v':
348             verbose = 1
349 
350         elif arg == '-createonly':
351             createonly = 1
352 
353         elif arg == '-separate':
354             separate = 1
355 
356         elif arg == '-seperate':
357             separate = 1
358 
359         elif arg == '-pct':
360             copy_pct = 1
361 
362         elif arg == '-ot':
363             i = i + 1
364             band_type = gdal.GetDataTypeByName( argv[i] )
365             if band_type == gdal.GDT_Unknown:
366                 print 'Unknown GDAL data type: ', argv[i]
367                 sys.exit( 1 )
368 
369         elif arg == '-init':
370             i = i + 1
371             pre_init = float(argv[i])
372 
373         elif arg == '-n':
374             i = i + 1
375             nodata = float(argv[i])
376 
377         elif arg == '-f':
378             # for backward compatibility.
379             i = i + 1
380             format = argv[i]
381 
382         elif arg == '-of':
383             i = i + 1
384             format = argv[i]
385 
386         elif arg == '-co':
387             i = i + 1
388             create_options.append( argv[i] )
389 
390         elif arg == '-ps':
391             psize_x = float(argv[i+1])
392             psize_y = -1 * abs(float(argv[i+2]))
393             i = i + 2
394 
395         elif arg == '-ul_lr':
396             ulx = float(argv[i+1])
397             uly = float(argv[i+2])
398             lrx = float(argv[i+3])
399             lry = float(argv[i+4])
400             i = i + 4
401 
402         elif arg[:1== '-':
403             print 'Unrecognised command option: ', arg
404             Usage()
405             sys.exit( 1 )
406 
407         else:
408             names.append( arg )
409             
410         i = i + 1
411 
412     if len(names) == 0:
413         print 'No input files selected.'
414         Usage()
415         sys.exit( 1 )
416 
417     Driver = gdal.GetDriverByName(format)
418     if Driver is None:
419         print 'Format driver %s not found, pick a supported driver.' % format
420         sys.exit( 1 )
421 
422     DriverMD = Driver.GetMetadata()
423     if not DriverMD.has_key('DCAP_CREATE'):
424         print 'Format driver %s does not support creation and piecewise writing.\nPlease select a format that does, such as GTiff (the default) or HFA (Erdas Imagine).' % format
425         sys.exit( 1 )
426 
427     # Collect information on all the source files.
428     file_infos = names_to_fileinfos( names )
429 
430     if ulx is None:
431         ulx = file_infos[0].ulx
432         uly = file_infos[0].uly
433         lrx = file_infos[0].lrx
434         lry = file_infos[0].lry
435         
436         for fi in file_infos:
437             ulx = min(ulx, fi.ulx)
438             uly = max(uly, fi.uly)
439             lrx = max(lrx, fi.lrx)
440             lry = min(lry, fi.lry)
441 
442     if psize_x is None:
443         psize_x = file_infos[0].geotransform[1]
444         psize_y = file_infos[0].geotransform[5]
445 
446     if band_type is None:
447         band_type = file_infos[0].band_type
448 
449     # Try opening as an existing file.
450     gdal.PushErrorHandler( 'CPLQuietErrorHandler' )
451     t_fh = gdal.Open( out_file, gdal.GA_ReadOnly )
452     gdal.PopErrorHandler()
453     
454     # Create output file if it does not already exist.
455     if t_fh is None:
456         geotransform = [ulx, psize_x, 0, uly, 0, psize_y]
457 
458         xsize = int((lrx - ulx) / geotransform[1+ 0.5)
459         ysize = int((lry - uly) / geotransform[5+ 0.5)
460 
461         if separate != 0:
462             bands = len(file_infos)
463         else:
464             bands = file_infos[0].bands
465 
466         t_fh = Driver.Create( out_file, xsize, ysize, bands,
467                               band_type, create_options )
468         if t_fh is None:
469             print 'Creation failed, terminating gdal_merge.'
470             sys.exit( 1 )
471             
472         t_fh.SetGeoTransform( geotransform )
473         t_fh.SetProjection( file_infos[0].projection )
474 
475         if copy_pct:
476             t_fh.GetRasterBand(1).SetRasterColorTable(file_infos[0].ct)
477     else:
478         if separate != 0:
479             bands = len(file_infos)
480         else:
481             bands = min(file_infos[0].bands,t_fh.RasterCount)
482 
483     # Do we need to pre-initialize the whole mosaic file to some value?
484     if pre_init is not None:
485         for i in range(t_fh.RasterCount):
486             t_fh.GetRasterBand(i+1).Fill( pre_init )
487 
488     # Copy data from source files into output file.
489     t_band = 1
490     for fi in file_infos:
491         if createonly != 0:
492             continue
493         
494         if verbose != 0:
495             print
496             fi.report()
497 
498         if separate == 0 :
499             for band in range(1, bands+1):
500                 fi.copy_into( t_fh, band, band, nodata )
501         else:
502             fi.copy_into( t_fh, 1, t_band, nodata )
503             t_band = t_band+1
504             
505     # Force file to be closed.
506     t_fh = None
507