python类Open()的实例源码

def test_gdal_vs_numpy(self):
        for t in self.tifs:
            ds = gdal.Open(t)
            n_list = geoinfo.numpy_band_stats(ds, t, 1)
            g_list = geoinfo.band_stats(ds, t, 1)
            self.assertEqual(n_list[0], g_list[0])
            self.assertEqual(n_list[-3], g_list[-3])
            self.assertEqual(n_list[-2], g_list[-2])

            np.testing.assert_almost_equal(
                np.array([float(t) for t in n_list[1:-3]]),
                np.array([float(t) for t in g_list[1:-3]]),
                decimal=4
            )

def test_geotransform_projection_nodata(self):
        tmp_output = tempfile.mktemp(suffix='.tif')
        extents = [str(s) for s in [920000, -4300000, 929000, -4290000]]

        # the mock is already geotransformed, so this will have no effect
        # to projection and nodata, but will be cropped making the
        # geotransform tuple different
        crop.crop_reproject_resample(self.std2000_no_mask, tmp_output,
                                     sampling='bilinear',
                                     extents=extents,
                                     reproject=True)

        ds = gdal.Open(tmp_output)
        gt = ds.GetGeoTransform()
        projection = ds.GetProjection()
        nodata = ds.GetRasterBand(1).GetNoDataValue()
        ds = None

        ds = gdal.Open(self.std2000_no_mask)
        projection_input = ds.GetProjection()
        nodata_input = ds.GetRasterBand(1).GetNoDataValue()
        ds = None

        self.assertEqual(nodata, nodata_input)
        self.assertEqual(projection, projection_input)
        self.assertEqual(gt[1], float(crop.OUTPUT_RES[0]))
        self.assertEqual(gt[0], float(extents[0]))
        self.assertEqual(gt[3], float(extents[3]))
        os.remove(tmp_output)

def test_apply_mask(self):
        output_file = tempfile.mktemp(suffix='.tif')
        jpeg = False
        tmp_out_file = tempfile.mktemp(suffix='.tif')
        shutil.copy(self.std2000_no_mask, tmp_out_file)
        crop.apply_mask(mask_file=self.mask,
                        tmp_output_file=tmp_out_file,
                        output_file=output_file,
                        jpeg=jpeg)

        ds = gdal.Open(output_file)
        gt = ds.GetGeoTransform()
        projection = ds.GetProjection()
        nodata = ds.GetRasterBand(1).GetNoDataValue()
        ds = None

        ds = gdal.Open(self.std2000)
        projection_input = ds.GetProjection()
        nodata_input = ds.GetRasterBand(1).GetNoDataValue()
        ds = None

        self.assertEqual(nodata, nodata_input)
        self.assertEqual(projection, projection_input)
        self.assertEqual(gt[1], float(crop.OUTPUT_RES[0]))
        self.assertEqual(gt[0], self.extents[0])
        self.assertEqual(gt[3], self.extents[3])
        os.remove(output_file)

def test_mpi_vs_serial(self):

        tmpdir1 = tempfile.mkdtemp()
        tmpdir2 = tempfile.mkdtemp()

        for n, partitions in product(range(1, 3), range(2, 100, 10)):
            size = 2 * n + 1
            raster_average.treat_file(self.test_tif,
                                      out_dir=tmpdir1,
                                      size=size,
                                      func='nanmean',
                                      partitions=partitions)
            arr1 = gdal.Open(
                join(tmpdir1, basename(self.test_tif))).ReadAsArray()

            raster_average.treat_file(self.test_tif,
                                      out_dir=tmpdir2,
                                      size=size,
                                      func='nanmean',
                                      partitions=partitions)
            arr2 = gdal.Open(join(tmpdir2,
                                  basename(self.test_tif))).ReadAsArray()

            np.testing.assert_array_almost_equal(arr1, arr2)

        shutil.rmtree(tmpdir2)
        shutil.rmtree(tmpdir1)

