from osgeo import ogr
import os
import argparse

def shapefile_to_dota(shp_path, output_path, class_field='class', difficulty_value=1):
    """
    将Shapefile转换为DOTA格式
    :param shp_path: Shapefile文件路径
    :param output_path: 输出目录
    :param class_field: 类别字段名
    :param difficulty_value: 难度默认字段
    """

    # 注册所有驱动
    ogr.RegisterAll()

    # 打开Shapefile文件
    driver = ogr.GetDriverByName('ESRI Shapefile')
    datasource = driver.Open(shp_path, 0)
    if datasource is None:
        print("无法打开Shapefile文件")
        return

    # 获取图层
    layer = datasource.GetLayer()

    output_file = output_path

    with open(output_file, 'w',encoding="utf-8") as f:
        # 写入DOTA格式头信息(可选)
        # f.write('imagesource:unknown\n')
        # f.write('gsd:1.0\n')

        # 遍历所有要素
        for feature in layer:
            # 获取几何对象
            geom = feature.GetGeometryRef()

            # 获取类别和难度
            class_name = feature.GetField(class_field) if feature.GetField(class_field) else 'unknown'
            difficulty = difficulty_value

            # 处理不同类型的几何图形
            if geom.GetGeometryName() == 'POLYGON':
                # 获取多边形外环
                ring = geom.GetGeometryRef(0)
                # 获取所有点
                points = []
                for i in range(ring.GetPointCount()):
                    points.append(ring.GetPoint(i))

                # 确保有足够的点(至少4个)
                if len(points) >= 4:
                    # 取前4个点作为DOTA格式的四个角点
                    # 注意: DOTA要求按顺序排列(顺时针或逆时针)
                    x1, y1 = points[0][0], points[0][1]
                    x2, y2 = points[1][0], points[1][1]
                    x3, y3 = points[2][0], points[2][1]
                    x4, y4 = points[3][0], points[3][1]

                    # 写入DOTA格式行
                    line = f"{x1} {y1} {x2} {y2} {x3} {y3} {x4} {y4} {class_name} {difficulty}\n"
                    f.write(line)

    # 释放资源
    datasource.Destroy()
    print("转换完毕")

def getParams():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i','--infile',type=str,default=r'F:\天仪SAR卫星数据集\德清院-测试-天仪提供数据\bc2-sp-org-vv-20250210t160723-022200-000120-0056b8-01_石佳宁.shp', help='输入shapefile文件')
    parser.add_argument('-o', '--outfile',type=str,default=r'F:\天仪SAR卫星数据集\德清院-测试-天仪提供数据\bc2-sp-org-vv-20250210t160723-022200-000120-0056b8-01_石佳宁.txt', help='输出geojson文件')
    parser.add_argument('-c', '--clsname',type=str,default=r'Label', help='输出geojson文件')
    parser.add_argument('-d', '--difficulty',type=int,default=1, help='输出geojson文件')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    parser = getParams()
    inFilePath=parser.infile
    outpath=parser.outfile
    clsname=parser.clsname
    difficulty=parser.difficulty
    print('infile=',inFilePath)
    print('outfile=',outpath)
    print('clsname=',clsname)
    print('difficulty=',difficulty)
    shapefile_to_dota(inFilePath, outpath, clsname, difficulty)
# 使用示例