91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
import docker
|
|
import os
|
|
import logging
|
|
from collections import deque
|
|
|
|
|
|
class DockerRunner:
|
|
def __init__(self, project_path: str):
|
|
"""
|
|
初始化 DockerRunner
|
|
|
|
Args:
|
|
project_path (str): 项目路径，将挂载到 Docker 容器中
|
|
"""
|
|
self.project_path = project_path
|
|
self.logger = logging.getLogger("UAV_Preprocess.DockerRunner")
|
|
self.docker_client = docker.from_env()
|
|
|
|
def run_odm_container(self):
|
|
"""
|
|
使用 Docker SDK 运行 OpenDroneMap 容器
|
|
"""
|
|
try:
|
|
self.logger.info("开始运行docker run指令")
|
|
# 挂载路径
|
|
volume_mapping = {
|
|
self.project_path: {
|
|
'bind': '/datasets',
|
|
'mode': 'rw'
|
|
}
|
|
}
|
|
|
|
# Docker 命令参数
|
|
command = [
|
|
"--project-path", "/datasets",
|
|
"project",
|
|
"--max-concurrency", "15",
|
|
"--force-gps",
|
|
"--split-overlap", "0",
|
|
"--rerun-all"
|
|
]
|
|
|
|
# 运行容器
|
|
container = self.docker_client.containers.run(
|
|
image="opendronemap/odm:gpu",
|
|
command=command,
|
|
volumes=volume_mapping,
|
|
device_requests=[
|
|
docker.types.DeviceRequest(
|
|
count=-1, capabilities=[["gpu"]])
|
|
], # 添加 GPU 支持
|
|
remove=False, # 容器运行结束后不自动删除，便于获取日志
|
|
tty=True,
|
|
detach=True # 后台运行
|
|
)
|
|
|
|
# 等待容器运行完成
|
|
exit_status = container.wait()
|
|
if exit_status["StatusCode"] != 0:
|
|
self.logger.error(f"容器运行失败，退出状态码: {exit_status['StatusCode']}")
|
|
|
|
# 获取容器的错误日志
|
|
error_logs = container.logs(
|
|
stderr=True).decode("utf-8").splitlines()
|
|
self.logger.error("容器运行失败的详细错误日志：")
|
|
for line in error_logs:
|
|
self.logger.error(line)
|
|
|
|
else:
|
|
# 获取所有日志
|
|
logs = container.logs().decode("utf-8").splitlines()
|
|
|
|
# 输出最后 50 行日志
|
|
self.logger.info("容器运行完成，以下是最后 50 行日志：")
|
|
for line in logs[-50:]:
|
|
self.logger.info(line)
|
|
|
|
# 删除容器
|
|
container.remove()
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"运行 Docker 容器时发生错误: {str(e)}", exc_info=True)
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 示例用法
|
|
project_path = r"E:\datasets\UAV\199"
|
|
docker_runner = DockerRunner(project_path)
|
|
docker_runner.run_odm_container()
|