Qam2 / mesh_processor.py
GlimmaryKarl's picture
Upload folder using huggingface_hub
3a8a707 verified
Raw
History Blame Contribute Delete
4.43 kB
import os
import numpy as np
import pandas as pd
import ezdxf
from pathlib import Path
from scipy.spatial import KDTree
def fit_and_solve_closed_polynomial(pts: np.ndarray) -> dict:
"""
Fits 5 points into a 2D quadric polynomial curve: Ax^2 + Bxy + Cy^2 + Dx + Ey + F = 0
Returns solved algebraic coefficients and roots for the dataset.
"""
if len(pts) < 5:
# Pad with zeros if fewer than 5 points exist
return {"A": 0.0, "B": 0.0, "C": 0.0, "D": 0.0, "E": 0.0, "F": 0.0, "root_x": 0.0, "root_y": 0.0}
x = pts[:, 0]
y = pts[:, 1]
# Form design matrix for implicit conic curve fitting: Ax^2 + Bxy + Cy^2 + Dx + Ey = 1 (F = -1)
M = np.column_stack([x**2, x * y, y**2, x, y])
rhs = np.ones(len(pts))
try:
# Solve algebraic coefficients via least squares SVD
coeffs, _, _, _ = np.linalg.lstsq(M, rhs, rcond=None)
A, B, C, D, E = coeffs
F = -1.0
except Exception:
A, B, C, D, E, F = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
# Solve critical manifold center point (root solutions where derivatives = 0)
# dP/dx = 2Ax + By + D = 0 ; dP/dy = Bx + 2Cy + E = 0
det = 4 * A * C - B**2
if abs(det) > 1e-6:
root_x = (B * E - 2 * C * D) / det
root_y = (B * D - 2 * A * E) / det
else:
root_x, root_y = np.mean(x), np.mean(y)
return {
"poly_A": float(A),
"poly_B": float(B),
"poly_C": float(C),
"poly_D": float(D),
"poly_E": float(E),
"poly_F": float(F),
"root_x": float(root_x),
"root_y": float(root_y)
}
def process_mesh(input_csv_path: str, dxf_dir_path: str, max_faces=500, precision=2) -> pd.DataFrame:
"""
Reads DXF files, interpolates closed curves from 5 nearest neighbor points,
solves quadric polynomials, and returns the solved polynomial dataset.
"""
print(f" -> Processing Polynomial Geometry directly from DXF: {dxf_dir_path}")
dxf_dir = Path(dxf_dir_path)
if not dxf_dir.exists():
raise FileNotFoundError(f"DXF directory not found at {dxf_dir}")
all_points = []
dxf_files = sorted(list(dxf_dir.glob("*.dxf")))
if len(dxf_files) > 60:
dxf_files = dxf_files[::2]
for dxf_file in dxf_files:
try:
frame_id = int(''.join(filter(str.isdigit, dxf_file.stem)) or 0)
doc = ezdxf.readfile(dxf_file)
msp = doc.modelspace()
entities = msp.query('LINE')
sampling_rate = 2 if len(entities) > 500 else 1
for i, entity in enumerate(entities):
if i % sampling_rate != 0:
continue
all_points.append({'t': frame_id, 'x': entity.dxf.start.x, 'y': entity.dxf.start.y, 'z': entity.dxf.start.z})
all_points.append({'t': frame_id, 'x': entity.dxf.end.x, 'y': entity.dxf.end.y, 'z': entity.dxf.end.z})
except Exception:
continue
if not all_points:
raise ValueError("No valid point geometry vectors could be parsed from DXF layers.")
df_vertices = pd.DataFrame(all_points)
solved_dataset_rows = []
for frame_id, group in df_vertices.groupby('t'):
if len(group) < 5:
continue
pts_2d = group[['x', 'y']].drop_duplicates().values
if len(pts_2d) < 5:
continue
# Construct KD-Tree to locate 5 closest points for interpolated closed curves
tree = KDTree(pts_2d)
visited = set()
for idx, pt in enumerate(pts_2d):
if idx in visited or len(solved_dataset_rows) >= max_faces:
continue
# Query 5 nearest neighbors to form interpolated closed boundary loop
distances, indices = tree.query(pt, k=min(5, len(pts_2d)))
cluster_pts = pts_2d[indices]
# Solve quadric polynomial equations for the closed boundary
poly_sol = fit_and_solve_closed_polynomial(cluster_pts)
poly_sol['frame_id'] = frame_id
poly_sol['mean_z'] = float(group['z'].mean())
solved_dataset_rows.append(poly_sol)
# Compile solved polynomial representations directly into pandas dataset
dataset_df = pd.DataFrame(solved_dataset_rows)
print(f" ✅ Generated {len(dataset_df)} solved polynomial curves for dataset.")
return dataset_df