def apply_mask(mask_file, tmp_output_file, output_file, jpeg):
    """
    Parameters
    ----------
    mask_file: mask file path
    tmp_output_file: intermediate cropped geotiff before mask application
    output_file: output geotiff path
    jpeg: boolean, whether to produce jpeg or not
    -------

    """
    mask = get_mask(mask_file)

    out_ds = gdal.Open(tmp_output_file, gdal.GA_Update)
    out_band = out_ds.GetRasterBand(1)
    out_data = out_band.ReadAsArray()
    no_data_value = out_band.GetNoDataValue()
    log.info('Found NoDataValue {} for file {}'.format(
        no_data_value, os.path.basename(tmp_output_file)))
    if no_data_value is not None:
        out_data[mask] = no_data_value
    else:
        log.warning('NoDataValue was not set for {}'.format(tmp_output_file))
        log.info('Manually setting NoDataValue for {}'.format(tmp_output_file))
    out_band.WriteArray(out_data)
    out_ds = None  # close dataset and flush cache

    # move file to output file
    shutil.move(tmp_output_file, output_file)
    log.info('Output file {} created'.format(tmp_output_file))

    if jpeg:
        dir_name = os.path.dirname(output_file)
        jpeg_file = os.path.basename(output_file).split('.')[0] + '.jpg'
        jpeg_file = os.path.join(dir_name, jpeg_file)
        cmd_jpg = ['gdal_translate', '-ot', 'Byte', '-of', 'JPEG', '-scale',
                   output_file,
                   jpeg_file] + COMMON
        subprocess.check_call(cmd_jpg)
        log.info('Created {}'.format(jpeg_file))

def average(input_dir, out_dir, size):
    input_dir = abspath(input_dir)
    log.info('Reading tifs from {}'.format(input_dir))
    tifs = glob.glob(join(input_dir, '*.tif'))

    for tif in tifs:
        data_source = gdal.Open(tif, gdal.GA_ReadOnly)
        band = data_source.GetRasterBand(1)
        # data_type = gdal.GetDataTypeName(band.DataType)
        data = band.ReadAsArray()
        no_data_val = band.GetNoDataValue()
        averaged_data = filter_data(data, size, no_data_val)
        log.info('Calculated average for {}'.format(basename(tif)))

        output_file = join(out_dir, 'average_' + basename(tif))
        out_ds = gdal.GetDriverByName('GTiff').Create(
            output_file, data_source.RasterXSize, data_source.RasterYSize,
            1, band.DataType)
        out_band = out_ds.GetRasterBand(1)
        out_band.WriteArray(averaged_data)
        out_ds.SetGeoTransform(data_source.GetGeoTransform())
        out_ds.SetProjection(data_source.GetProjection())
        out_band.FlushCache()  # Write data to disc
        out_ds = None  # close out_ds
        data_source = None  # close dataset

        log.info('Finished converting {}'.format(basename(tif)))

def gdalaverage(input_dir, out_dir, size):
    """
    average data using gdal's averaging method.
    Parameters
    ----------
    input_dir: str
        input dir name of the tifs that needs to be averaged
    out_dir: str
        output dir name
    size: int, optional
        size of kernel
    Returns
    -------

    """
    input_dir = abspath(input_dir)
    log.info('Reading tifs from {}'.format(input_dir))
    tifs = glob.glob(join(input_dir, '*.tif'))

    process_tifs = np.array_split(tifs, mpiops.chunks)[mpiops.chunk_index]

    for tif in process_tifs:
        data_set = gdal.Open(tif, gdal.GA_ReadOnly)
        # band = data_set.GetRasterBand(1)
        # data_type = gdal.GetDataTypeName(band.DataType)
        # data = band.ReadAsArray()
        # no_data_val = band.GetNoDataValue()
        # averaged_data = filter_data(data, size, no_data_val)
        log.info('Calculated average for {}'.format(basename(tif)))

        output_file = join(out_dir, 'average_' + basename(tif))
        src_gt = data_set.GetGeoTransform()
        tmp_file = '/tmp/tmp_{}.tif'.format(mpiops.chunk_index)
        resample_cmd = [TRANSLATE] + [tif, tmp_file] + \
            ['-tr', str(src_gt[1]*size), str(src_gt[1]*size)] + \
            ['-r', 'bilinear']
        check_call(resample_cmd)
        rollback_cmd = [TRANSLATE] + [tmp_file, output_file] + \
            ['-tr', str(src_gt[1]), str(src_gt[1])]
        check_call(rollback_cmd)
        log.info('Finished converting {}'.format(basename(tif)))

