foodforthought-cli 0.2.7__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/urdf/kinematics.py

@@ -0,0 +1,330 @@
+"""
+Joint parameter estimation for URDF generation.
+
+This module handles Phase 3 of the pipeline:
+1. Fit kinematic parameters to point cloud trajectories
+2. Estimate joint axes, pivots, and limits
+3. Generate kinematics.json output
+
+Implements the core kinematic optimization described in the research.
+"""
+
+import logging
+from typing import Dict, List, Optional, Tuple, Set
+from pathlib import Path
+from dataclasses import dataclass, asdict
+
+logger = logging.getLogger(__name__)
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+    np = None
+
+
+class KinematicsError(Exception):
+    """Error during kinematic estimation."""
+    pass
+
+
+@dataclass
+class JointParameters:
+    """Estimated parameters for a single joint."""
+    name: str
+    parent_link: str
+    child_link: str
+    joint_type: str  # revolute, prismatic, fixed
+    axis: List[float]  # [x, y, z] unit vector
+    origin: List[float]  # [x, y, z] position
+    limits: Dict[str, float]  # lower, upper (radians or meters)
+    confidence: float  # Estimation confidence [0, 1]
+
+    def to_dict(self) -> Dict:
+        return {
+            "name": self.name,
+            "parent_link": self.parent_link,
+            "child_link": self.child_link,
+            "type": self.joint_type,
+            "axis": self.axis,
+            "origin": self.origin,
+            "limits": self.limits,
+            "confidence": self.confidence,
+        }
+
+
+def load_per_frame_clouds(
+    session: "ScanSession",
+    masks: Dict[str, Dict[int, "np.ndarray"]],
+    depth_maps: Dict[int, "np.ndarray"],
+) -> Dict[str, List["np.ndarray"]]:
+    """
+    Load per-frame point clouds for each link.
+
+    Args:
+        session: ScanSession
+        masks: Per-frame segmentation masks
+        depth_maps: Per-frame depth maps
+
+    Returns:
+        Dict mapping link_name -> list of point clouds per frame
+    """
+    import cv2
+    from .lifting import lift_to_3d
+    from .scale import estimate_intrinsics_from_resolution
+
+    # Get intrinsics
+    w, h = session.metadata.resolution
+    intrinsics = estimate_intrinsics_from_resolution(w, h)
+
+    # Open video
+    cap = cv2.VideoCapture(str(session.video_path))
+    if not cap.isOpened():
+        raise KinematicsError(f"Could not open video: {session.video_path}")
+
+    frame_indices = sorted(depth_maps.keys())
+    link_clouds = {name: [] for name in masks.keys()}
+
+    try:
+        for frame_idx in frame_indices:
+            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+            ret, frame = cap.read()
+            if not ret:
+                for name in masks.keys():
+                    link_clouds[name].append(None)
+                continue
+
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            depth = depth_maps[frame_idx]
+
+            for link_name in masks.keys():
+                if frame_idx in masks[link_name]:
+                    mask = masks[link_name][frame_idx]
+                    points, _ = lift_to_3d(rgb, depth, mask, intrinsics)
+                    link_clouds[link_name].append(points if len(points) > 0 else None)
+                else:
+                    link_clouds[link_name].append(None)
+
+    finally:
+        cap.release()
+
+    return link_clouds
+
+
+def estimate_joint_parameters(
+    parent_traj: "LinkTrajectory",
+    child_traj: "LinkTrajectory",
+    joint_name: str,
+    joint_type_hint: Optional[str] = None,
+) -> JointParameters:
+    """
+    Estimate joint parameters between parent and child links.
+
+    Args:
+        parent_traj: Parent link trajectory
+        child_traj: Child link trajectory
+        joint_name: Name for this joint
+        joint_type_hint: Optional hint for joint type
+
+    Returns:
+        JointParameters with estimated values
+    """
+    from .motion_analysis import (
+        compute_relative_motion,
+        estimate_revolute_axis,
+        estimate_prismatic_axis,
+        estimate_joint_limits,
+    )
+
+    # Compute relative motion
+    rel_motion = compute_relative_motion(parent_traj, child_traj)
+
+    # Classify joint type
+    if joint_type_hint:
+        joint_type = joint_type_hint
+        confidence = 0.9  # User-provided hint
+    else:
+        joint_type, confidence = rel_motion.analyze_joint_type()
+
+    # Estimate axis and origin
+    if joint_type == "revolute":
+        axis, pivot, axis_conf = estimate_revolute_axis(rel_motion)
+        confidence = min(confidence, axis_conf)
+        origin = pivot.tolist()
+    elif joint_type == "prismatic":
+        axis, axis_conf = estimate_prismatic_axis(rel_motion)
+        confidence = min(confidence, axis_conf)
+        # Origin at parent centroid
+        origin = parent_traj.centroids[0].tolist()
+    else:  # fixed
+        axis = [0.0, 0.0, 1.0]
+        origin = child_traj.centroids[0].tolist()
+
+    # Normalize axis
+    axis_array = np.array(axis)
+    axis_array = axis_array / (np.linalg.norm(axis_array) + 1e-8)
+    axis = axis_array.tolist()
+
+    # Estimate limits
+    if joint_type != "fixed":
+        lower, upper = estimate_joint_limits(
+            rel_motion, joint_type, np.array(axis), np.array(origin)
+        )
+    else:
+        lower, upper = 0.0, 0.0
+
+    return JointParameters(
+        name=joint_name,
+        parent_link=parent_traj.name,
+        child_link=child_traj.name,
+        joint_type=joint_type,
+        axis=axis,
+        origin=origin,
+        limits={"lower": lower, "upper": upper},
+        confidence=confidence,
+    )
+
+
+def run_kinematic_optimization(
+    session: "ScanSession",
+    masks: Dict[str, Dict[int, "np.ndarray"]],
+    depth_maps: Dict[int, "np.ndarray"],
+    joint_type_hints: Optional[Dict[str, str]] = None,
+    progress_callback: Optional[callable] = None,
+) -> List[JointParameters]:
+    """
+    Run kinematic optimization to estimate joint parameters.
+
+    Args:
+        session: ScanSession
+        masks: Per-frame segmentation masks
+        depth_maps: Per-frame depth maps
+        joint_type_hints: Optional dict of link_name -> joint_type
+        progress_callback: Optional progress callback
+
+    Returns:
+        List of JointParameters
+    """
+    from .motion_analysis import extract_trajectories
+    from .topology import build_kinematic_tree, infer_joint_names
+
+    logger.info("Starting kinematic optimization...")
+
+    # Load per-frame clouds
+    logger.info("Loading per-frame point clouds...")
+    frame_indices = sorted(depth_maps.keys())
+    link_clouds = load_per_frame_clouds(session, masks, depth_maps)
+
+    # Extract trajectories
+    logger.info("Extracting link trajectories...")
+    trajectories = extract_trajectories(link_clouds, frame_indices)
+
+    if len(trajectories) < 2:
+        raise KinematicsError(
+            f"Need at least 2 links for kinematics, got {len(trajectories)}"
+        )
+
+    # Get fixed link hints from session
+    fixed_hints = {
+        link.name for link in session.links if link.is_fixed
+    }
+
+    # Build kinematic tree
+    logger.info("Building kinematic tree...")
+    tree = build_kinematic_tree(trajectories, fixed_hints)
+
+    # Get joint names
+    joint_names = infer_joint_names(tree)
+
+    # Estimate parameters for each joint
+    logger.info("Estimating joint parameters...")
+    joints = []
+    total = len(joint_names)
+
+    for i, ((parent, child), joint_name) in enumerate(joint_names.items()):
+        parent_traj = trajectories[parent]
+        child_traj = trajectories[child]
+
+        # Get type hint if available
+        type_hint = None
+        if joint_type_hints and child in joint_type_hints:
+            type_hint = joint_type_hints[child]
+
+        params = estimate_joint_parameters(
+            parent_traj, child_traj, joint_name, type_hint
+        )
+        joints.append(params)
+
+        logger.info(
+            f"  {joint_name}: {params.joint_type} "
+            f"(conf={params.confidence:.2f}, axis={params.axis})"
+        )
+
+        if progress_callback:
+            progress_callback(i + 1, total)
+
+    # Update session
+    from .scan_session import JointInfo
+    session.joints = [
+        JointInfo(
+            name=j.name,
+            parent_link=j.parent_link,
+            child_link=j.child_link,
+            joint_type=j.joint_type,
+            axis=j.axis,
+            origin=j.origin,
+            limits=j.limits,
+        )
+        for j in joints
+    ]
+    session.save_kinematics()
+    session.metadata.optimize_complete = True
+    session.save_metadata()
+
+    logger.info(f"Kinematic optimization complete: {len(joints)} joints")
+    return joints
+
+
+def save_kinematics_json(
+    output_path: Path,
+    links: List[str],
+    joints: List[JointParameters],
+    tree: "KinematicTree",
+) -> None:
+    """
+    Save kinematics to JSON file.
+
+    Args:
+        output_path: Path to output file
+        links: List of link names
+        joints: List of JointParameters
+        tree: KinematicTree structure
+    """
+    import json
+
+    data = {
+        "links": [
+            {
+                "name": name,
+                "parent": tree.get_parent(name),
+            }
+            for name in tree.get_ordered_links()
+        ],
+        "joints": [j.to_dict() for j in joints],
+    }
+
+    with open(output_path, 'w') as f:
+        json.dump(data, f, indent=2)
+
+    logger.info(f"Saved kinematics to {output_path}")
+
+
+__all__ = [
+    "KinematicsError",
+    "JointParameters",
+    "load_per_frame_clouds",
+    "estimate_joint_parameters",
+    "run_kinematic_optimization",
+    "save_kinematics_json",
+]