# __author: GeoTech_QingQian`s Daddy
# __data: 2021-07-23

import Metashape
import os
import time

'''
    0. 预设常量
        WORK_SPACE： 工作空间，存放工程目录，数据输出位置（注意存储空间大小）；
        DATA_DIRECTORY： 数据目录，指向原始影像的父级文件夹路径。
        以上注意路径中反斜杠的转义
'''
WORK_SPACE = "E:\\workspace"
DATA_DIRECTORY = "E:\\100_0033"

def log(ctx):
        with open(os.path.join(WORK_SPACE, "log","log"+ current_time +".txt"), "a") as f:
            f.write("%s     %s\n" %(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),ctx))
        f.close()

def logEmpty():
        with open(os.path.join(WORK_SPACE, "log","log"+ current_time +".txt"), "a") as f:
            f.write("\n")
        f.close()

'''1. 创建 Metashape app object'''

# 授时
current_time= time.strftime("%Y_%m_%d_%H_%M_%S", time.localtime())

# 创建目录文件
if os.path.exists(WORK_SPACE):
    if os.path.exists(os.path.join(WORK_SPACE, "tileZip")):
        pass
    else:
        os.mkdir(os.path.join(WORK_SPACE, "tileZip"))
    
    if os.path.exists(os.path.join(WORK_SPACE, "tifFile")):
        pass
    else:
        os.mkdir(os.path.join(WORK_SPACE, "tifFile"))

    if os.path.exists(os.path.join(WORK_SPACE, "log")):
        pass
    else:
        os.mkdir(os.path.join(WORK_SPACE, "log"))

else:
    os.mkdir(WORK_SPACE)
    os.mkdir(os.path.join(WORK_SPACE, "tileZip"))
    os.mkdir(os.path.join(WORK_SPACE, "tifFile"))

doc = Metashape.app.document

app = Metashape.Application()

projection_name = current_time + ".psx"

doc.save(os.path.join(WORK_SPACE, projection_name))
# 初始化日志文件
log("success    创建 Metashape app object")
logEmpty()


'''2.初始化, init chunk
        2.1 移除原始chunk
        2.2 遍历文件夹，并添加chunk
'''


# 2.1 移除原始chunk
for chunk in doc.chunks:
    if chunk:
        doc.remove(chunk)

    log("success     移除原始chunk")
    logEmpty()


# 2.2 遍历文件夹，并添加chunk
for root, dirs, files in os.walk(DATA_DIRECTORY):
    if root != DATA_DIRECTORY:

        chunk = doc.addChunk()  # 添加chunk

        chunk.label = root.split('\\')[-1]  # chuunk 以文件夹重命名

        chunk.crs = Metashape.CoordinateSystem("EPSG::4326")  # 定义坐标框架

        log("success     " + chunk.label + "_____chunk创建成功")

        # 装配文件路径列表
        photoList = []
        for file in files:
            photoList.append(os.path.join(root, file))

        # 为 chunk添加照片
        chunk.addPhotos(photoList)

        log("success     " + chunk.label + "______chunk相片添加成功")

doc.save()
logEmpty()


''' 
    3.空三处理: algin photos
        3.1 匹配照片
        3.2 对齐相机
        3.3 建立深度
'''
for chunk in doc.chunks:
    try:
        # 3.1 匹配照片
        chunk.matchPhotos(downscale=8, generic_preselection = True, reference_preselection=False)
        # 特征点提取精度  downscale [1（原图提取特征点）, 2（图像缩小到1/2提取导入特征点）.....]

        log("success     " + chunk.label + "______照片匹配成功")

        #3.2 对齐相机
        chunk.alignCameras()

        log("success     " + chunk.label + "______照片对齐成功")

        #3.3 建立深度图
        chunk.buildDepthMaps(downscale=8, filter_mode=Metashape.AggressiveFiltering)
        # 深度贴图  downscale [1（原图贴图）, 2（图像缩小到1/2贴图）.....]

        log("success     " + chunk.label + "______照片深度建立成功")

    except:
        log("success     " + chunk.label + "______空三阶段失败") 

doc.save()
logEmpty()


'''
    4. 建立密集点云 dense cloud
'''
for chunk in doc.chunks:
    try:
        chunk.buildDenseCloud()

        log("success     " + chunk.label + "______密集点云创建成功")
    except:
        log("fail     " + chunk.label + "______密集点云创建失败")
doc.save()
logEmpty()


'''
    5. 建立几何结构 build Model
'''
for chunk in doc.chunks:
    try: 
        chunk.buildModel(surface_type = Metashape.Arbitrary, interpolation = Metashape.EnabledInterpolation)
        
        log("success     " + chunk.label + "______几何结构创建成功")
    except:
        log("fail     " + chunk.label + "______几何结构创建失败")  
doc.save()
logEmpty()


'''
    6. 建立正射数据 build Orthomosaic
'''
for chunk in doc.chunks:
    try:
        
        chunk.buildOrthomosaic(surface_data = Metashape.DataSource.ModelData)
        log("success     " + chunk.label + "______正射影像创建成功")   
    except:
        log("fail     " + chunk.label + "______正射影像创建失败")   
doc.save()
logEmpty()


'''
    7. 导出谷歌切片
'''
for chunk in doc.chunks:
    try:
        EXPORT_FILE = os.path.join(WORK_SPACE, "tileZip",chunk.label)
        zipFile = EXPORT_FILE+".zip"
        chunk.exportRaster(zipFile, Metashape.RasterFormat.RasterFormatXYZ,Metashape.ImageFormat.ImageFormatPNG, min_zoom_level=15, max_zoom_level=23)
        
        log("success     " + chunk.label + "______谷歌切片导出成功")
    except:
        log("fail     " + chunk.label + "______谷歌切片导出失败")


doc.save()
logEmpty()


'''
    8. 导出Tiff影像
'''
for chunk in doc.chunks:
    try:
        EXPORT_FILE = os.path.join(WORK_SPACE, "tifFile",chunk.label)
        tiffFile = EXPORT_FILE+".tif"
        chunk.exportRaster(tiffFile,image_format = Metashape.ImageFormat.ImageFormatTIFF)
        
        log("success     " + chunk.label + "______geoTiff导出成功")
    except:
        log("fail     " + chunk.label + "______geoTiff导出失败")


doc.save()
logEmpty()

log("全部处理完成-----0")

Metashape.app.messageBox("全部处理完成")