examples/example_tactile_feature_extraction.py

"""
Example script demonstrating tactile-visual feature extraction.

This script:
1. Generates fake wrist and tactile images
2. Generates fake camera parameters
3. Computes 3D pose of tactile sensor using forward kinematics
4. Projects the pose to 2D bounding boxes
5. Extracts features h_tau and h_w using CLIP-like backbone
"""

import numpy as np
import matplotlib.pyplot as plt
from dobot_control.tactile_feature_extraction import (
    TactileFeatureExtractor,
    ForwardKinematics,
    CameraProjection,
    generate_fake_camera_params,
    generate_fake_images
)


def main():
    """Main example function."""
    print("=" * 60)
    print("Tactile-Visual Feature Extraction Example")
    print("=" * 60)
    
    # ========== Step 1: Generate Fake Data ==========
    print("\n[Step 1] Generating fake images and camera parameters...")
    
    img_size = (640, 480)
    wrist_image, tactile_image = generate_fake_images(img_size)
    K_w, E_w = generate_fake_camera_params(img_size)  # Wrist camera params
    K_tp, E_tp = generate_fake_camera_params(img_size)  # Third-person camera params
    
    print(f"  - Wrist image shape: {wrist_image.shape}")
    print(f"  - Tactile image shape: {tactile_image.shape}")
    print(f"  - Camera intrinsic K:\n{K_w}")
    
    # ========== Step 2: Compute Tactile Sensor 3D Pose ==========
    print("\n[Step 2] Computing tactile sensor 3D pose using forward kinematics...")
    
    # Fake joint angles (in radians)
    joint_angles = np.array([
        0.1,   # q0
        0.2,   # q1
        -0.3,  # q2
        0.4,   # q3
        -0.1,  # q4
        0.2    # q5
    ])
    
    # Sensor offset from end-effector (in meters)
    sensor_offset = np.array([0.0, 0.0, 0.02])  # 2cm offset in z-direction
    
    # Compute sensor pose
    P_sensor = ForwardKinematics.compute_tactile_sensor_pose(
        joint_angles=joint_angles,
        robot_type="Nova 2",
        sensor_offset=sensor_offset
    )
    
    print(f"  - Joint angles (rad): {joint_angles}")
    print(f"  - Sensor pose (position): {P_sensor[:3, 3]}")
    print(f"  - Sensor pose matrix shape: {P_sensor.shape}")
    
    # ========== Step 3: Project 3D Pose to 2D Bounding Boxes ==========
    print("\n[Step 3] Projecting 3D pose to 2D bounding boxes...")
    
    # Define sensor size (width, height in meters)
    sensor_size = (0.04, 0.04)  # 4cm x 4cm sensor
    
    # Compute bounding box in wrist view
    bbox_w = CameraProjection.compute_sensor_bounding_box(
        sensor_pose=P_sensor,
        sensor_size=sensor_size,
        K=K_w,
        E=E_w,
        img_size=img_size
    )
    
    # Compute bounding box in third-person view
    bbox_tp = CameraProjection.compute_sensor_bounding_box(
        sensor_pose=P_sensor,
        sensor_size=sensor_size,
        K=K_tp,
        E=E_tp,
        img_size=img_size
    )
    
    print(f"  - Wrist bounding box: x=[{bbox_w['x_min']:.1f}, {bbox_w['x_max']:.1f}], "
          f"y=[{bbox_w['y_min']:.1f}, {bbox_w['y_max']:.1f}]")
    print(f"  - Third-person bounding box: x=[{bbox_tp['x_min']:.1f}, {bbox_tp['x_max']:.1f}], "
          f"y=[{bbox_tp['y_min']:.1f}, {bbox_tp['y_max']:.1f}]")
    
    # ========== Step 4: Extract Features ==========
    print("\n[Step 4] Extracting features using CLIP-like backbone...")
    
    # Initialize feature extractor
    extractor = TactileFeatureExtractor(
        img_size=640,  # Will resize images to this size
        patch_size=16,
        embed_dim=768,
        device='cpu'  # Use 'cuda' if GPU available
    )
    
    # Extract features
    features = extractor.extract_features(
        wrist_image=wrist_image,
        tactile_image=tactile_image,
        bbox_wrist=bbox_w,
        bbox_tp=bbox_tp
    )
    
    h_tau = features['h_tau']
    h_w = features['h_w']
    
    print(f"  - h_tau shape: {h_tau.shape}")
    print(f"  - h_w shape: {h_w.shape}")
    print(f"  - h_tau mean: {h_tau.mean().item():.4f}, std: {h_tau.std().item():.4f}")
    print(f"  - h_w mean: {h_w.mean().item():.4f}, std: {h_w.std().item():.4f}")
    
    # ========== Step 5: Visualize Results ==========
    print("\n[Step 5] Visualizing results...")
    
    # Create visualization
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    
    # Plot wrist image with bounding box
    axes[0, 0].imshow(wrist_image)
    rect_w = plt.Rectangle(
        (bbox_w['x_min'], bbox_w['y_min']),
        bbox_w['x_max'] - bbox_w['x_min'],
        bbox_w['y_max'] - bbox_w['y_min'],
        linewidth=2, edgecolor='r', facecolor='none'
    )
    axes[0, 0].add_patch(rect_w)
    axes[0, 0].set_title('Wrist Camera Image\n(Red box: projected sensor)')
    axes[0, 0].axis('off')
    
    # Plot tactile image
    axes[0, 1].imshow(tactile_image)
    axes[0, 1].set_title('Tactile Sensor Image')
    axes[0, 1].axis('off')
    
    # Plot feature vectors
    axes[1, 0].plot(h_tau.detach().numpy())
    axes[1, 0].set_title('h_tau (Tactile Features)')
    axes[1, 0].set_xlabel('Feature Dimension')
    axes[1, 0].set_ylabel('Feature Value')
    axes[1, 0].grid(True)
    
    axes[1, 1].plot(h_w.detach().numpy())
    axes[1, 1].set_title('h_w (Wrist Visual Features)')
    axes[1, 1].set_xlabel('Feature Dimension')
    axes[1, 1].set_ylabel('Feature Value')
    axes[1, 1].grid(True)
    
    plt.tight_layout()
    plt.savefig('tactile_feature_extraction_example.png', dpi=150)
    print("  - Visualization saved to 'tactile_feature_extraction_example.png'")
    
    # ========== Step 6: Demonstrate Feature Grid Mapping ==========
    print("\n[Step 6] Demonstrating feature grid mapping...")
    
    # Map bounding box to feature grid
    feat_x_min, feat_y_min, feat_x_max, feat_y_max = extractor.map_bbox_to_feature_grid(
        bbox_w, (wrist_image.shape[1], wrist_image.shape[0])
    )
    
    feat_h, feat_w = extractor.feature_grid_shape
    print(f"  - Original image size: {wrist_image.shape[1]} x {wrist_image.shape[0]}")
    print(f"  - Feature grid size: {feat_w} x {feat_h}")
    print(f"  - Original bbox: x=[{bbox_w['x_min']:.1f}, {bbox_w['x_max']:.1f}], "
          f"y=[{bbox_w['y_min']:.1f}, {bbox_w['y_max']:.1f}]")
    print(f"  - Feature grid bbox: x=[{feat_x_min}, {feat_x_max}], "
          f"y=[{feat_y_min}, {feat_y_max}]")
    print(f"  - Number of tokens selected: {(feat_x_max - feat_x_min + 1) * (feat_y_max - feat_y_min + 1)}")
    
    print("\n" + "=" * 60)
    print("Example completed successfully!")
    print("=" * 60)
    
    # Show plot (optional)
    try:
        plt.show()
    except:
        pass


if __name__ == "__main__":
    main()