76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
import pandas as pd
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from pathlib import Path
|
|
|
|
def load_csv_with_standard_decimal(file_path):
|
|
with open(file_path, 'r') as f:
|
|
scene_line = f.readline().strip()
|
|
scene_name = scene_line.split(': ')[1].split('_')[0]
|
|
df = pd.read_csv(file_path, skiprows=1)
|
|
return df, scene_name
|
|
|
|
def process_position_data(data_dir):
|
|
# Group by scene
|
|
scene_data = {}
|
|
csv_files = list(Path(data_dir).glob('**/P*_PositionData*.csv'))
|
|
for file_path in csv_files:
|
|
df, scene_name = load_csv_with_standard_decimal(str(file_path))
|
|
if scene_name not in scene_data:
|
|
scene_data[scene_name] = []
|
|
scene_data[scene_name].append(df)
|
|
return scene_data
|
|
|
|
def plot_relative_positions(scene_data):
|
|
scenes = list(scene_data.keys())
|
|
# Create one figure with 4 subplots
|
|
fig, axes = plt.subplots(2, 2, figsize=(16, 14))
|
|
axes = axes.flatten()
|
|
for idx, (scene, recordings) in enumerate(scene_data.items()):
|
|
all_rel_x = []
|
|
all_rel_z = []
|
|
avg_obj_xs = []
|
|
avg_obj_zs = []
|
|
for df in recordings:
|
|
obj_x = df['TrackedObject_X'].values
|
|
obj_z = df['TrackedObject_Z'].values
|
|
avg_obj_x = np.mean(obj_x)
|
|
avg_obj_z = np.mean(obj_z)
|
|
avg_obj_xs.append(avg_obj_x)
|
|
avg_obj_zs.append(avg_obj_z)
|
|
rel_x = df['XR_Origin_X'].values - avg_obj_x
|
|
rel_z = df['XR_Origin_Z'].values - avg_obj_z
|
|
all_rel_x.extend(rel_x)
|
|
all_rel_z.extend(rel_z)
|
|
ax = axes[idx]
|
|
ax.scatter(all_rel_x, all_rel_z, alpha=0.2, c='blue', label='Player Positions (relative)')
|
|
ax.scatter(0, 0, c='red', s=800, marker='*', label='Agent', edgecolor='black', linewidths=2, zorder=10)
|
|
ax.set_title(f'Scene: {scene}')
|
|
ax.set_xlabel('X Position (relative to Agent)')
|
|
ax.set_ylabel('Z Position (relative to Agent)')
|
|
ax.legend()
|
|
ax.axis('equal')
|
|
ax.grid(True, linestyle='--', alpha=0.7)
|
|
# Also save individual images
|
|
plt.figure(figsize=(8, 8))
|
|
plt.scatter(all_rel_x, all_rel_z, alpha=0.2, c='blue', label='Player Positions (relative)')
|
|
plt.scatter(0, 0, c='red', s=800, marker='*', label='Agent', edgecolor='black', linewidths=2, zorder=10)
|
|
plt.title(f'Player Positions Relative to Agent\nScene: {scene}')
|
|
plt.xlabel('X Position (relative to Agent)')
|
|
plt.ylabel('Z Position (relative to Agent)')
|
|
plt.legend()
|
|
plt.axis('equal')
|
|
plt.grid(True, linestyle='--', alpha=0.7)
|
|
plt.tight_layout()
|
|
plt.savefig(f'room_layout_{scene}.png', dpi=300)
|
|
plt.close()
|
|
plt.tight_layout()
|
|
plt.savefig('room_layout_all_scenes.png', dpi=300)
|
|
plt.close()
|
|
|
|
def main():
|
|
scene_data = process_position_data('Recordings')
|
|
plot_relative_positions(scene_data)
|
|
|
|
if __name__ == "__main__":
|
|
main() |