def get_stats(tif, partitions=100):
    ds = gdal.Open(tif, gdal.GA_ReadOnly)
    number_of_bands = ds.RasterCount

    if ds.RasterCount > 1:
        d = {}
        log.info('Found multibanded geotif {}'.format(basename(tif)))
        for b in range(number_of_bands):
            d['{tif}_{b}'.format(tif=tif, b=b)] = band_stats(ds, tif, b + 1,
                                                             partitions)
        return d
    else:
        log.info('Found single band geotif {}'.format(basename(tif)))
        return {tif: band_stats(ds, tif, 1, partitions)}

def get_numpy_stats(tif, writer):
    ds = gdal.Open(tif, gdal.GA_ReadOnly)
    number_of_bands = ds.RasterCount

    if ds.RasterCount > 1:
        log.info('Found multibanded geotif {}'.format(basename(tif)))
        for b in range(number_of_bands):
            write_rows(stats=number_of_bands(ds, tif, b + 1), writer=writer)
    else:
        log.info('Found single band geotif {}'.format(basename(tif)))
        write_rows(stats=number_of_bands(ds, tif, 1), writer=writer)

def VPMprmt (directory,tile):
    tempfile=None
    for path, subdirs, files in os.walk(directory):
        for name in files:
            if (tile in name) and (".tif" == name[-4:]) : tempfile=os.path.join(path,name)
    print tempfile
    inprmtdata=gdal.Open(tempfile)
    prmtdata=inprmtdata.ReadAsArray()
    return prmtdata


#Function to build vrt

def movingaverage(tile):
    # Output directories for moving average LSWImax

    # if the output directories don't exist, create the new directories
    if not os.path.exists(dirLSWIMA):
        os.makedirs(dirLSWIMA)
    # build LSWImax vrt file and read as an array
    file=buildVrtFile(year,tile)
    vrtLSWImax=os.path.join(os.path.dirname(file),year+tile+'LSWImax_vrt.vrt')
    #print "Building the vrt file: ", year+tile+vrtLSWImax
    os.system('gdalbuildvrt -separate -input_file_list '+file+' '+vrtLSWImax)

    global rows, cols, geoProj,geoTran
    inLSWImax=gdal.Open(vrtLSWImax)
    #print "reading the multi-LSWI..."
    LSWImax=inLSWImax.ReadAsArray()

    rows = 2400
    cols = 2400
    geoTran=inLSWImax.GetGeoTransform()
    geoProj=inLSWImax.GetProjection()

    #find the second largest LSWImax
    secLSWImax = numpy.sort(LSWImax, axis=0, kind='quicksort', order=None)[3,:,:]

    write_file(dirLSWIMA+'/'+tile+'.'+year+'.maxLSWI_MA.tif',secLSWImax,geoTran,rows,cols,geoProj,driverName='GTiff')

    secLSWImax=None
    LSWImax=None

