Heatmap_3D.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  1. import UnityEngine as ue
  2. import pandas as pd
  3. import numpy as np
  4. import matplotlib.pyplot as plt
  5. import matplotlib.cm as cm
  6. # import seaborn as sns
  7. # import matplotlib.patches as patches
  8. # import matplotlib.colors as colors
  9. # from sklearn import preprocessing
  10. # from pylab import *
  11. # from matplotlib import style
  12. from matplotlib.colors import LinearSegmentedColormap
  13. from mpl_toolkits.mplot3d import Axes3D
  14. WIDTH = int(70)
  15. HEIGHT = int(35)
  16. OBSTACLE_PATH = "Assets/Data_image/obstacle.pkl"
  17. POSITION_PATH = ue.Application.dataPath + "/Data_position/80/Walk4.csv"
  18. HEATMAP_PATH = "Assets/Data_image/80/heatmap3D4.png"
  19. # 1. Get position data from csv file
  20. data = pd.read_csv(POSITION_PATH, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
  21. data = data.round(0)
  22. # 2. Group by positions and count appearance
  23. data_count = data.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
  24. # 3. Assign x, y, z, width, depth, height
  25. x = data_count["Position x"].tolist()
  26. y = data_count["Position z"].tolist()
  27. z = np.zeros_like(len(x))
  28. dz = data_count["counts"].tolist() # Change
  29. # 4. Create figure and axes
  30. fig = plt.figure()
  31. ax = fig.add_subplot(111, projection='3d')
  32. # 5.1 Create custom colormap Day1, Day2, Day3
  33. cmap = LinearSegmentedColormap.from_list(name='day1', colors=[(0.40,0.76,0.65), (0.11,0.62,0.47)])
  34. # cmap = LinearSegmentedColormap.from_list(name='day2', colors=['grey', (0.99,0.55,0.38), (0.85,0.37,0.01)])
  35. # cmap = LinearSegmentedColormap.from_list(name='day3', colors=['grey', (0.55,0.63,0.80), (0.46,0.44,0.70)])
  36. # 5.2 Initialize array for coloring the bars
  37. dz_array = np.array(data_count['counts'])
  38. fracs = dz_array.astype(float) / dz_array.max()
  39. color_values = cmap(fracs.tolist())
  40. # 6. Create the bars
  41. # for i in range(len(x)):
  42. # img = ax.bar3d(x[i], y[i], z, 1, 1, dz[i], color=color_values[i], shade=False)
  43. img = ax.bar3d(x, y, z, 1, 1, dz, color=color_values, shade=False)
  44. # 7. Create Colorbar
  45. color_map = cm.ScalarMappable(cmap=cmap)
  46. color_map.set_array(dz)
  47. fig.colorbar(color_map)
  48. # 8.
  49. plt.show()
  50. # 10. Save 3D Heatmap
  51. # heatmap.get_figure().savefig(HEATMAP_PATH, transparent=True)
  52. fig.savefig(HEATMAP_PATH, transparent=True)