# -*- encoding: utf-8 -*-

from osgeo import gdal
from osgeo import osr
import numpy as np

def getSRSPair(dataset):
    '''
    獲得給定數(shù)據(jù)的投影參考系和地理參考系
    :param dataset: GDAL地理數(shù)據(jù)
    :return: 投影參考系和地理參考系
    '''
    prosrs = osr.SpatialReference()
    prosrs.ImportFromWkt(dataset.GetProjection())
    geosrs = prosrs.CloneGeogCS()
    return prosrs, geosrs

def geo2lonlat(dataset, x, y):
    '''
    將投影坐標(biāo)轉(zhuǎn)為經(jīng)緯度坐標(biāo)（具體的投影坐標(biāo)系由給定數(shù)據(jù)確定）
    :param dataset: GDAL地理數(shù)據(jù)
    :param x: 投影坐標(biāo)x
    :param y: 投影坐標(biāo)y
    :return: 投影坐標(biāo)(x, y)對(duì)應(yīng)的經(jīng)緯度坐標(biāo)(lon, lat)
    '''
    prosrs, geosrs = getSRSPair(dataset)
    ct = osr.CoordinateTransformation(prosrs, geosrs)
    coords = ct.TransformPoint(x, y)
    return coords[:2]


def lonlat2geo(dataset, lon, lat):
    '''
    將經(jīng)緯度坐標(biāo)轉(zhuǎn)為投影坐標(biāo)（具體的投影坐標(biāo)系由給定數(shù)據(jù)確定）
    :param dataset: GDAL地理數(shù)據(jù)
    :param lon: 地理坐標(biāo)lon經(jīng)度
    :param lat: 地理坐標(biāo)lat緯度
    :return: 經(jīng)緯度坐標(biāo)(lon, lat)對(duì)應(yīng)的投影坐標(biāo)
    '''
    prosrs, geosrs = getSRSPair(dataset)
    ct = osr.CoordinateTransformation(geosrs, prosrs)
    coords = ct.TransformPoint(lon, lat)
    return coords[:2]

def imagexy2geo(dataset, row, col):
    '''
    根據(jù)GDAL的六參數(shù)模型將影像圖上坐標(biāo)（行列號(hào)）轉(zhuǎn)為投影坐標(biāo)或地理坐標(biāo)（根據(jù)具體數(shù)據(jù)的坐標(biāo)系統(tǒng)轉(zhuǎn)換）
    :param dataset: GDAL地理數(shù)據(jù)
    :param row: 像素的行號(hào)
    :param col: 像素的列號(hào)
    :return: 行列號(hào)(row, col)對(duì)應(yīng)的投影坐標(biāo)或地理坐標(biāo)(x, y)
    '''
    trans = dataset.GetGeoTransform()
    px = trans[0] + col * trans[1] + row * trans[2]
    py = trans[3] + col * trans[4] + row * trans[5]
    return px, py


def geo2imagexy(dataset, x, y):
    '''
    根據(jù)GDAL的六 參數(shù)模型將給定的投影或地理坐標(biāo)轉(zhuǎn)為影像圖上坐標(biāo)（行列號(hào)）
    :param dataset: GDAL地理數(shù)據(jù)
    :param x: 投影或地理坐標(biāo)x
    :param y: 投影或地理坐標(biāo)y
    :return: 影坐標(biāo)或地理坐標(biāo)(x, y)對(duì)應(yīng)的影像圖上行列號(hào)(row, col)
    '''
    trans = dataset.GetGeoTransform()
    a = np.array([[trans[1], trans[2]], [trans[4], trans[5]]])
    b = np.array([x - trans[0], y - trans[3]])
    return np.linalg.solve(a, b)  # 使用numpy的linalg.solve進(jìn)行二元一次方程的求解


if __name__ == '__main__':
    gdal.AllRegister()
    dataset = gdal.Open(r"F:\2016\Data\Great Khingan\DEM\Projection\strm_6102_UTM.tif")
    print('數(shù)據(jù)投影：')
    print(dataset.GetProjection())
    print('數(shù)據(jù)的大小（行，列）：')
    print('(%s %s)' % (dataset.RasterYSize, dataset.RasterXSize))

    x = 464201
    y = 5818760
    lon = 122.47242
    lat = 52.51778
    row = 2399
    col = 3751

    print('投影坐標(biāo) -> 經(jīng)緯度：')
    coords = geo2lonlat(dataset, x, y)
    print('(%s, %s)->(%s, %s)' % (x, y, coords[0], coords[1]))
    print('經(jīng)緯度 -> 投影坐標(biāo)：')
    coords = lonlat2geo(dataset, lon, lat)
    print('(%s, %s)->(%s, %s)' % (lon, lat, coords[0], coords[1]))

    print('像素坐標(biāo) -> 投影坐標(biāo)：')
    coords = imagexy2geo(dataset, row, col)
    print('(%s, %s)->(%s, %s)' % (row, col, coords[0], coords[1]))
    print('投影坐標(biāo) -> 像素坐標(biāo)：')
    coords = geo2imagexy(dataset, x, y)
    print('(%s, %s)->(%s, %s)' % (x, y, coords[0], coords[1]))