plane_generate.py

import os
import numpy as np
import pandas as pd
import random
import plotly.graph_objects as go
import matplotlib.pyplot as plt

# Directory containing the .txt files
directory = os.path.expanduser("/home/ma1614/zero123/zero123/dataset/ob_in_cam")

# List to store all pose matrices
pose_matrices = []

# Iterate over each file in the directory
for filename in sorted(os.listdir(directory)):
    if filename.endswith(".txt"):
        file_path = os.path.join(directory, filename)
        
        # Read the content of the file
        with open(file_path, "r") as file:
            lines = file.readlines()
        
        # Convert the pose data into a NumPy array
        pose_matrix = []
        for line in lines:
            pose_matrix.append([float(x) for x in line.strip().split()])
        
        # Append to the list as a NumPy array
        pose_matrices.append(np.array(pose_matrix))

# Convert the list of poses into a NumPy array for structured data access
pose_matrices = np.array(pose_matrices)

# Parameters for data generation
num_points_per_pose = 10  # Number of points to generate per reference point
max_radius = 0.3  # Fraction of the radius for clustering closer to the center

# Generate unique colors for each reference point
num_references = len(pose_matrices)
colors = plt.cm.tab20(np.linspace(0, 1, num_references))  # Generate colors from a colormap

# Initialize storage for all generated data
all_data = []

# Initialize a Plotly 3D figure
fig = go.Figure()

# Process each pose matrix
for idx, pose_matrix in enumerate(pose_matrices):
    # Extract the reference point (translation component from pose matrix)
    ref_point = pose_matrix[:3, 3]  # Assuming the last column is [x, y, z, 1]

    # Calculate the radius dynamically for the current reference point
    radius = np.linalg.norm(ref_point)
    print(f"Processing Pose Matrix {idx + 1}, Reference Point: {ref_point}, Radius: {radius:.4f}")

    # Save the reference point itself as a row in the desired format
    ref_azimuth = 0  # Assume reference point azimuth is 0
    ref_elevation = 0  # Assume reference point elevation is 0
    reference_row = [ref_azimuth, ref_elevation, radius, ref_point[0], ref_point[1], ref_point[2]]
    all_data.append(reference_row)

    # Plot the reference point with its unique color
    fig.add_trace(go.Scatter3d(
        x=[ref_point[0]],
        y=[ref_point[1]],
        z=[ref_point[2]],
        mode='markers',
        marker=dict(size=8, color=f'rgba({colors[idx][0] * 255},{colors[idx][1] * 255},{colors[idx][2] * 255},1)'),
        name=f'Reference Point {idx + 1}'
    ))

    # Generate points surrounding the current reference point
    generated_points = []
    for _ in range(num_points_per_pose):
        # Generate azimuth (phi) randomly in [0, 2*pi)
        azimuth = random.uniform(0, 2 * np.pi)
        # Use a constant elevation (0) to maintain the same z-plane
        elevation = 0
        # Non-linear scaling to cluster points near the center
        r = radius * max_radius * (random.uniform(0, 1) ** 0.5)  # Square root for denser center

        # Convert spherical to Cartesian coordinates
        x = r * np.cos(elevation) * np.cos(azimuth)
        y = r * np.cos(elevation) * np.sin(azimuth)
        z = 0  # Keep z constant relative to the reference point's z-plane

        # Adjust coordinates relative to the current reference point
        x += ref_point[0]
        y += ref_point[1]
        z += ref_point[2]

        # Save the generated point in the same format
        generated_points.append([azimuth, elevation, r, x, y, z])

    # Add generated points to all_data
    all_data.extend(generated_points)

    # Extract Cartesian coordinates for plotting
    generated_points = np.array(generated_points)
    fig.add_trace(go.Scatter3d(
        x=generated_points[:, 3],
        y=generated_points[:, 4],
        z=generated_points[:, 5],
        mode='markers',
        marker=dict(size=4, color=f'rgba({colors[idx][0] * 255},{colors[idx][1] * 255},{colors[idx][2] * 255},1)'),
        name=f'Generated Points {idx + 1}'
    ))

# Convert all data to a DataFrame
columns = ['Azimuth (rad)', 'Elevation (rad)', 'Radius', 'X', 'Y', 'Z']
final_df = pd.DataFrame(all_data, columns=columns)

# Save the DataFrame to a CSV file
output_file = "data_with_pose_matrices.csv"
final_df.to_csv(output_file, index=False)
print(f"Data saved to '{output_file}'.")

# Update layout for better interaction
fig.update_layout(
    title='Interactive 3D Points from Pose Matrices with Unique Colors',
    scene=dict(
        xaxis_title='X-axis',
        yaxis_title='Y-axis',
        zaxis_title='Z-axis',
        aspectmode='cube'  # Ensure equal scaling
    )
)

# Show the plot
fig.show()