File size: 5,854 Bytes
d99a295 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | def export_model(self, export_format="both", progress=gr.Progress()):
"""Export model - supports point cloud, mesh, or both"""
progress(0.90, desc="💾 Exporting 3D model...")
try:
# Find the latest model
models_dir = os.path.join(self.output_dir, "models")
if not os.path.exists(models_dir):
return {"status": "error", "message": "No trained model found"}
# Get the most recent model directory
model_dirs = [d for d in os.listdir(models_dir) if os.path.isdir(os.path.join(models_dir, d))]
if not model_dirs:
return {"status": "error", "message": "No model directories found"}
latest_model = sorted(model_dirs)[-1]
model_path = os.path.join(models_dir, latest_model)
config_path = os.path.join(model_path, "config.yml")
if not os.path.exists(config_path):
return {"status": "error", "message": "Model config not found"}
exports = {}
# Export textured mesh (like the Mantis example!)
if export_format in ["mesh", "both"]:
progress(0.90, desc="🎨 Exporting textured mesh...")
output_mesh = os.path.join(self.output_dir, "textured_mesh.ply")
# Use Poisson surface reconstruction to create mesh
result = subprocess.run([
"ns-export", "poisson",
"--load-config", config_path,
"--output-dir", self.output_dir,
"--target-num-faces", "500000", # 500k faces for good detail
"--num-pixels-per-side", "2048", # 2K texture resolution
"--normal-method", "model_output",
"--save-point-cloud", "False"
], capture_output=True, text=True, timeout=900)
if result.returncode == 0 and os.path.exists(output_mesh):
file_size = os.path.getsize(output_mesh) / (1024 * 1024)
exports["mesh"] = {
"path": output_mesh,
"size_mb": round(file_size, 2),
"type": "Textured Mesh (Poisson)"
}
# Also try to export as OBJ with textures (more compatible)
output_obj = os.path.join(self.output_dir, "textured_mesh.obj")
result_obj = subprocess.run([
"ns-export", "poisson",
"--load-config", config_path,
"--output-dir", self.output_dir,
"--target-num-faces", "500000",
"--num-pixels-per-side", "2048",
"--normal-method", "model_output",
"--save-point-cloud", "False",
"--output-format", "obj"
], capture_output=True, text=True, timeout=900)
if result_obj.returncode == 0 and os.path.exists(output_obj):
exports["mesh_obj"] = {
"path": output_obj,
"type": "OBJ with Texture"
}
# Export point cloud (your current method)
if export_format in ["pointcloud", "both"]:
progress(0.92, desc="☁️ Exporting point cloud...")
output_ply = os.path.join(self.output_dir, "point_cloud.ply")
result = subprocess.run([
"ns-export", "pointcloud",
"--load-config", config_path,
"--output-dir", self.output_dir,
"--num-points", "1000000",
"--remove-outliers", "True",
"--use-bounding-box", "True"
], capture_output=True, text=True, timeout=600)
if result.returncode == 0 and os.path.exists(output_ply):
file_size = os.path.getsize(output_ply) / (1024 * 1024)
exports["pointcloud"] = {
"path": output_ply,
"size_mb": round(file_size, 2),
"type": "Point Cloud"
}
if exports:
return {"status": "success", "exports": exports}
# Fallback to COLMAP export
colmap_sparse = os.path.join(self.colmap_dir, "sparse", "0")
if os.path.exists(colmap_sparse):
output_ply = os.path.join(self.output_dir, "colmap_points.ply")
result = subprocess.run([
"colmap", "model_converter",
"--input_path", colmap_sparse,
"--output_path", output_ply,
"--output_type", "PLY"
], capture_output=True, text=True, timeout=300)
if os.path.exists(output_ply):
file_size = os.path.getsize(output_ply) / (1024 * 1024)
return {
"status": "success",
"exports": {
"pointcloud": {
"path": output_ply,
"size_mb": round(file_size, 2),
"type": "COLMAP Points"
}
}
}
return {"status": "error", "message": "All export methods failed"}
except Exception as e:
return {"status": "error", "message": str(e)}
|