def movingaverage(tile):
    # Output directories for moving average LSWImax

    # if the output directories don't exist, create the new directories
    if not os.path.exists(dirLSWIMA):
        os.makedirs(dirLSWIMA)
    # build LSWImax vrt file and read as an array
    file=buildVrtFile(year,tile)
    vrtLSWImax=os.path.join(os.path.dirname(file),year+tile+'LSWImax_vrt.vrt')
    #print "Building the vrt file: ", year+tile+vrtLSWImax
    os.system('gdalbuildvrt -separate -input_file_list '+file+' '+vrtLSWImax)

    global rows, cols, geoProj,geoTran
    inLSWImax=gdal.Open(vrtLSWImax)
    #print "reading the multi-LSWI..."
    LSWImax=inLSWImax.ReadAsArray()

    rows = 2400
    cols = 2400
    geoTran=inLSWImax.GetGeoTransform()
    geoProj=inLSWImax.GetProjection()

    #find the second largest LSWImax
    secLSWImax = numpy.sort(LSWImax, axis=0, kind='quicksort', order=None)[3,:,:]

    write_file(dirLSWIMA+'/'+tile+'.'+year+'.maxLSWI_MA.tif',secLSWImax,geoTran,rows,cols,geoProj,driverName='GTiff')

    secLSWImax=None
    LSWImax=None

def buildTindex(rasterFolder, rasterExtention='.tif'):
    rasterList = glob.glob(os.path.join(rasterFolder, '*{}'.format(rasterExtention)))
    print(rasterList)

    print(os.path.join(rasterFolder, '*{}'.format(rasterExtention)))

    memDriver = ogr.GetDriverByName('MEMORY')
    gTindex = memDriver.CreateDataSource('gTindex')
    srcImage = gdal.Open(rasterList[0])
    spat_ref = osr.SpatialReference()
    spat_ref.SetProjection(srcImage.GetProjection())
    gTindexLayer = gTindex.CreateLayer("gtindexlayer", spat_ref, geom_type=ogr.wkbPolygon)

    # Add an ID field
    idField = ogr.FieldDefn("location", ogr.OFTString)
    gTindexLayer.CreateField(idField)

    # Create the feature and set values
    featureDefn = gTindexLayer.GetLayerDefn()



    for rasterFile in rasterList:
        srcImage = gdal.Open(rasterFile)

        geoTrans, polyToCut, ulX, ulY, lrX, lrY = gT.getRasterExtent(srcImage)

        feature = ogr.Feature(featureDefn)
        feature.SetGeometry(polyToCut)
        feature.SetField("location", rasterFile)
        gTindexLayer.CreateFeature(feature)
        feature = None


    return gTindex, gTindexLayer

def createTiledGeoJsonFromSrc(rasterFolderLocation, vectorSrcFile, geoJsonOutputDirectory, rasterTileIndex='',
                              geoJsonPrefix='GEO', rasterFileExtenstion='.tif',
                              rasterPrefixToReplace='PAN'
                              ):
    if rasterTileIndex == '':
        gTindex, gTindexLayer = buildTindex(rasterFolderLocation, rasterExtention=rasterFileExtenstion)
    else:
        gTindex = ogr.Open(rasterTileIndex,0)
        gTindexLayer = gTindex.GetLayer()

    shapeSrc = ogr.Open(vectorSrcFile,0)
    chipSummaryList = []
    for feature in gTindexLayer:
        featureGeom = feature.GetGeometryRef()
        rasterFileName = feature.GetField('location')
        rasterFileBaseName = os.path.basename(rasterFileName)
        outGeoJson = rasterFileBaseName.replace(rasterPrefixToReplace, geoJsonPrefix)
        outGeoJson = outGeoJson.replace(rasterFileExtenstion, '.geojson')
        outGeoJson = os.path.join(geoJsonOutputDirectory, outGeoJson)

        gT.clipShapeFile(shapeSrc, outGeoJson, featureGeom, minpartialPerc=0.0, debug=False)
        imageId = rasterFileBaseName.replace(rasterPrefixToReplace+"_", "")
        chipSummary = {'chipName': rasterFileName,
                           'geoVectorName': outGeoJson,
                           'imageId': os.path.splitext(imageId)[0]}

        chipSummaryList.append(chipSummary)

    return chipSummaryList

