批量旋转mesh文件

"""
Batch-rotate OBJ models in a directory by a specified quaternion.

- Loads OBJ using trimesh
- Applies rotation using scipy.spatial.transform.Rotation
- Saves to a sibling output directory (same level as input)

Example:
  python rot.py ./models/meshes --quat 0 0 0 1 --quat-format wxyz --suffix _rot

Requirements:
  pip install trimesh scipy numpy
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Iterable, List, Tuple

import numpy as np


def parse_args(argv: List[str]) -> argparse.Namespace:
	parser = argparse.ArgumentParser(
		description="Rotate all OBJ files in a directory by a quaternion and save to a sibling directory.",
		formatter_class=argparse.ArgumentDefaultsHelpFormatter,
	)

	parser.add_argument(
		"input_dir",
		type=Path,
		help="Directory containing OBJ files",
	)

	out_group = parser.add_argument_group("output")
	out_group.add_argument(
		"--outdir",
		type=Path,
		default=None,
		help="Output directory path. If not set, use sibling '<input_dir.name><suffix>'",
	)
	out_group.add_argument(
		"--suffix",
		type=str,
		default="_rotated",
		help="Suffix to append to input_dir.name when --outdir is not provided",
	)
	out_group.add_argument(
		"--overwrite",
		action="store_true",
		help="Allow writing into an existing output directory",
	)

	scan_group = parser.add_argument_group("scanning")
	scan_group.add_argument(
		"--recursive",
		action="store_true",
		help="Recurse into subdirectories to find OBJ files",
	)
	scan_group.add_argument(
		"--pattern",
		type=str,
		default="*.obj",
		help="Glob pattern for files to process",
	)

	rot_group = parser.add_argument_group("rotation")
	rot_group.add_argument(
		"--quat",
		type=float,
		nargs=4,
		required=True,
		metavar=("q0", "q1", "q2", "q3"),
		help="Quaternion components in the order specified by --quat-format",
	)
	rot_group.add_argument(
		"--quat-format",
		choices=["xyzw", "wxyz"],
		default="xyzw",
		help="Interpretation of --quat entries; scipy expects xyzw",
	)
	rot_group.add_argument(
		"--center",
		choices=["origin", "centroid", "bounds", "pivot"],
		default="origin",
		help="Rotation center: world origin, mesh centroid, AABB center, or a custom pivot",
	)
	rot_group.add_argument(
		"--pivot",
		type=float,
		nargs=3,
		metavar=("px", "py", "pz"),
		default=None,
		help="Custom pivot point (used when --center pivot)",
	)

	misc_group = parser.add_argument_group("misc")
	misc_group.add_argument(
		"--dry-run",
		action="store_true",
		help="Don't write files; just print planned actions",
	)
	misc_group.add_argument(
		"--verbose",
		"-v",
		action="store_true",
		help="Verbose logging",
	)

	args = parser.parse_args(argv)
	return args


def _to_xyzw(quat: Iterable[float], fmt: str) -> np.ndarray:
	q = np.asarray(list(quat), dtype=float).reshape(4)
	if fmt == "xyzw":
		return q
	elif fmt == "wxyz":
		w, x, y, z = q
		return np.array([x, y, z, w], dtype=float)
	else:
		raise ValueError(f"Unsupported quat format: {fmt}")


def _rotation_matrix_from_quat_xyzw(q_xyzw: np.ndarray) -> np.ndarray:
	# Lazy import to avoid mandatory dependency if user only runs --help
	from scipy.spatial.transform import Rotation as R  # type: ignore

	q = np.array(q_xyzw, dtype=float)
	if not np.isfinite(q).all():
		raise ValueError("Quaternion has non-finite entries")
	if np.linalg.norm(q) <= 0:
		raise ValueError("Quaternion has zero magnitude")
	# SciPy will normalize internally; we can leave as-is
	R3 = R.from_quat(q).as_matrix()  # xyzw
	M = np.eye(4, dtype=float)
	M[:3, :3] = R3
	return M


def _pivot_for_geometry(geom, center: str) -> np.ndarray:
	import trimesh

	if center == "origin":
		return np.zeros(3, dtype=float)

	if isinstance(geom, trimesh.Trimesh):
		if center == "centroid":
			return np.asarray(geom.centroid, dtype=float)
		elif center == "bounds":
			b = np.asarray(geom.bounds, dtype=float)  # (2,3)
			return (b[0] + b[1]) / 2.0
		else:
			raise ValueError("center 'pivot' requires explicit --pivot")
	else:
		# Treat as Scene (or unknown): use bounds center when centroid not defined
		try:
			b = np.asarray(geom.bounds, dtype=float)
			c = (b[0] + b[1]) / 2.0
		except Exception:
			c = np.zeros(3, dtype=float)

		if center in ("centroid", "bounds"):
			return c
		else:
			raise ValueError("center 'pivot' requires explicit --pivot")


def _compose_about_pivot(R4: np.ndarray, pivot: np.ndarray) -> np.ndarray:
	Tm = np.eye(4, dtype=float)
	Tp = np.eye(4, dtype=float)
	Tm[:3, 3] = -np.asarray(pivot, dtype=float)
	Tp[:3, 3] = np.asarray(pivot, dtype=float)
	return Tp @ R4 @ Tm


def _collect_files(root: Path, pattern: str, recursive: bool) -> List[Path]:
	if recursive:
		return sorted([p for p in root.rglob(pattern) if p.is_file()])
	else:
		return sorted([p for p in root.glob(pattern) if p.is_file()])


def _export_any(geom, out_path: Path) -> None:
	import trimesh

	out_path.parent.mkdir(parents=True, exist_ok=True)
	if isinstance(geom, trimesh.Trimesh):
		geom.export(out_path)
	else:
		# Scene or other exportables
		geom.export(out_path)


def process_directory(
	input_dir: Path,
	outdir: Path,
	quat: Iterable[float],
	quat_format: str = "xyzw",
	center: str = "origin",
	pivot: Tuple[float, float, float] | None = None,
	pattern: str = "*.obj",
	recursive: bool = False,
	dry_run: bool = False,
	verbose: bool = False,
) -> Tuple[int, int]:
	"""Process all matching files, returning (ok_count, fail_count)."""
	import trimesh

	files = _collect_files(input_dir, pattern, recursive)
	if verbose:
		print(f"Found {len(files)} file(s) under {input_dir}")

	q_xyzw = _to_xyzw(quat, quat_format)
	R4 = _rotation_matrix_from_quat_xyzw(q_xyzw)

	ok = 0
	fail = 0

	for f in files:
		# Preserve relative path structure if recursive
		rel = f.relative_to(input_dir)
		out_f = outdir / rel
		out_f = out_f.with_suffix(".obj")

		if verbose:
			print(f"Processing: {f} -> {out_f}")

		try:
			geom = trimesh.load(f, force=None)  # Trimesh or Scene

			pv: np.ndarray
			if center == "pivot":
				if pivot is None:
					raise ValueError("--center pivot requires --pivot px py pz")
				pv = np.asarray(pivot, dtype=float)
			else:
				pv = _pivot_for_geometry(geom, center)

			M = _compose_about_pivot(R4, pv)

			if not dry_run:
				geom.apply_transform(M)
				_export_any(geom, out_f)

			ok += 1
		except Exception as e:
			fail += 1
			print(f"[WARN] Failed processing {f}: {e}")

	return ok, fail


def main(argv: List[str]) -> int:
	args = parse_args(argv)

	input_dir: Path = args.input_dir
	if not input_dir.exists() or not input_dir.is_dir():
		print(f"Input directory not found or not a directory: {input_dir}")
		return 2

	if args.outdir is not None:
		outdir = args.outdir
	else:
		outdir = input_dir.parent / f"{input_dir.name}{args.suffix}"

	if outdir.exists() and not args.overwrite:
		# Allow reusing empty dir without overwrite flag
		try:
			is_empty = next(outdir.iterdir(), None) is None
		except FileNotFoundError:
			is_empty = True
		if not is_empty:
			print(f"Output directory exists and not empty: {outdir}. Use --overwrite to proceed.")
			return 3

	if args.center == "pivot" and args.pivot is None:
		print("--center pivot requires --pivot px py pz")
		return 4

	ok, fail = process_directory(
		input_dir=input_dir,
		outdir=outdir,
		quat=args.quat,
		quat_format=args.quat_format,
		center=args.center,
		pivot=tuple(args.pivot) if args.pivot is not None else None,
		pattern=args.pattern,
		recursive=args.recursive,
		dry_run=args.dry_run,
		verbose=args.verbose,
	)

	if args.dry_run:
		print(f"Dry-run complete. Would process {ok} file(s); failures encountered: {fail}")
	else:
		print(f"Done. Processed {ok} file(s); failures: {fail}. Output: {outdir}")

	return 0 if fail == 0 else 1


if __name__ == "__main__":
	sys.exit(main(sys.argv[1:]))