pytexrecon.py

import subprocess
import os
import tempfile
from typing import Optional, Dict, Tuple

class TexRecon:
    def __init__(self, verbose: bool = True):
        """初始化TexRecon绑定，适配mvs-texturing项目和Hugging Face Space环境"""
        self.verbose = verbose
        self.texrecon_path = self._find_texrecon()
        
        # 检查是否可用，若未找到则尝试自动编译
        if not self.texrecon_path:
            # Hugging Face Space特征路径 
            self.texrecon_path = self._compile_texrecon()
        
        self._check_texrecon()
        
        # 创建临时目录用于处理中间文件
        self.temp_dir = tempfile.TemporaryDirectory()
        self.work_dir = self.temp_dir.name
    
    def _find_texrecon(self) -> Optional[str]:
        """在常见路径中查找mvs-texturing编译的texrecon可执行文件"""
        # 重点：mvs-texturing的可执行文件在build/apps/texrecon路径下
        search_paths = [
            "./mvs-texturing/build/apps/texrecon/texrecon",
            "/home/user/app/mvs-texturing/build/apps/texrecon/texrecon",
            os.path.expanduser("~/mvs-texturing/build/apps/texrecon/texrecon")
        ]
        
        for path in search_paths:
            if os.path.exists(path) and os.access(path, os.X_OK):
                if self.verbose:
                    print(f"找到TexRecon可执行文件：{path}")
                return path
        return None
    
    def _compile_texrecon(self) -> str:
        """在Hugging Face Space中自动编译mvs-texturing项目"""
        try:
            if self.verbose:
                print("=== 开始自动编译mvs-texturing ===")
            
            # 1. 安装项目所需依赖（根据官方文档）
            if self.verbose:
                print("=== 安装编译依赖 ===")
            deps = [
                "git", "cmake", "make", "g++",  # 基础编译工具
                "libpng-dev", "libjpeg-dev", "libtiff-dev", "libtbb-dev"  # 图像和并行库
            ]
            
            # Space环境无需sudo
            subprocess.run(
                ["apt-get", "update"],
                check=True,
                stdout=subprocess.PIPE if not self.verbose else None,
                stderr=subprocess.STDOUT if not self.verbose else None
            )
            subprocess.run(
                ["apt-get", "install", "-y"] + deps,
                check=True,
                stdout=subprocess.PIPE if not self.verbose else None,
                stderr=subprocess.STDOUT if not self.verbose else None
            )
            
            # 2. 克隆mvs-texturing源码（使用官方仓库）
            if self.verbose:
                print("=== 克隆mvs-texturing源码 ===")
            repo_dir = "/home/user/app/mvs-texturing"
            src_dir = os.path.join(repo_dir, "src")
            build_dir = os.path.join(repo_dir, "build")
            
            # 确保目录存在
            os.makedirs(repo_dir, exist_ok=True)
            
            # 若目录为空则克隆
            if not os.listdir(repo_dir):
                subprocess.run(
                    ["git", "clone", "https://github.com/nmoehrle/mvs-texturing.git", repo_dir],
                    check=True,
                    stdout=subprocess.PIPE if not self.verbose else None,
                    stderr=subprocess.STDOUT if not self.verbose else None
                )
            else:
                if self.verbose:
                    print("mvs-texturing源码已存在，跳过克隆")
            
            # 3. 按照官方文档步骤编译
            if self.verbose:
                print("=== 编译mvs-texturing ===")
            os.makedirs(build_dir, exist_ok=True)
            
            # 运行cmake（官方命令）
            subprocess.run(
                ["cmake", ".."],
                cwd=build_dir,
                check=True,
                stdout=subprocess.PIPE if not self.verbose else None,
                stderr=subprocess.STDOUT if not self.verbose else None
            )
            
            # 运行make（官方命令，-j4并行编译）
            subprocess.run(
                ["make", "-j4"],
                cwd=build_dir,
                check=True,
                stdout=subprocess.PIPE if not self.verbose else None,
                stderr=subprocess.STDOUT if not self.verbose else None
            )
            
            # 4. 验证可执行文件（mvs-texturing的texrecon在build/apps/texrecon下）
            texrecon_path = os.path.join(build_dir, "apps", "texrecon", "texrecon")
            if os.path.exists(texrecon_path) and os.access(texrecon_path, os.X_OK):
                if self.verbose:
                    print(f"=== mvs-texturing编译成功：{texrecon_path} ===")
                return texrecon_path
            else:
                raise RuntimeError(f"编译成功但未找到可执行文件：{texrecon_path}")
                
        except subprocess.CalledProcessError as e:
            error_msg = f"命令执行失败（返回码：{e.returncode}）"
            if self.verbose and e.output:
                error_msg += f"\n输出：{e.output.decode('utf-8', errors='ignore')}"
            raise RuntimeError(error_msg)
        except Exception as e:
            raise RuntimeError(f"自动编译失败：{str(e)}")
    
    def _check_texrecon(self) -> None:
        """检查TexRecon是否可执行"""
        try:
            subprocess.run(
                [self.texrecon_path, "--help"],
                check=True,
                stdout=subprocess.PIPE if not self.verbose else None,
                stderr=subprocess.STDOUT if not self.verbose else None
            )
        except FileNotFoundError:
            raise RuntimeError(f"未找到TexRecon可执行文件：{self.texrecon_path}")
        except subprocess.CalledProcessError:
            raise RuntimeError(f"TexRecon可执行文件无效（可能编译不完整）：{self.texrecon_path}")
    
    def run_texturing(self,
                      mesh_data: bytes,
                      images: Dict[str, bytes],
                      intrinsics: str,
                      poses: Dict[str, str],
                      texture_resolution: int = 2048,
                      padding: int = 5,
                      fill_holes: bool = True) -> Tuple[bytes, bytes]:
        """运行纹理重建（适配mvs-texturing的参数要求）"""
        # 创建工作目录结构
        mesh_path = os.path.join(self.work_dir, "mesh.ply")
        images_dir = os.path.join(self.work_dir, "images")
        intrinsics_path = os.path.join(self.work_dir, "intrinsics.txt")
        poses_dir = os.path.join(self.work_dir, "poses")
        output_dir = os.path.join(self.work_dir, "output")
        
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(poses_dir, exist_ok=True)
        os.makedirs(output_dir, exist_ok=True)
        
        # 写入输入文件
        with open(mesh_path, "wb") as f:
            f.write(mesh_data)
        
        for img_name, img_data in images.items():
            img_path = os.path.join(images_dir, img_name)
            with open(img_path, "wb") as f:
                f.write(img_data)
        
        with open(intrinsics_path, "w") as f:
            f.write(intrinsics)
        
        for img_name, pose_data in poses.items():
            pose_path = os.path.join(poses_dir, f"{img_name}.txt")
            with open(pose_path, "w") as f:
                f.write(pose_data)
        
        # 构建mvs-texturing的命令（符合官方文档参数）
        cmd = [
            self.texrecon_path,
            "--mesh", mesh_path,
            "--images", images_dir,
            "--intrinsics", intrinsics_path,
            "--poses", poses_dir,
            "--output", output_dir,
            "--resolution", str(texture_resolution),
            "--padding", str(padding)
        ]
        
        if fill_holes:
            cmd.append("--fill-holes")
        
        try:
            subprocess.run(
                cmd,
                stdout=subprocess.PIPE if not self.verbose else None,
                stderr=subprocess.STDOUT if not self.verbose else None,
                check=True
            )
            
            # mvs-texturing的输出文件命名
            textured_mesh_path = os.path.join(output_dir, "mesh_textured.obj")
            texture_path = os.path.join(output_dir, "texture_0.png")
            
            # 读取输出结果
            with open(textured_mesh_path, "rb") as f:
                mesh_bytes = f.read()
            
            with open(texture_path, "rb") as f:
                texture_bytes = f.read()
                
            return mesh_bytes, texture_bytes
            
        except subprocess.CalledProcessError as e:
            error_msg = f"纹理重建命令失败（返回码：{e.returncode}）"
            if self.verbose and e.output:
                error_msg += f"\n输出：{e.output.decode('utf-8', errors='ignore')}"
            raise RuntimeError(error_msg)
    
    def __del__(self):
        """清理临时目录"""
        if hasattr(self, 'temp_dir'):
            self.temp_dir.cleanup()