def createTruthPixelLinePickle(truthLineFile, pickleLocation=''):
    if pickleLocation=='':
        extension = os.path.splitext(truthLineFile)[1]
        pickleLocation = truthLineFile.replace(extension, 'Pixline.p')
    if truthLineFile != '':
        # get Source Line File Information
        shapef = ogr.Open(truthLineFile, 0)
        truthLayer = shapef.GetLayer()
        pt1X = []
        pt1Y = []
        pt2X = []
        pt2Y = []
        for tmpFeature in truthLayer:
            tmpGeom = tmpFeature.GetGeometryRef()
            for i in range(0, tmpGeom.GetPointCount()):
                pt = tmpGeom.GetPoint(i)

                if i == 0:
                    pt1X.append(pt[0])
                    pt1Y.append(pt[1])
                elif i == 1:
                    pt2X.append(pt[0])
                    pt2Y.append(pt[1])

        lineData = {'pt1X': np.asarray(pt1X),
                    'pt1Y': np.asarray(pt1Y),
                    'pt2X': np.asarray(pt2X),
                    'pt2Y': np.asarray(pt2Y)
                    }

        with open(pickleLocation, 'wb') as f:
            pickle.dump(lineData, f)
            # get Source Line File Information

def createTruthPixelPolyPickle(truthPoly, pickleLocation=''):
    # returns dictionary with list of minX, maxX, minY, maxY

    if pickleLocation=='':
        extension = os.path.splitext(truthPoly)[1]
        pickleLocation = truthPoly.replace(extension, 'PixPoly.p')
    if truthPoly != '':
        # get Source Line File Information
        shapef = ogr.Open(truthPoly, 0)
        truthLayer = shapef.GetLayer()
        envList = []

        for tmpFeature in truthLayer:
            tmpGeom = tmpFeature.GetGeometryRef()
            env = tmpGeom.GetEvnelope()
            envList.append(env)

        envArray = np.asarray(envList)
        envelopeData = {'minX': envArray[:,0],
                        'maxX': envArray[:,1],
                        'minY': envArray[:,2],
                        'maxY': envArray[:,3]
                        }


        with open(pickleLocation, 'wb') as f:
            pickle.dump(envelopeData, f)
            # get Source Line File Information

def openm3(input_data):
    if input_data.split('.')[-1] == 'hdr':
        # GDAL wants the img, but many users aim at the .hdr
        input_data = input_data.split('.')[0] + '.img'
    ds = gdal.Open(input_data)
    ref_array = ds.GetRasterBand(1).ReadAsArray()
    metadata = ds.GetMetadata()
    wv_array = metadatatoband(metadata)
    return wv_array, ref_array, ds

def openmi(input_data):
    ds = gdal.Open(input_data)
    band_pointers = []
    nbands = ds.RasterCount

    for b in xrange(1, nbands + 1):
        band_pointers.append(ds.GetRasterBand(b))

    ref_array = ds.GetRasterBand(1).ReadAsArray()
    wv_array = None
    return wv_array, ref_array[::3, ::3], ds

def readBandMetaData(inputFile, band, name):
    # Open the dataset in Read Only mode
    dataset = gdal.Open(inputFile, gdal.GA_ReadOnly)
    # Check that the image has opened
    if not dataset is None:
        # Get the image band
        imgBand = dataset.GetRasterBand(band)
        # Check the image band was available
        if not imgBand is None:
            # Get the meta data value specified
            metaData = imgBand.GetMetaDataItem(name)
            # Check that it is present
            if metaData == None:
                # If not present, print error.
                print("Could not find \'", name, "\' item.")
            else:
                # Else print out the metadata value
                print(name, "=\'", metaData, "\'")

        else:
            # Print out an error message.
            print("Could not open the image band: ", band)

    else:
        # Print an error message if the file
        # could not be opened.
        print("Could not open the input image file: ", inputFile)

# A function to read a meta-data item
# from an image